#!/usr/bin/env python3
"""Kana for Linux: a tiny agent that makes true local statements and joins the team screen."""

import argparse
import datetime
import hashlib
import json
import os
import platform
import re
import shutil
import socket
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from pathlib import Path

AGENT_VERSION = "0.2.3"


def default_base_url() -> str:
    channel_base = "https://staging.kana.io" if os.environ.get("KANA_CHANNEL") == "staging" else "https://app.kana.io"
    return os.environ.get("KANA_SERVER_URL") or channel_base


DEFAULT_BASE_URL = default_base_url()
STATE_PATH = Path(os.environ.get("KANA_STATE_PATH", Path.home() / ".config" / "kana" / "state.json"))
LOGIND_CONF_PATH = Path("/etc/systemd/logind.conf")
LOGIND_CONF_D_DIR = Path("/etc/systemd/logind.conf.d")
APT_AUTO_UPGRADES_PATH = Path("/etc/apt/apt.conf.d/20auto-upgrades")
CLAMAV_MAX_FILE_SIZE = 50 * 1024 * 1024
CLAMAV_MAX_SCAN_SIZE = 200 * 1024 * 1024
CLAMAV_TIMEOUT_SECONDS = 120


class EnrollmentError(Exception):
    pass


def normalize_base_url(url: str) -> str:
    url = url.strip()
    while url.endswith("/"):
        url = url[:-1]
    return url


def parse_enrollment(raw_input: str, default_base: str = None) -> dict:
    s = raw_input.strip()
    if not s:
        raise EnrollmentError("Enrollment input cannot be empty")
    fallback_base = normalize_base_url(default_base or default_base_url())

    if s.lower().startswith("kana://"):
        parsed = urllib.parse.urlparse(s)
        if (parsed.netloc or "") != "enroll" and (parsed.path or "") not in ("enroll", "/enroll"):
            raise EnrollmentError("Unsupported kana action; expected kana://enroll")
        query = urllib.parse.parse_qs(parsed.query)
        codes = query.get("code")
        if not codes or not codes[0].strip():
            raise EnrollmentError("Missing 'code' query parameter in kana:// link")
        token = codes[0].strip()
        bases = query.get("base")
        base = normalize_base_url(bases[0].strip()) if (bases and bases[0].strip()) else fallback_base
        return {"token": token, "base_url": base}

    if s.lower().startswith(("https://", "http://")):
        parsed = urllib.parse.urlparse(s)
        if not parsed.netloc:
            raise EnrollmentError("Invalid URL: missing host")
        query = urllib.parse.parse_qs(parsed.query)
        codes = query.get("code")
        if not codes or not codes[0].strip():
            raise EnrollmentError("Missing 'code' query parameter in enrollment URL")
        return {"token": codes[0].strip(), "base_url": normalize_base_url(f"{parsed.scheme}://{parsed.netloc}")}

    if s.startswith("KANA|"):
        parts = s.split("|")
        if len(parts) != 3 or not parts[1].strip() or not parts[2].strip():
            raise EnrollmentError("KANA code must have format KANA|<token>|<base>")
        parsed = urllib.parse.urlparse(parts[2].strip())
        if not parsed.scheme or not parsed.netloc:
            raise EnrollmentError(f"Invalid base URL in KANA code: '{parts[2].strip()}'")
        return {"token": parts[1].strip(), "base_url": normalize_base_url(parts[2].strip())}

    if any(c in s for c in (" ", "\t", "\n", "|")):
        raise EnrollmentError("Raw token contains invalid characters or whitespace")
    if "://" in s:
        raise EnrollmentError(f"Unknown URL scheme in '{s}'")
    return {"token": s, "base_url": fallback_base}


def check_disk_encryption() -> dict:
    lsblk_bin = shutil.which("lsblk")
    if not lsblk_bin:
        return {"state": "unknown", "source": "lsblk", "detail": "lsblk not found"}
    try:
        findmnt_bin = shutil.which("findmnt")
        root_dev = None
        if findmnt_bin:
            res = subprocess.run([findmnt_bin, "-no", "SOURCE", "/"], capture_output=True, text=True, timeout=5)
            if res.returncode == 0 and res.stdout.strip():
                root_dev = res.stdout.strip()

        if root_dev and root_dev.startswith("/dev"):
            res = subprocess.run([lsblk_bin, "-s", "-no", "TYPE,NAME", root_dev], capture_output=True, text=True, timeout=5)
            if res.returncode == 0:
                for line in (l.strip() for l in res.stdout.splitlines() if l.strip()):
                    if line.split() and line.split()[0].lower() == "crypt":
                        return {"state": "enabled", "source": "lsblk", "detail": f"crypt device found in / chain: {line}"}
                return {"state": "disabled", "source": "lsblk", "detail": f"no crypt device in chain for {root_dev}"}

        res = subprocess.run([lsblk_bin, "-no", "TYPE,NAME"], capture_output=True, text=True, timeout=5)
        if res.returncode == 0:
            for line in (l.strip() for l in res.stdout.splitlines() if l.strip()):
                if line.split() and line.split()[0].lower() == "crypt":
                    return {"state": "enabled", "source": "lsblk", "detail": f"crypt device found: {line}"}
            return {"state": "disabled", "source": "lsblk", "detail": "no crypt device found"}
        return {"state": "unknown", "source": "lsblk", "detail": f"lsblk failed with exit code {res.returncode}"}
    except Exception as e:
        return {"state": "unknown", "source": "lsblk", "detail": f"disk encryption check error: {e}"}


def check_firewall() -> dict:
    tried = []
    ufw_bin = shutil.which("ufw")
    if ufw_bin:
        tried.append("ufw")
        try:
            res = subprocess.run([ufw_bin, "status"], capture_output=True, text=True, timeout=5)
            if "status: active" in res.stdout.lower():
                return {"state": "enabled", "source": "ufw", "detail": "ufw is active"}
        except Exception:
            pass

    systemctl_bin = shutil.which("systemctl")
    if systemctl_bin:
        tried.append("firewalld")
        try:
            res = subprocess.run([systemctl_bin, "is-active", "firewalld"], capture_output=True, text=True, timeout=5)
            if res.returncode == 0 and res.stdout.strip().lower() == "active":
                return {"state": "enabled", "source": "firewalld", "detail": "firewalld is active"}
        except Exception:
            pass

    nft_bin = shutil.which("nft")
    if nft_bin:
        tried.append("nft")
        try:
            res = subprocess.run([nft_bin, "list", "ruleset"], capture_output=True, text=True, timeout=5)
            if res.returncode == 0:
                out = res.stdout.lower()
                if "type filter hook input" in out or "hook input" in out or ("policy drop" in out and "input" in out):
                    return {"state": "enabled", "source": "nft", "detail": "nftables input chain configured"}
        except Exception:
            pass

    if not tried:
        return {"state": "unknown", "source": "firewall", "detail": "no firewall tools available (tried: ufw, firewalld, nft)"}
    return {"state": "disabled", "source": "firewall", "detail": f"firewall inactive (checked: {', '.join(tried)})"}


def check_screen_lock() -> dict:
    tried = []
    gsettings_state = None
    gsettings_detail = None

    gsettings_bin = shutil.which("gsettings")
    if gsettings_bin:
        tried.append("gsettings")
        try:
            res = subprocess.run(
                [gsettings_bin, "get", "org.gnome.desktop.screensaver", "lock-enabled"],
                capture_output=True,
                text=True,
                timeout=5,
            )
            if res.returncode == 0:
                val = res.stdout.strip().strip("'\"").lower()
                if val == "true":
                    gsettings_state = "enabled"
                    gsettings_detail = "gsettings lock-enabled is true"
                elif val == "false":
                    gsettings_state = "disabled"
                    gsettings_detail = "gsettings lock-enabled is false"
        except Exception:
            pass

    if gsettings_state == "enabled":
        return {"state": "enabled", "source": "gsettings", "detail": gsettings_detail}

    logind_present = False
    logind_idle_action = None
    conf_files = []

    p_conf = Path(LOGIND_CONF_PATH)
    if p_conf.exists():
        logind_present = True
        conf_files.append(p_conf)

    p_dir = Path(LOGIND_CONF_D_DIR)
    if p_dir.exists() and p_dir.is_dir():
        logind_present = True
        try:
            for f in sorted(p_dir.glob("*.conf")):
                if f.is_file():
                    conf_files.append(f)
        except Exception:
            pass

    if logind_present:
        tried.append(str(p_conf))
        for f in conf_files:
            try:
                with open(f, "r", encoding="utf-8", errors="replace") as fh:
                    for line in fh:
                        line = line.strip()
                        if not line or line.startswith(("#", ";")):
                            continue
                        if "=" in line:
                            k, v = line.split("=", 1)
                            if k.strip() == "IdleAction":
                                logind_idle_action = v.strip().lower()
            except Exception:
                pass

        if logind_idle_action == "lock":
            return {"state": "enabled", "source": "logind", "detail": "logind IdleAction=lock"}

    if gsettings_state == "disabled":
        return {"state": "disabled", "source": "gsettings", "detail": gsettings_detail}

    if logind_present:
        action_desc = logind_idle_action if logind_idle_action else "empty"
        return {"state": "disabled", "source": "logind", "detail": f"logind IdleAction is {action_desc}"}

    detail_str = f"neither gsettings nor logind exists (tried: {', '.join(tried) if tried else 'gsettings, ' + str(p_conf)})"
    return {"state": "unknown", "source": "screen_lock", "detail": detail_str}


def check_automatic_updates() -> dict:
    tried = []
    unattended_upgrades_state = None
    apt_file_state = None
    apt_file_present = False
    dnf_state = None

    systemctl_bin = shutil.which("systemctl")
    if systemctl_bin:
        tried.append("systemctl")
        try:
            res = subprocess.run(
                [systemctl_bin, "is-enabled", "unattended-upgrades"],
                capture_output=True,
                text=True,
                timeout=5,
            )
            out = res.stdout.strip().lower()
            if out == "enabled":
                unattended_upgrades_state = "enabled"
            elif out in ("disabled", "static", "indirect", "masked"):
                unattended_upgrades_state = "disabled"
        except Exception:
            pass

        for timer in ("dnf-automatic.timer", "dnf-automatic-install.timer"):
            try:
                res = subprocess.run(
                    [systemctl_bin, "is-enabled", timer],
                    capture_output=True,
                    text=True,
                    timeout=5,
                )
                out = res.stdout.strip().lower()
                if out == "enabled":
                    dnf_state = ("enabled", timer)
                    break
                elif out in ("disabled", "static", "indirect", "masked") and dnf_state is None:
                    dnf_state = ("disabled", timer)
            except Exception:
                pass

    apt_file = Path(APT_AUTO_UPGRADES_PATH)
    if apt_file.exists():
        apt_file_present = True
        tried.append(str(apt_file))
        try:
            content = apt_file.read_text(encoding="utf-8", errors="replace")
            if re.search(r'APT::Periodic::Unattended-Upgrade\s+"1"', content):
                apt_file_state = "enabled"
            else:
                apt_file_state = "disabled"
        except Exception:
            pass

    if unattended_upgrades_state == "enabled":
        return {"state": "enabled", "source": "unattended-upgrades", "detail": "systemctl is-enabled unattended-upgrades is enabled"}
    if apt_file_state == "enabled":
        return {"state": "enabled", "source": "unattended-upgrades", "detail": f"{apt_file} has APT::Periodic::Unattended-Upgrade 1"}
    if dnf_state and dnf_state[0] == "enabled":
        return {"state": "enabled", "source": "dnf-automatic", "detail": f"systemctl is-enabled {dnf_state[1]} is enabled"}

    if unattended_upgrades_state == "disabled":
        return {"state": "disabled", "source": "unattended-upgrades", "detail": "systemctl is-enabled unattended-upgrades is disabled"}
    if apt_file_present and apt_file_state == "disabled":
        return {"state": "disabled", "source": "unattended-upgrades", "detail": f"{apt_file} present but Unattended-Upgrade not 1"}
    if dnf_state and dnf_state[0] == "disabled":
        return {"state": "disabled", "source": "dnf-automatic", "detail": f"systemctl is-enabled {dnf_state[1]} is disabled"}

    tried_items = ["unattended-upgrades", str(apt_file), "dnf-automatic"]
    return {"state": "unknown", "source": "automatic_updates", "detail": f"no automatic updates mechanism found (tried: {', '.join(tried_items)})"}


def check_launch_at_login() -> dict:
    tried = ["systemctl --user is-enabled kana.timer", "systemctl --user is-enabled kana.service"]
    systemctl_bin = shutil.which("systemctl")
    if not systemctl_bin:
        return {
            "state": "unknown",
            "source": "systemctl",
            "detail": f"systemctl not found (tried: {', '.join(tried)})",
        }

    units = ["kana.timer", "kana.service"]
    user_queryable = False

    for unit in units:
        try:
            res = subprocess.run(
                [systemctl_bin, "--user", "is-enabled", unit],
                capture_output=True,
                text=True,
                timeout=5,
            )
            out = res.stdout.strip().lower()
            err = res.stderr.strip().lower()

            if "failed to connect to bus" in err or "no medium found" in err:
                continue

            if out == "enabled":
                return {
                    "state": "enabled",
                    "source": "systemctl",
                    "detail": f"systemctl --user is-enabled {unit} is enabled",
                }

            if out in ("disabled", "static", "indirect", "masked", "not-found", "linked", "bad"):
                user_queryable = True
            elif res.returncode == 0:
                user_queryable = True
            elif "not-found" in err or "no such file or directory" in err:
                user_queryable = True
        except Exception:
            pass

    if user_queryable:
        return {
            "state": "disabled",
            "source": "systemctl",
            "detail": "systemd user units kana.timer and kana.service are not enabled",
        }

    return {
        "state": "unknown",
        "source": "systemctl",
        "detail": f"systemctl --user cannot be queried (tried: {', '.join(tried)})",
    }


def get_os_version() -> str:
    for path in ("/etc/os-release", "/usr/lib/os-release"):
        if os.path.exists(path):
            try:
                with open(path, "r", encoding="utf-8", errors="replace") as f:
                    for line in f:
                        if line.startswith("PRETTY_NAME="):
                            val = line.split("=", 1)[1].strip().strip('"\'')
                            if val:
                                return val
            except Exception:
                pass
    return platform.platform() or "Linux"


def _clamav_summary(output: str) -> dict:
    values = {"scanned_count": 0, "infected_count": 0, "errors_count": 0}
    fields = {
        "scanned files": "scanned_count",
        "infected files": "infected_count",
        "total errors": "errors_count",
    }
    for line in output.splitlines():
        if ":" not in line:
            continue
        label, raw_value = line.split(":", 1)
        key = fields.get(label.strip().lower())
        match = re.search(r"\d+", raw_value)
        if key and match:
            values[key] = min(int(match.group(0)), 1_000_000)
    return values


def _clamav_version(output: str) -> tuple:
    parts = output.strip().split("/")
    engine = "unknown"
    signatures = "unknown"
    if parts and parts[0].startswith("ClamAV "):
        engine = parts[0][7:][:80] or "unknown"
    if len(parts) > 1 and re.fullmatch(r"\d+", parts[1]):
        signatures = parts[1][:80]
    return engine, signatures


def _clamav_definition_evidence(now: datetime.datetime) -> tuple:
    configured = os.environ.get("CLAMAV_DB_DIR")
    roots = [Path(configured)] if configured else [Path("/var/lib/clamav"), Path("/var/lib/clamav-data")]
    newest = None
    newest_name = "unknown"
    for root in roots:
        try:
            for pattern in ("*.cvd", "*.cld"):
                for candidate in root.glob(pattern):
                    modified = candidate.stat().st_mtime
                    if newest is None or modified > newest:
                        newest = modified
                        newest_name = candidate.stem[:80] or "unknown"
        except (OSError, ValueError):
            continue
    if newest is None:
        return "unknown", "unknown"
    age = now.timestamp() - newest
    state = "fresh" if -86_400 <= age < 7 * 86_400 else "stale"
    return newest_name, state


def _clamav_local_counts(scan_root: Path) -> tuple:
    scanned = 0
    skipped = 0
    errors = 0
    total_bytes = 0
    try:
        candidates = scan_root.rglob("*")
        for candidate in candidates:
            try:
                if candidate.is_symlink() or not candidate.is_file():
                    continue
                size = candidate.stat().st_size
                if size > CLAMAV_MAX_FILE_SIZE or total_bytes + size > CLAMAV_MAX_SCAN_SIZE:
                    skipped += 1
                    continue
                total_bytes += size
                scanned += 1
            except OSError:
                errors += 1
    except OSError:
        errors += 1
    return min(scanned, 1_000_000), min(skipped, 1_000_000), min(errors, 1_000_000)


def _clamav_detections(output: str, scan_root: Path) -> list:
    detections = []
    resolved_root = scan_root.resolve()
    for line in output.splitlines():
        if not line.endswith(" FOUND") or ": " not in line:
            continue
        path_text, signature_text = line.rsplit(": ", 1)
        signature = signature_text[:-6]
        if not re.fullmatch(r"[A-Za-z0-9_.() -]{1,200}", signature):
            continue
        try:
            candidate = Path(path_text).resolve(strict=True)
            candidate.relative_to(resolved_root)
            if not candidate.is_file() or candidate.is_symlink():
                continue
            digest = hashlib.sha256(candidate.read_bytes()).hexdigest()
        except (OSError, ValueError):
            continue
        detections.append({"sha256": digest, "signature": signature, "quarantined": False})
        if len(detections) == 100:
            break
    return detections


def collect_malware_scan(now: datetime.datetime = None) -> dict:
    executable = shutil.which("clamscan")
    if not executable:
        return {"state": "unavailable", "reason": "no_scanner"}

    now = now or datetime.datetime.now(datetime.timezone.utc)
    if now.tzinfo is None:
        now = now.replace(tzinfo=datetime.timezone.utc)
    now = now.astimezone(datetime.timezone.utc)
    observed_at = now.isoformat()
    scan_root = Path.home() / "Downloads"
    signature_version, definitions_state = _clamav_definition_evidence(now)
    report = {
        "state": "error",
        "engine": "clamav",
        "engine_version": "unknown",
        "signature_version": signature_version,
        "definitions_state": definitions_state,
        "scope": "approved_folders",
        "scope_labels": ["folder_1"],
        "observed_at": observed_at,
        "scanned_count": 0,
        "skipped_count": 0,
        "errors_count": 0,
        "incomplete": False,
        "detections": [],
    }
    if not scan_root.is_dir():
        report.update(reason="scan_root_unavailable", errors_count=1)
        return report

    try:
        version = subprocess.run(
            [executable, "--version"], capture_output=True, text=True, timeout=5
        )
    except (OSError, subprocess.TimeoutExpired):
        report.update(reason="engine_version_failed", errors_count=1)
        return report
    if version.returncode != 0:
        report.update(reason="engine_version_failed", errors_count=1)
        return report
    report["engine_version"], version_signatures = _clamav_version(version.stdout)
    if version_signatures != "unknown":
        report["signature_version"] = version_signatures

    command = [
        executable,
        "--recursive=yes",
        "--infected",
        "--no-summary",
        "--max-filesize=50M",
        "--max-scansize=200M",
        str(scan_root),
    ]
    local_scanned, local_skipped, local_errors = _clamav_local_counts(scan_root)
    try:
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            timeout=CLAMAV_TIMEOUT_SECONDS,
        )
    except subprocess.TimeoutExpired:
        report.update(
            reason="scan_timeout",
            scanned_count=local_scanned,
            skipped_count=local_skipped,
            errors_count=local_errors + 1,
            incomplete=True,
        )
        return report
    except OSError:
        report.update(reason="engine_scan_failed", errors_count=local_errors + 1)
        return report

    output = "\n".join(part for part in (result.stdout, result.stderr) if part)
    summary = _clamav_summary(output)
    report["scanned_count"] = summary["scanned_count"] or local_scanned
    report["skipped_count"] = local_skipped
    report["errors_count"] = min(summary["errors_count"] + local_errors, 1_000_000)
    report["detections"] = _clamav_detections(output, scan_root)

    if result.returncode == 0 and report["errors_count"] == 0 and report["scanned_count"] > 0:
        report["state"] = "clean"
    elif result.returncode == 1 and report["detections"]:
        report["state"] = "detected"
    elif result.returncode == 0 and report["scanned_count"] == 0:
        report["state"] = "unavailable"
        report["reason"] = "no_files_scanned"
    else:
        report["reason"] = "detection_evidence_unreadable" if result.returncode == 1 else "engine_scan_failed"
        report["errors_count"] = min(report["errors_count"] + 1, 1_000_000)
    return report


def build_security_report() -> dict:
    return {
        "type": "security_report",
        "schema_version": 1,
        "platform": "linux",
        "report_id": str(uuid.uuid4()),
        "observed_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "agent_version": AGENT_VERSION,
        "os_version": get_os_version(),
        "posture": {
            "disk_encryption": check_disk_encryption(),
            "firewall": check_firewall(),
            "screen_lock": check_screen_lock(),
            "automatic_updates": check_automatic_updates(),
            "launch_at_login": check_launch_at_login(),
            "gatekeeper": {"state": "unavailable", "source": "kana", "detail": "no Linux equivalent"},
            "sip": {"state": "unavailable", "source": "kana", "detail": "no Linux equivalent"},
        },
        "malware_scan": collect_malware_scan(),
    }


def load_state() -> dict:
    p = Path(STATE_PATH)
    if not p.exists():
        return None
    try:
        with open(p, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def save_state(state: dict):
    p = Path(STATE_PATH)
    p.parent.mkdir(parents=True, exist_ok=True)
    tmp = p.with_suffix(".tmp")
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(state, f, indent=2)
    os.chmod(tmp, 0o600)
    os.replace(tmp, p)


def enroll(code_or_link: str) -> dict:
    target = parse_enrollment(code_or_link)
    base_url = target["base_url"]
    token = target["token"]
    hostname = socket.gethostname()
    os_ver = get_os_version()

    payload = {
        "product": "kana",
        "platform": "linux",
        "enrollment_token": token,
        "hostname": hostname,
        "os_version": os_ver,
        "app_version": AGENT_VERSION,
        "device_info": {
            "bundle_id": "pro.endpoint.kana.linux",
            "hostname": hostname,
            "os": os_ver,
            "os_version": os_ver,
            "app_version": AGENT_VERSION,
        },
    }

    req = urllib.request.Request(
        f"{base_url}/endpoint/register",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json", "Accept": "application/json", "User-Agent": f"Kana-Linux/{AGENT_VERSION}"},
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            resp_data = json.loads(resp.read().decode("utf-8"))
            dev_tok = resp_data.get("device_token")
            if not dev_tok:
                raise EnrollmentError("Registration response missing device token")
            host = urllib.parse.urlparse(base_url).netloc or base_url
            state = {
                "base_url": base_url,
                "device_token": dev_tok,
                "enrolled_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                "host": host,
            }
            save_state(state)
            return {"enrolled": True, "host": host}
    except urllib.error.HTTPError as e:
        err_body = e.read().decode("utf-8", errors="replace")
        try:
            data = json.loads(err_body)
            msg = data.get("message") or data.get("error") or err_body
        except Exception:
            msg = err_body
        raise EnrollmentError(f"Server error ({e.code}): {msg}")
    except Exception as e:
        raise EnrollmentError(f"Registration failed: {e}")


def post_report() -> dict:
    state = load_state()
    if not state or not state.get("device_token"):
        raise ValueError("not_enrolled")

    base_url = state["base_url"]
    device_token = state["device_token"]
    host = state.get("host") or urllib.parse.urlparse(base_url).netloc or base_url
    report = build_security_report()

    req = urllib.request.Request(
        f"{base_url}/endpoint/security_reports",
        data=json.dumps(report).encode("utf-8"),
        headers={
            "Content-Type": "application/json",
            "Accept": "application/json",
            "X-Device-Token": device_token,
            "User-Agent": f"Kana-Linux/{AGENT_VERSION}",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            now_iso = report["observed_at"]
            state["last_reported_at"] = now_iso
            state["last_reported_host"] = host
            state.pop("last_report_error", None)
            save_state(state)
            return {"reported": True, "at": now_iso, "to": host}
    except Exception as e:
        state["last_report_error"] = str(e)
        save_state(state)
        raise


def get_status() -> dict:
    state = load_state()
    if not state or not state.get("device_token"):
        return {"enrolled": False, "status_line": "Not enrolled", "last_reported_at": None}

    host = state.get("host") or urllib.parse.urlparse(state["base_url"]).netloc or state["base_url"]
    last_reported = state.get("last_reported_at")
    if last_reported:
        try:
            dt = datetime.datetime.fromisoformat(last_reported)
            local_dt = dt.astimezone() if dt.tzinfo else dt
            time_str = local_dt.strftime("%H:%M")
        except Exception:
            time_str = last_reported[11:16] if len(last_reported) >= 16 else last_reported
        status_line = f"Reported {time_str} to {host}"
    elif state.get("last_report_error"):
        status_line = f"Not reported: {state['last_report_error']}"
    else:
        status_line = "Enrolled: no report yet"

    return {
        "enrolled": True,
        "host": host,
        "base_url": state.get("base_url"),
        "status_line": status_line,
        "last_reported_at": last_reported,
    }


def run_loop():
    while True:
        try:
            res = post_report()
            print(json.dumps(res))
        except Exception as e:
            print(json.dumps({"reported": False, "error": str(e)}), file=sys.stderr)
        time.sleep(15 * 60)


def fix_control(control: str) -> dict:
    needs_root = ["firewall", "automatic_updates", "malware_protection"]
    if control in needs_root and os.geteuid() != 0:
        executable = os.path.abspath(sys.argv[0])
        try:
            res = subprocess.run(["pkexec", executable, "--fix", control], capture_output=True, text=True)
            if res.returncode == 0 and res.stdout:
                return json.loads(res.stdout)
            return {"state": "error", "detail": f"pkexec failed: {res.stderr}"}
        except Exception as e:
            return {"state": "error", "detail": str(e)}

    if control == "firewall":
        ufw = shutil.which("ufw")
        systemctl = shutil.which("systemctl")
        nft = shutil.which("nft")
        if ufw:
            subprocess.run([ufw, "--force", "enable"], capture_output=True)
        elif systemctl:
            subprocess.run([systemctl, "enable", "--now", "firewalld"], capture_output=True)
        elif nft:
            subprocess.run([nft, "add", "rule", "inet", "filter", "input", "drop"], capture_output=True)
        return check_firewall()

    elif control == "screen_lock":
        gsettings = shutil.which("gsettings")
        if gsettings:
            subprocess.run([gsettings, "set", "org.gnome.desktop.screensaver", "lock-enabled", "true"], capture_output=True)
            subprocess.run([gsettings, "set", "org.gnome.desktop.screensaver", "lock-delay", "uint32", "0"], capture_output=True)
            subprocess.run([gsettings, "set", "org.gnome.desktop.session", "idle-delay", "uint32", "300"], capture_output=True)
        else:
            if os.geteuid() != 0:
                executable = os.path.abspath(sys.argv[0])
                res = subprocess.run(["pkexec", executable, "--fix", "screen_lock_logind"], capture_output=True, text=True)
                if res.returncode == 0 and res.stdout:
                    return json.loads(res.stdout)
        return check_screen_lock()
    
    elif control == "screen_lock_logind":
        p_dir = Path("/etc/systemd/logind.conf.d")
        if not p_dir.exists():
            p_dir.mkdir(parents=True, exist_ok=True)
        (p_dir / "kana_lock.conf").write_text("[Login]\nIdleAction=lock\nIdleActionSec=5min\n")
        systemctl = shutil.which("systemctl")
        if systemctl:
            subprocess.run([systemctl, "restart", "systemd-logind"], capture_output=True)
        return check_screen_lock()

    elif control == "automatic_updates":
        apt = shutil.which("apt-get")
        dnf = shutil.which("dnf")
        if apt:
            subprocess.run([apt, "update"], capture_output=True)
            subprocess.run([apt, "install", "-y", "unattended-upgrades"], capture_output=True)
            dpkg = shutil.which("dpkg-reconfigure")
            if dpkg:
                subprocess.run([dpkg, "-f", "noninteractive", "unattended-upgrades"], capture_output=True)
        elif dnf:
            subprocess.run([dnf, "install", "-y", "dnf-automatic"], capture_output=True)
            systemctl = shutil.which("systemctl")
            if systemctl:
                subprocess.run([systemctl, "enable", "--now", "dnf-automatic.timer"], capture_output=True)
        return check_automatic_updates()

    elif control == "malware_protection":
        apt = shutil.which("apt-get")
        dnf = shutil.which("dnf")
        if apt:
            subprocess.run([apt, "update"], capture_output=True)
            subprocess.run([apt, "install", "-y", "clamav", "clamav-daemon"], capture_output=True)
        elif dnf:
            subprocess.run([dnf, "install", "-y", "clamav", "clamd"], capture_output=True)
        freshclam = shutil.which("freshclam")
        if freshclam:
            subprocess.run([freshclam], capture_output=True)
        return collect_malware_scan()

    elif control == "starts_at_login":
        systemctl = shutil.which("systemctl")
        if systemctl:
            subprocess.run([systemctl, "--user", "enable", "--now", "kana.timer"], capture_output=True)
        return check_launch_at_login()

    elif control == "disk_encryption":
        fscrypt = shutil.which("fscrypt")
        supported = False
        if fscrypt:
            res = subprocess.run([fscrypt, "status"], capture_output=True, text=True)
            if res.returncode == 0 and "supported" in res.stdout.lower() and "not supported" not in res.stdout.lower() and "unsupported" not in res.stdout.lower():
                supported = True
        if supported:
            subprocess.run([fscrypt, "setup"], capture_output=True)
            subprocess.run([fscrypt, "encrypt", str(Path.home())], capture_output=True)
            return check_disk_encryption()
        else:
            return {"state": "pending", "detail": "Needs a reinstall with encryption; ask your admin"}

    return {"state": "error", "detail": "Unknown control"}


def main():
    parser = argparse.ArgumentParser(description="Kana Linux Agent")
    parser.add_argument("--security-report", action="store_true", help="Print SecurityReport JSON")
    parser.add_argument("--enroll", metavar="CODE", help="Enroll device with team code or link")
    parser.add_argument("--post-report", action="store_true", help="Post security report to enrolled server")
    parser.add_argument("--status", action="store_true", help="Show current enrollment and report status")
    parser.add_argument("--run", action="store_true", help="Post report now and every 15 minutes")
    parser.add_argument("--fix", metavar="CONTROL", help="Attempt to fix a failing control")
    args = parser.parse_args()

    if args.fix:
        try:
            res = fix_control(args.fix)
            print(json.dumps(res))
            sys.exit(0)
        except Exception as e:
            print(json.dumps({"state": "error", "detail": str(e)}))
            sys.exit(1)

    if args.security_report:
        print(json.dumps(build_security_report(), indent=2))
        sys.exit(0)
    if args.enroll:
        try:
            res = enroll(args.enroll)
            print(json.dumps(res))
            sys.exit(0)
        except Exception as e:
            print(json.dumps({"enrolled": False, "error": str(e)}))
            sys.exit(1)
    if args.post_report:
        try:
            res = post_report()
            print(json.dumps(res))
            sys.exit(0)
        except Exception as e:
            print(json.dumps({"status": "failed", "reason": str(e)}))
            sys.exit(1)
    if args.status:
        print(json.dumps(get_status()))
        sys.exit(0)
    if args.run:
        run_loop()

    parser.print_help()
    sys.exit(1)


if __name__ == "__main__":
    main()
