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

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

import pgpy

try:
    import win32evtlog
except ImportError:
    win32evtlog = None

CONFIG_FILE_NAME = "client_config.json"
STATE_FILE_NAME = "client_state.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
):
    """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
    }
    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 get_state_path(config_path: str, custom_state_path: Optional[str] = None) -> str:
    if custom_state_path:
        return custom_state_path
    config_dir = os.path.dirname(os.path.abspath(config_path))
    return os.path.join(config_dir, STATE_FILE_NAME)


def load_state(state_path: str) -> dict:
    if os.path.exists(state_path):
        try:
            with open(state_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"[!] Warning: Failed to read state file '{state_path}': {e}")
            return {}
    return {}


def save_state(state_path: str, state: dict):
    try:
        temp_path = f"{state_path}.tmp"
        with open(temp_path, "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2)
        os.replace(temp_path, state_path)
    except Exception as e:
        print(f"[!] Warning: Could not save client state to '{state_path}': {e}")


def commit_state(state: dict, state_path: str):
    if "new_last_record_number" in state:
        val = state.pop("new_last_record_number")
        if val:
            state["last_record_number"] = val
    if "new_sent_record_ids" in state:
        state["sent_record_ids"] = state.pop("new_sent_record_ids")
    save_state(state_path, state)


def load_config(config_path: str = CONFIG_FILE_NAME):
    if not os.path.exists(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(config_path, "r", encoding="utf-8") as f:
        return json.load(f)


def get_machine_identifier() -> str:
    """
    Returns the hostname of the machine sending the logs,
    and appends the network/DNS domain if available.
    """
    fqdn = socket.getfqdn()
    if fqdn and "." in fqdn and not fqdn.startswith("localhost"):
        return fqdn

    hostname = socket.gethostname()
    user_dns_domain = os.environ.get("USERDNSDOMAIN")
    if user_dns_domain and user_dns_domain.lower() != hostname.lower():
        return f"{hostname}.{user_dns_domain.lower()}"

    try:
        host_ip = socket.gethostbyname(hostname)
        canonical_name = socket.gethostbyaddr(host_ip)[0]
        if canonical_name and "." in canonical_name and not canonical_name.startswith("localhost"):
            return canonical_name
    except Exception:
        pass

    return hostname


def get_recent_windows_logs(hours: int = 24, state: Optional[dict] = None) -> list:
    """
    Scans the Windows Application Event Log backwards for events within the window.
    Edge Filtering: Retains INFO, WARNING, and ERROR. Drops Audit and Debug noise.
    State Tracking: Skips events older than lookback window (default 24h) and events
    that have already been sent in previous runs.
    """
    if win32evtlog is None:
        print("[!] pywin32 is not installed or not running on Windows. Returning mock/empty candidate list.")
        return []

    server = "localhost"
    log_type = "Application"
    flags = win32evtlog.EVENTLOG_BACKWARDS_READ | win32evtlog.EVENTLOG_SEQUENTIAL_READ

    try:
        hand = win32evtlog.OpenEventLog(server, log_type)
    except Exception as e:
        print(f"[!] Error opening Windows event log: {e}")
        return []

    logs = []
    cutoff_time = datetime.now() - timedelta(hours=hours)
    machine_id = get_machine_identifier()

    last_record_number = 0
    sent_record_ids = set()
    if state:
        last_record_number = int(state.get("last_record_number", 0))
        sent_record_ids = set(state.get("sent_record_ids", []))

    # Windows Event Log EventTypes:
    # 1: EVENTLOG_ERROR_TYPE -> ERROR
    # 2: EVENTLOG_WARNING_TYPE -> WARNING
    # 4: EVENTLOG_INFORMATION_TYPE -> INFO
    # Excludes: 8 (Audit Success), 16 (Audit Failure), and other verbose noise
    sev_map = {
        1: "ERROR",
        2: "WARNING",
        4: "INFO"
    }

    newest_record_number = 0
    collected_record_ids = []

    while True:
        events = win32evtlog.ReadEventLog(hand, flags, 0)
        if not events:
            break

        for event in events:
            rec_num = int(event.RecordNumber)
            if newest_record_number == 0:
                newest_record_number = rec_num

            # 1. Skip entries older than lookback window (default: 24h)
            if event.TimeGenerated < cutoff_time:
                break

            # 2. Skip already sent events if we've reached records <= last_record_number
            # (unless the log was cleared and numbers wrapped, i.e. newest_record_number < last_record_number)
            if last_record_number > 0 and newest_record_number >= last_record_number:
                if rec_num <= last_record_number:
                    break

            rec_id = f"{rec_num}:{event.TimeGenerated.isoformat()}"
            if rec_id in sent_record_ids:
                continue

            # Filter: upload everything from INFO to ERROR only
            if event.EventType in sev_map:
                msg = " ".join(event.StringInserts) if event.StringInserts else "Event Log Entry"
                logs.append({
                    "server": machine_id,
                    "os_type": "windows",
                    "signature": event.SourceName or "Windows-Event",
                    "severity": sev_map[event.EventType],
                    "message": msg[:2048]  # Limit message length
                })
                collected_record_ids.append(rec_id)

        if events[-1].TimeGenerated < cutoff_time:
            break
        if last_record_number > 0 and newest_record_number >= last_record_number and events[-1].RecordNumber <= last_record_number:
            break

    win32evtlog.CloseEventLog(hand)

    if state is not None:
        target_rec = max(newest_record_number, last_record_number)
        state["new_last_record_number"] = target_rec
        state["new_sent_record_ids"] = (list(sent_record_ids) + collected_record_ids)[-1000:]
        state["last_run_timestamp"] = datetime.now(timezone.utc).isoformat()

    return logs


def send_encrypted_logs_over_socket(config: dict, logs: list):
    """
    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.
    """
    server_host = config["server_host"]
    server_port = int(config["server_port"])
    cert_dir = config.get("cert_dir", "certs")
    enrollment_secret = config.get("enrollment_secret")
    machine_id = 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)
        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, machine_id, machine_id, os_type="windows")
        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 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, machine_id, machine_id, os_type="windows", 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": machine_id,
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "logs": logs
            }
            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": machine_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "logs": logs
    }
    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 main():
    parser = argparse.ArgumentParser(description="LOGAR Windows Edge Log Forwarder with State Tracking")
    parser.add_argument("--config", default=CONFIG_FILE_NAME, help="Path to client_config.json")
    parser.add_argument("--hours", type=int, default=24, help="Lookback window in hours for event logs (default: 24)")
    parser.add_argument("--state-file", default=None, help="Path to state tracking file (default: client_state.json next to config)")
    parser.add_argument("--no-state", action="store_true", help="Disable state tracking and send all events matching lookback window")
    args = parser.parse_args()

    try:
        config = load_config(args.config)
    except Exception as e:
        print(f"[!] Configuration error: {e}")
        sys.exit(1)

    state_path = get_state_path(args.config, args.state_file)
    state = None if args.no_state else load_state(state_path)

    machine_id = get_machine_identifier()
    print(f"[*] Edge Forwarder Node: {machine_id}")
    if state and "last_record_number" in state:
        print(f"[*] State tracking active: resuming from record #{state['last_record_number']} (state file: {state_path})")
    elif not args.no_state:
        print(f"[*] State tracking initialized (state file: {state_path})")

    print(f"[*] Scanning Windows Application event log for unsent entries (last {args.hours} hours)...")
    candidate_logs = get_recent_windows_logs(hours=args.hours, state=state)
    print(f"[*] Found {len(candidate_logs)} unsent candidate entries (INFO to ERROR, entries > {args.hours}h and already-sent skipped).")

    if not candidate_logs:
        print("[*] No new unsent events to transmit.")
        if state is not None:
            commit_state(state, state_path)
        return

    try:
        resp = send_encrypted_logs_over_socket(config, candidate_logs)
        if state is not None and resp and resp.get("status") == "success":
            commit_state(state, state_path)
            print(f"[+] State successfully committed to {state_path}")
    except Exception as e:
        print(f"[!] Failed to stream logs to server: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()