"""
ARIA Host-Agent — Direktzugriff auf EINEN Rechner.

Laeuft direkt auf dem Ziel-Rechner (Linux) und verbindet sich AUSGEHEND als
RVS-Client in Stefans Raum (gleicher Token). Damit kann ARIA diesen Rechner
direkt steuern, auch wenn er sonst aus dem Netz nicht erreichbar ist (hinter
NAT/Firewall, kein offener Port). Anders als der Satellit (der ein LAN
entdeckt/steuert) ist beim Agent das "Geraet" der Rechner selbst.

Als reine Binary verteilbar (PyInstaller onefile) + .env fuer die Zugangsdaten.

Faehigkeiten (host_command → host_result):
  exec        Shell-Kommando ausfuehren (optional sudo)
  read/write  Datei lesen/schreiben (Base64)
  info        OS / CPU / RAM / Disk / Uptime / Netz
  screenshot  Bildschirmfoto (X11/Wayland, wenn grafische Session da ist)

Sicherheit: reagiert nur auf den eigenen RVS-Raum (Token) und nur, wenn
CONTROL_ENABLED=true. Alles wird geloggt. Keine offenen Ports.
"""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import os
import platform
import re
import shutil
import socket
import subprocess
import time
from pathlib import Path

import websockets

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [host-agent] %(levelname)s %(message)s",
)
logger = logging.getLogger("host-agent")


def _load_dotenv() -> None:
    """Laedt eine .env neben der Binary/dem Script (oder im CWD) in os.environ.
    Bereits gesetzte Werte gewinnen. Kein python-dotenv noetig."""
    here = os.path.dirname(os.path.abspath(__file__))
    for path in (os.path.join(here, ".env"), os.path.join(os.getcwd(), ".env")):
        if not os.path.isfile(path):
            continue
        try:
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith("#") or "=" not in line:
                        continue
                    key, _, val = line.partition("=")
                    key = key.strip()
                    if key.startswith("export "):
                        key = key[len("export "):].strip()
                    val = val.strip()
                    if val[:1] in ("'", '"'):
                        q = val[0]
                        end = val.find(q, 1)
                        val = val[1:end] if end != -1 else val[1:]
                    else:
                        m = re.search(r"\s+#", val)
                        if m:
                            val = val[:m.start()]
                        val = val.strip()
                    if key and key not in os.environ:
                        os.environ[key] = val
        except Exception as exc:
            logger.warning(".env laden fehlgeschlagen (%s): %s", path, exc)
        break


_load_dotenv()


# ─── Konfiguration ──────────────────────────────────────────────────

def _env_bool(name: str, default: bool) -> bool:
    v = os.environ.get(name)
    if v is None:
        return default
    return v.strip().lower() in ("1", "true", "yes", "on", "ja")


def _default_id() -> str:
    host = socket.gethostname() or "host"
    slug = re.sub(r"[^a-zA-Z0-9_-]+", "-", host).strip("-").lower()
    return slug or "host"


RVS_HOST = os.environ.get("RVS_HOST", "")
RVS_PORT = int(os.environ.get("RVS_PORT", "443") or "443")
RVS_TLS = _env_bool("RVS_TLS", True)
RVS_TOKEN = os.environ.get("RVS_TOKEN", "")
# TLS-Hostname (SNI + Zertifikatspruefung), falls RVS_HOST eine IP ist — z.B. der
# Agent laeuft im selben Netz wie der RVS und verbindet direkt auf die interne IP,
# das Caddy-Zertifikat gilt aber fuer den Namen. Dann: RVS_HOST=<interne-ip>,
# RVS_SNI=<zert-name>. Leer = SNI = RVS_HOST (Normalfall).
RVS_SNI = os.environ.get("RVS_SNI", "").strip()

# In-Process-DNS-Override: wenn RVS_SNI gesetzt ist, verbindet die URI ueber den
# HOSTNAMEN (Host-Header + SNI + Cert stimmen), waehrend getaddrinfo den Namen auf
# die echte IP in RVS_HOST aufloest. Versionsunabhaengig — host=/port= kollidiert
# in der Legacy-websockets-API mit dem aus der URI abgeleiteten Host.
if RVS_TLS and RVS_SNI:
    import socket as _socket
    _orig_getaddrinfo = _socket.getaddrinfo
    def _sni_getaddrinfo(host, *a, **k):
        return _orig_getaddrinfo(RVS_HOST if host == RVS_SNI else host, *a, **k)
    _socket.getaddrinfo = _sni_getaddrinfo


HOST_ID = (os.environ.get("HOST_ID") or _default_id()).strip()
HOST_NAME = (os.environ.get("HOST_NAME") or HOST_ID).strip()

# Steuerung ist der Sinn des Agents — aber bewusst opt-in (Sicherheit).
CONTROL_ENABLED = _env_bool("CONTROL_ENABLED", False)

# sudo: 1) Agent laeuft als root -> direkt. 2) SUDO_PASSWORD gesetzt -> sudo -S.
# 3) SUDO_NOPASSWD=true (Live-ISO / NOPASSWD-sudoers) -> sudo -n. 4) sonst Fehler.
SUDO_PASSWORD = os.environ.get("SUDO_PASSWORD", "")
SUDO_NOPASSWD = _env_bool("SUDO_NOPASSWD", False)

EXEC_TIMEOUT = float(os.environ.get("EXEC_TIMEOUT", "60") or "60")
# Ausgabe-Fenster (wie beim Satelliten): grosse stdout klein halten fuers RVS.
OUT_MAX_CHARS = int(os.environ.get("OUT_MAX_CHARS", "20000") or "20000")
OUT_MAX_CHARS_HARD = int(os.environ.get("OUT_MAX_CHARS_HARD", "200000") or "200000")
# Datei-Transfer-Limit (Base64 durchs RVS).
FILE_MAX_BYTES = int(os.environ.get("FILE_MAX_BYTES", str(10 * 1024 * 1024)) or str(10 * 1024 * 1024))

HEARTBEAT_SEC = 25
CAPS = ["exec", "read", "write", "info", "screenshot"]


# ─── Text-/Zahl-Helfer ──────────────────────────────────────────────

def _to_int(s):
    try:
        return int(str(s).strip())
    except (TypeError, ValueError):
        return None


def _to_float(s):
    try:
        return float(str(s).strip())
    except (TypeError, ValueError):
        return None


def _window_text(text: str, params: dict) -> dict:
    """contains-Zeilenfilter + offset/max_chars-Fenster. Gibt body + Metadaten."""
    total = len(text)
    contains = params.get("contains")
    if contains:
        terms = [contains] if isinstance(contains, str) else list(contains)
        terms = [str(t).lower() for t in terms if str(t).strip()]
        if terms:
            text = "\n".join(ln for ln in text.splitlines()
                             if any(t in ln.lower() for t in terms))
    offset = max(0, _to_int(params.get("offset")) or 0)
    max_chars = _to_int(params.get("max_chars")) or OUT_MAX_CHARS
    max_chars = max(1, min(max_chars, OUT_MAX_CHARS_HARD))
    body = text[offset:offset + max_chars]
    return {"body": body, "total_chars": total, "filtered": bool(contains),
            "offset": offset, "returned_chars": len(body),
            "truncated": offset + len(body) < len(text)}


# ─── Aktionen ───────────────────────────────────────────────────────

def _wrap_sudo(cmd: str, use_sudo: bool):
    """Gibt (argv, stdin_data) oder (None, fehlertext) wenn sudo nicht moeglich."""
    if not use_sudo or os.geteuid() == 0:
        return ["bash", "-lc", cmd], None
    if SUDO_PASSWORD:
        return ["sudo", "-S", "-p", "", "bash", "-lc", cmd], SUDO_PASSWORD + "\n"
    if SUDO_NOPASSWD:
        return ["sudo", "-n", "bash", "-lc", cmd], None
    return None, ("sudo verlangt ein Passwort. Setze SUDO_PASSWORD in der .env, "
                  "oder SUDO_NOPASSWD=true (Live-ISO / passwortloses sudo), oder "
                  "starte den Agent als root.")


def _do_exec(params: dict) -> dict:
    cmd = params.get("cmd") or params.get("command") or ""
    if not cmd:
        return {"ok": False, "error": "cmd (Kommando) erforderlich."}
    argv, stdin_data = _wrap_sudo(cmd, bool(params.get("sudo")))
    if argv is None:
        return {"ok": False, "error": stdin_data}
    timeout = _to_float(params.get("timeout")) or EXEC_TIMEOUT
    try:
        r = subprocess.run(argv, input=stdin_data, capture_output=True,
                           text=True, timeout=timeout)
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": f"Timeout nach {timeout:.0f}s."}
    except Exception as exc:
        return {"ok": False, "error": f"exec fehlgeschlagen: {exc}"}
    win = _window_text(r.stdout, params)
    return {"ok": True, "result": {"exit_code": r.returncode,
                                   "stderr": (r.stderr or "")[:4000], **win}}


def _do_read(params: dict) -> dict:
    path = params.get("path") or ""
    if not path:
        return {"ok": False, "error": "path erforderlich."}
    p = Path(path).expanduser()
    if not p.is_file():
        return {"ok": False, "error": f"Datei nicht gefunden: {path}"}
    size = p.stat().st_size
    offset = max(0, _to_int(params.get("offset")) or 0)
    max_bytes = _to_int(params.get("max_bytes")) or FILE_MAX_BYTES
    max_bytes = max(1, min(max_bytes, FILE_MAX_BYTES))
    try:
        with p.open("rb") as f:
            f.seek(offset)
            data = f.read(max_bytes)
    except Exception as exc:
        return {"ok": False, "error": f"Lesen fehlgeschlagen: {exc}"}
    return {"ok": True, "result": {
        "path": str(p), "size": size, "offset": offset,
        "returned_bytes": len(data), "truncated": offset + len(data) < size,
        "base64": base64.b64encode(data).decode("ascii"),
    }}


def _do_write(params: dict) -> dict:
    path = params.get("path") or ""
    if not path:
        return {"ok": False, "error": "path erforderlich."}
    b64 = params.get("base64")
    text = params.get("text")
    if b64 is None and text is None:
        return {"ok": False, "error": "base64 ODER text erforderlich."}
    try:
        data = base64.b64decode(b64) if b64 is not None else str(text).encode("utf-8")
    except Exception as exc:
        return {"ok": False, "error": f"base64 ungueltig: {exc}"}
    p = Path(path).expanduser()
    try:
        p.parent.mkdir(parents=True, exist_ok=True)
        mode = "ab" if params.get("append") else "wb"
        with p.open(mode) as f:
            f.write(data)
        if params.get("chmod"):
            os.chmod(p, int(str(params["chmod"]), 8))
    except Exception as exc:
        return {"ok": False, "error": f"Schreiben fehlgeschlagen: {exc}"}
    return {"ok": True, "result": {"path": str(p), "bytes": len(data)}}


def _do_info(params: dict) -> dict:
    info = {
        "host": HOST_NAME, "hostname": socket.gethostname(),
        "os": platform.platform(), "kernel": platform.release(),
        "arch": platform.machine(), "python": platform.python_version(),
        "user": os.environ.get("USER") or "", "is_root": os.geteuid() == 0,
    }
    try:
        import psutil
        info["cpu_percent"] = psutil.cpu_percent(interval=0.3)
        info["cpu_count"] = psutil.cpu_count()
        vm = psutil.virtual_memory()
        info["ram_used_mb"] = round(vm.used / 1024 / 1024)
        info["ram_total_mb"] = round(vm.total / 1024 / 1024)
        info["ram_percent"] = vm.percent
        du = psutil.disk_usage("/")
        info["disk_used_gb"] = round(du.used / 1024 / 1024 / 1024, 1)
        info["disk_total_gb"] = round(du.total / 1024 / 1024 / 1024, 1)
        info["disk_percent"] = du.percent
        info["uptime_s"] = round(time.time() - psutil.boot_time())
        info["load_avg"] = list(os.getloadavg()) if hasattr(os, "getloadavg") else None
    except Exception:
        # Fallback ohne psutil: das Noetigste aus os/shutil.
        try:
            info["load_avg"] = list(os.getloadavg())
        except Exception:
            info["load_avg"] = None
        try:
            total, used, free = shutil.disk_usage("/")
            info["disk_used_gb"] = round(used / 1024 ** 3, 1)
            info["disk_total_gb"] = round(total / 1024 ** 3, 1)
        except Exception:
            pass
    # Primaere IP (best effort).
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        info["primary_ip"] = s.getsockname()[0]
        s.close()
    except Exception:
        info["primary_ip"] = ""
    return {"ok": True, "result": info}


def _do_screenshot(params: dict) -> dict:
    """Bildschirmfoto via System-Tool (Wayland: grim; X11: scrot/maim/import/
    gnome-screenshot). Braucht eine aktive grafische Session (DISPLAY/WAYLAND)."""
    import tempfile
    tmp = os.path.join(tempfile.gettempdir(), f"aria_shot_{int(time.time())}.png")
    wayland = bool(os.environ.get("WAYLAND_DISPLAY"))
    candidates = []
    if wayland and shutil.which("grim"):
        candidates.append(["grim", tmp])
    for tool, argv in (("scrot", ["scrot", "-o", tmp]),
                       ("maim", ["maim", tmp]),
                       ("gnome-screenshot", ["gnome-screenshot", "-f", tmp]),
                       ("import", ["import", "-window", "root", tmp])):
        if shutil.which(tool):
            candidates.append(argv)
    if not candidates:
        return {"ok": False, "error":
                "Kein Screenshot-Tool gefunden. Installiere grim (Wayland) oder "
                "scrot/maim (X11)."}
    last_err = ""
    for argv in candidates:
        try:
            r = subprocess.run(argv, capture_output=True, text=True, timeout=15)
            if r.returncode == 0 and os.path.isfile(tmp) and os.path.getsize(tmp) > 0:
                with open(tmp, "rb") as f:
                    b = f.read()
                os.remove(tmp)
                return {"ok": True, "result": {"format": "png", "bytes": len(b),
                                               "base64": base64.b64encode(b).decode("ascii")}}
            last_err = (r.stderr or r.stdout or "").strip()[:200]
        except Exception as exc:
            last_err = str(exc)[:200]
    return {"ok": False, "error": f"Screenshot fehlgeschlagen ({last_err}). "
            "Laeuft der Agent in derselben grafischen Session?"}


ACTIONS = {
    "exec": _do_exec, "read": _do_read, "write": _do_write,
    "info": _do_info, "screenshot": _do_screenshot,
}


# ─── RVS-Client ─────────────────────────────────────────────────────

class HostAgent:
    def __init__(self) -> None:
        self.ws = None

    def _for_me(self, payload: dict) -> bool:
        """Command gilt uns, wenn kein host-Feld gesetzt ist (Broadcast) oder es
        auf unsere ID/Name passt."""
        target = (payload.get("host") or payload.get("hostId") or "").strip()
        if not target:
            return True
        return target.lower() in (HOST_ID.lower(), HOST_NAME.lower())

    async def _send(self, message: dict) -> None:
        if self.ws is None:
            return
        try:
            await self.ws.send(json.dumps(message))
        except Exception as exc:
            logger.warning("Senden fehlgeschlagen: %s", exc)

    async def _hello(self, log: bool = False) -> None:
        if log:
            logger.info("host_hello: id=%s name=%s caps=%s control=%s",
                        HOST_ID, HOST_NAME, ",".join(CAPS), CONTROL_ENABLED)
        await self._send({"type": "host_hello", "payload": {
            "hostId": HOST_ID, "name": HOST_NAME, "os": platform.platform(),
            "caps": CAPS, "control": CONTROL_ENABLED,
        }, "timestamp": int(time.time() * 1000)})

    async def _heartbeat(self) -> None:
        while True:
            await asyncio.sleep(HEARTBEAT_SEC)
            await self._send({"type": "host_ping", "payload": {"hostId": HOST_ID},
                              "timestamp": int(time.time() * 1000)})
            await self._hello()

    async def _handle(self, raw: str) -> None:
        try:
            msg = json.loads(raw)
        except Exception:
            return
        if msg.get("type") != "host_command":
            return
        payload = msg.get("payload") or {}
        if not self._for_me(payload):
            return
        req_id = payload.get("requestId", "")
        action = (payload.get("action") or "").strip()
        params = payload.get("params") or {}
        if not CONTROL_ENABLED:
            result = {"ok": False, "error": "Steuerung ist deaktiviert (CONTROL_ENABLED=false)."}
        elif action not in ACTIONS:
            result = {"ok": False, "error": f"Aktion '{action}' unbekannt (bekannt: {', '.join(ACTIONS)})."}
        else:
            logger.info("[cmd] %s params=%s", action,
                        {k: str(v)[:60] for k, v in params.items() if k not in ("base64",)})
            loop = asyncio.get_event_loop()
            try:
                result = await loop.run_in_executor(None, ACTIONS[action], params)
            except Exception as exc:
                result = {"ok": False, "error": f"{action} fehlgeschlagen: {exc}"}
        await self._send({"type": "host_result",
                          "payload": {"requestId": req_id, "hostId": HOST_ID,
                                      "action": action, **result},
                          "timestamp": int(time.time() * 1000)})

    async def run(self) -> None:
        if not RVS_HOST or not RVS_TOKEN:
            logger.error("RVS_HOST und RVS_TOKEN sind Pflicht (siehe .env.example).")
            return
        backoff = 1
        # RVS_SNI gesetzt (Verbindung auf eine interne IP): die URI nutzt den
        # HOSTNAMEN, damit Host-Header + SNI + Cert-Pruefung stimmen (sonst weist
        # Caddy die Verbindung ab -> HTTP 200 statt WS-Upgrade). Der TCP-Connect
        # wird per host=/port= auf die IP in RVS_HOST umgebogen.
        proto = "wss" if RVS_TLS else "ws"
        uri_host = RVS_SNI if (RVS_TLS and RVS_SNI) else RVS_HOST
        url = f"{proto}://{uri_host}:{RVS_PORT}?token={RVS_TOKEN}"
        connect_kwargs = {"max_size": 16 * 1024 * 1024,
                          "ping_interval": 20, "ping_timeout": 20}
        while True:
            try:
                logger.info("Verbinde mit RVS %s://%s:%s%s …", proto, uri_host, RVS_PORT,
                            f" (TCP {RVS_HOST})" if (RVS_TLS and RVS_SNI) else "")
                async with websockets.connect(url, **connect_kwargs) as ws:
                    self.ws = ws
                    backoff = 1
                    await self._hello(log=True)
                    hb = asyncio.create_task(self._heartbeat())
                    try:
                        async for raw in ws:
                            await self._handle(raw)
                    finally:
                        hb.cancel()
            except Exception as exc:
                logger.warning("RVS-Verbindung verloren: %s", exc)
            finally:
                self.ws = None
            await asyncio.sleep(backoff)
            backoff = min(backoff * 2, 30)


def main() -> None:
    logger.info("ARIA Host-Agent startet — id=%s name=%s control=%s root=%s",
                HOST_ID, HOST_NAME, CONTROL_ENABLED, os.geteuid() == 0)
    try:
        asyncio.run(HostAgent().run())
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()