LOGAR/src/client_transport.py

"""
LOGAR - Unified Client Transport & PKI Enrollment Module
Encapsulates mutual TLS (mTLS 1.3), OpenPGP fallback envelope framing,
and dynamic certificate enrollment with the LOGAR Server Hub.
"""

import os
import sys
import json
import socket
import ssl
import struct
import urllib.request
import warnings
from datetime import datetime, timezone, timedelta
from typing import Optional, List, Dict, Any

# Suppress cryptography / pgpy deprecation notices
warnings.filterwarnings("ignore")

import pgpy
from log_collector import (
    get_machine_identifier,
    is_path_writable,
    get_default_data_dir,
    resolve_writable_path,
    collect_system_metrics
)
try:
    from src.version import LOGAR_VERSION
    from src.client_updater import ClientUpdater
except ImportError:
    from version import LOGAR_VERSION
    from client_updater import ClientUpdater

CONFIG_FILE_NAME = "client_config.json"


def is_cert_expiring_soon(cert_path: str, threshold_days: int = 30) -> bool:
    """Checks if client certificate at cert_path is expiring within threshold_days."""
    if not os.path.exists(cert_path):
        return True
    try:
        from cryptography import x509
        with open(cert_path, "r", encoding="utf-8") as f:
            cert = x509.load_pem_x509_certificate(f.read().encode("utf-8"))
        expiry = getattr(cert, "not_valid_after_utc", None)
        if expiry is None:
            expiry = cert.not_valid_after.replace(tzinfo=timezone.utc)
        now = datetime.now(timezone.utc)
        return expiry <= (now + timedelta(days=threshold_days))
    except Exception:
        return True


def enroll_client_if_needed(
    hub_url: str,
    enrollment_secret: str,
    cert_dir: str,
    client_id: str,
    hostname: str,
    os_type: str = "windows",
    force_renew: bool = False,
    threshold_days: int = 30
) -> bool:
    """Bootstraps client enrollment if certificates are missing or expiring soon."""
    os.makedirs(cert_dir, exist_ok=True)
    ca_path = os.path.join(cert_dir, "ca.crt")
    cert_path = os.path.join(cert_dir, "client.crt")
    key_path = os.path.join(cert_dir, "client.key")

    if not force_renew and os.path.exists(ca_path) and os.path.exists(cert_path) and os.path.exists(key_path):
        if not is_cert_expiring_soon(cert_path, threshold_days=threshold_days):
            return True
        print(f"[*] Client certificate at {cert_path} is expiring within {threshold_days} days. Auto-renewing...")

    action_name = "re-enrolling" if os.path.exists(cert_path) else "enrolling"
    print(f"[*] Bootstrapping client {action_name} with LOGAR Hub at {hub_url}...")
    enroll_endpoint = f"{hub_url.rstrip('/')}/api/client/enroll"
    payload = {
        "client_id": client_id,
        "hostname": hostname,
        "os": os_type,
        "enrollment_secret": enrollment_secret,
        "version": LOGAR_VERSION
    }
    req = urllib.request.Request(
        enroll_endpoint,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=10) as resp:
        if resp.status != 200:
            raise RuntimeError(f"Enrollment failed with status code {resp.status}")
        data = json.loads(resp.read().decode("utf-8"))

    with open(ca_path, "w", encoding="utf-8") as f:
        f.write(data["ca_cert"])
    with open(cert_path, "w", encoding="utf-8") as f:
        f.write(data["client_cert"])
    with open(key_path, "w", encoding="utf-8") as f:
        f.write(data["client_key"])

    try:
        os.chmod(key_path, 0o600)
    except Exception:
        pass

    print(f"[+] Client certificates updated successfully in {os.path.abspath(cert_dir)}")
    return True


def get_tls_socket(hub_host: str, hub_port: int, cert_dir: str):
    """Establishes an mTLS connection with the LOGAR hub using client certificates."""
    ca_path = os.path.join(cert_dir, "ca.crt")
    cert_path = os.path.join(cert_dir, "client.crt")
    key_path = os.path.join(cert_dir, "client.key")

    if not (os.path.exists(ca_path) and os.path.exists(cert_path) and os.path.exists(key_path)):
        raise FileNotFoundError(f"mTLS certificates not found in '{cert_dir}'. Enroll client first.")

    ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH, cafile=ca_path)
    ctx.load_cert_chain(certfile=cert_path, keyfile=key_path)
    ctx.minimum_version = ssl.TLSVersion.TLSv1_3
    ctx.check_hostname = False

    raw_sock = socket.create_connection((hub_host, hub_port), timeout=15)
    return ctx.wrap_socket(raw_sock, server_hostname=hub_host)


def load_config(config_path: str = CONFIG_FILE_NAME) -> dict:
    """Loads and normalizes client configuration JSON."""
    actual_config_path = config_path
    if not os.path.exists(actual_config_path):
        data_dir_cfg = os.path.join(get_default_data_dir("LOGAR"), os.path.basename(config_path))
        if os.path.exists(data_dir_cfg):
            actual_config_path = data_dir_cfg

    if not os.path.exists(actual_config_path):
        raise FileNotFoundError(
            f"Client configuration file not found at: {config_path}\n"
            f"Generate one from the server using: python Server.py --create-client-config --client-out {config_path}"
        )
    with open(actual_config_path, "r", encoding="utf-8") as f:
        config = json.load(f)

    # Resolve cert_dir to a writable directory
    cfg_dir = os.path.dirname(os.path.abspath(actual_config_path))
    raw_cert_dir = config.get("cert_dir", "certs")
    if os.path.isabs(raw_cert_dir) and is_path_writable(raw_cert_dir):
        config["cert_dir"] = raw_cert_dir
    else:
        candidate = os.path.join(cfg_dir, os.path.basename(raw_cert_dir))
        if is_path_writable(candidate) or is_path_writable(cfg_dir):
            config["cert_dir"] = candidate
        else:
            config["cert_dir"] = resolve_writable_path("certs", "LOGAR", "certs")

    # If ca_cert is present in config and ca.crt doesn't exist yet, save it
    ca_cert_pem = config.get("ca_cert")
    if ca_cert_pem and isinstance(ca_cert_pem, str):
        ca_file = os.path.join(config["cert_dir"], "ca.crt")
        if not os.path.exists(ca_file):
            try:
                os.makedirs(config["cert_dir"], exist_ok=True)
                with open(ca_file, "w", encoding="utf-8") as f:
                    f.write(ca_cert_pem)
            except Exception:
                pass

    return config


def ensure_enrolled(config: dict, machine_id: str, os_type: str = "windows") -> bool:
    """Proactively ensures client certificate enrollment with the Hub if enrollment_secret is available."""
    cert_dir = config.get("cert_dir", "certs")
    enrollment_secret = config.get("enrollment_secret")
    if not enrollment_secret:
        return False

    server_host = config.get("server_host", "127.0.0.1")
    hermes_host = config.get("hermes_host", server_host)
    if hermes_host in ("0.0.0.0", "::", None, ""):
        hermes_host = server_host if server_host not in ("0.0.0.0", "::", None, "") else "127.0.0.1"
    hermes_port = config.get("hermes_port", 8443)
    hub_url = f"http://{hermes_host}:{hermes_port}"
    try:
        return enroll_client_if_needed(hub_url, enrollment_secret, cert_dir, machine_id, machine_id, os_type=os_type)
    except Exception as e:
        print(f"[!] Warning: Initial client enrollment bootstrap with Hub at {hub_url} failed: {e}")
        return False


def send_encrypted_logs_over_socket(
    config: dict,
    logs: list,
    os_type: Optional[str] = None,
    machine_id: Optional[str] = None
) -> dict:
    """
    Streams logs to the LOGAR hub.
    Uses mutual TLS 1.3 (mTLS) with client certificates if available,
    or falls back to OpenPGP encrypted envelope over TCP socket.
    """
    node_os = os_type or ("windows" if sys.platform == "win32" else "linux")
    server_host = config["server_host"]
    server_port = int(config["server_port"])
    cert_dir = config.get("cert_dir", "certs")
    enrollment_secret = config.get("enrollment_secret")
    actual_machine_id = machine_id or config.get("client_id") or get_machine_identifier()

    # Attempt automatic enrollment bootstrap if certs are missing and secret is provided
    hub_url = None
    if enrollment_secret:
        hermes_host = config.get("hermes_host", server_host)
        if hermes_host in ("0.0.0.0", "::", None):
            hermes_host = server_host if server_host not in ("0.0.0.0", "::", None) else "127.0.0.1"
        hermes_port = config.get("hermes_port", 8443)
        hub_url = f"http://{hermes_host}:{hermes_port}"
        try:
            enroll_client_if_needed(hub_url, enrollment_secret, cert_dir, actual_machine_id, actual_machine_id, os_type=node_os)
        except Exception as e:
            print(f"[!] Warning: Enrollment bootstrap failed: {e}")

    ca_path = os.path.join(cert_dir, "ca.crt")
    cert_path = os.path.join(cert_dir, "client.crt")
    key_path = os.path.join(cert_dir, "client.key")
    has_mtls_certs = os.path.exists(ca_path) and os.path.exists(cert_path) and os.path.exists(key_path)

    if not has_mtls_certs and enrollment_secret and hub_url:
        try:
            enroll_client_if_needed(hub_url, enrollment_secret, cert_dir, actual_machine_id, actual_machine_id, os_type=node_os)
            has_mtls_certs = os.path.exists(ca_path) and os.path.exists(cert_path) and os.path.exists(key_path)
        except Exception as e:
            print(f"[!] Warning: Enrollment before TLS connection failed: {e}")

    if has_mtls_certs:
        print(f"[*] Connecting to LOGAR server at {server_host}:{server_port} over mTLS (TLS 1.3)...")
        sock = None
        try:
            sock = get_tls_socket(server_host, server_port, cert_dir)
        except (ssl.SSLError, ssl.CertificateError, ConnectionResetError) as tls_err:
            if enrollment_secret and hub_url:
                print(f"[!] TLS handshake error ({tls_err}). Re-enrolling with LOGAR Hub...")
                try:
                    enroll_client_if_needed(hub_url, enrollment_secret, cert_dir, actual_machine_id, actual_machine_id, os_type=node_os, force_renew=True)
                    sock = get_tls_socket(server_host, server_port, cert_dir)
                except Exception as retry_err:
                    print(f"[!] Re-enrollment or reconnection retry failed: {retry_err}")
                    raise
            else:
                raise

        with sock:
            payload = {
                "server": actual_machine_id,
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "logs": logs,
                "metrics": collect_system_metrics(),
                "version": LOGAR_VERSION
            }
            payload_bytes = json.dumps(payload).encode("utf-8")
            frame = struct.pack(">I", len(payload_bytes)) + payload_bytes
            sock.sendall(frame)

            resp_len_bytes = sock.recv(4)
            if not resp_len_bytes:
                raise ConnectionError("Server closed mTLS connection without response.")
            resp_len = struct.unpack(">I", resp_len_bytes)[0]
            resp_bytes = bytearray()
            while len(resp_bytes) < resp_len:
                chunk = sock.recv(min(4096, resp_len - len(resp_bytes)))
                if not chunk:
                    break
                resp_bytes.extend(chunk)

            response = json.loads(resp_bytes.decode("utf-8"))
            print(f"[+] Server response: {response}")
            return response

    # Fallback to OpenPGP envelope over plain TCP socket
    auth_token = config.get("auth_token", "")
    pub_key_armored = config.get("server_public_key")
    if not pub_key_armored:
        raise ValueError("No server public key or mTLS certificates available for connection.")

    expected_fp = config.get("server_fingerprint", "").replace(" ", "").upper()
    pub_key, _ = pgpy.PGPKey.from_blob(pub_key_armored)
    actual_fp = str(pub_key.fingerprint).replace(" ", "").upper()
    if expected_fp and actual_fp != expected_fp:
        raise ValueError(f"Server fingerprint mismatch! Expected {expected_fp}, but key has {actual_fp}.")

    payload = {
        "server": actual_machine_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "logs": logs,
        "metrics": collect_system_metrics(),
        "version": LOGAR_VERSION
    }
    payload_json = json.dumps(payload)
    pgp_msg = pgpy.PGPMessage.new(payload_json)
    encrypted_msg = pub_key.encrypt(pgp_msg)
    encrypted_armored = str(encrypted_msg)

    envelope = {
        "auth_token": auth_token,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "encrypted_payload": encrypted_armored
    }
    envelope_bytes = json.dumps(envelope).encode("utf-8")

    print(f"[*] Connecting to LOGAR server at {server_host}:{server_port} over secure TCP socket...")
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.settimeout(15.0)
        sock.connect((server_host, server_port))

        frame = struct.pack(">I", len(envelope_bytes)) + envelope_bytes
        sock.sendall(frame)

        resp_len_bytes = sock.recv(4)
        if not resp_len_bytes:
            raise ConnectionError("Server closed connection without response.")

        resp_len = struct.unpack(">I", resp_len_bytes)[0]
        resp_bytes = bytearray()
        while len(resp_bytes) < resp_len:
            chunk = sock.recv(min(4096, resp_len - len(resp_bytes)))
            if not chunk:
                break
            resp_bytes.extend(chunk)

        response = json.loads(resp_bytes.decode("utf-8"))
        print(f"[+] Server response: {response}")
        return response


def check_and_apply_auto_update(config: dict, machine_id: str, os_type: Optional[str] = None, auto_apply: bool = True) -> Dict[str, Any]:
    """Helper to check and apply auto updates for edge forwarders."""
    updater = ClientUpdater(config=config, os_type=os_type)
    return updater.check_and_apply_auto_update(client_id=machine_id, auto_apply=auto_apply)