"""
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 sys
import time
from pathlib import Path

import websockets

# ─── Plattform-Weichen (Linux / macOS / Windows) ────────────────────
IS_WINDOWS = os.name == "nt"
IS_MAC = sys.platform == "darwin"


def _is_admin() -> bool:
    """root (Unix) bzw. Administrator (Windows)."""
    try:
        return os.geteuid() == 0            # Unix (Linux/macOS)
    except AttributeError:
        try:
            import ctypes
            return ctypes.windll.shell32.IsUserAnAdmin() != 0   # Windows
        except Exception:
            return False

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__))
    # PyInstaller-onefile: __file__ liegt im Temp-Extract-Dir, NICHT beim .exe/
    # Binary — deshalb zusaetzlich sys.executable-Ordner (echter Binary-Ort) und
    # das CWD (z.B. der ProgramData-Ordner, den der Windows-Dienst als AppDir nutzt).
    exe_dir = os.path.dirname(os.path.abspath(sys.executable))
    seen = set()
    candidates = [os.path.join(d, ".env") for d in (exe_dir, here, os.getcwd())]
    for path in [p for p in candidates if not (p in seen or seen.add(p))]:
        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)
# Bei TLS-Fehlschlag einmal auf ws:// zurueckfallen (wie die Compute-Bridges).
# Hilft nur, wenn der RVS plaintext erreichbar ist; gegen Caddy-TLS bleibt wss.
RVS_TLS_FALLBACK = _env_bool("RVS_TLS_FALLBACK", 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))

# Version (wird von release_agent.sh beim Release gesetzt).
AGENT_VERSION = "0.0.0.2"

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):
    """Baut die Argv (OS-abhaengige Shell) + optional Root/Admin-Rechte.
    Gibt (argv, stdin_data) oder (None, fehlertext)."""
    if IS_WINDOWS:
        # PowerShell; kein sudo. Fuer Admin-Rechte muss der Agent SELBST als
        # Administrator laufen (UAC) — dann hat 'sudo:true' bereits volle Rechte.
        if use_sudo and not _is_admin():
            return None, ("Windows kennt kein sudo. Starte den Agent als "
                          "Administrator ('Als Administrator ausfuehren'), dann "
                          "laufen Kommandos mit vollen Rechten.")
        return ["powershell", "-NoProfile", "-NonInteractive", "-Command", cmd], None
    # Unix: Linux + macOS (bash vorhanden; macOS-sudo verhaelt sich wie Linux)
    if not use_sudo or _is_admin():
        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(),
        "agent_version": AGENT_VERSION,
        "os": platform.platform(), "kernel": platform.release(),
        "arch": platform.machine(), "python": platform.python_version(),
        "user": os.environ.get("USER") or os.environ.get("USERNAME") or "",
        "is_root": _is_admin(),
    }
    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 — OS-abhaengig. Windows: PowerShell/System.Drawing;
    macOS: screencapture; Linux: grim (Wayland) / scrot/maim/import (X11).
    Braucht eine aktive grafische Session."""
    import tempfile
    tmp = os.path.join(tempfile.gettempdir(), f"aria_shot_{int(time.time())}.png")
    candidates = []
    if IS_WINDOWS:
        ps = ("Add-Type -AssemblyName System.Windows.Forms,System.Drawing;"
              "$b=[System.Windows.Forms.SystemInformation]::VirtualScreen;"
              "$bmp=New-Object System.Drawing.Bitmap $b.Width,$b.Height;"
              "$g=[System.Drawing.Graphics]::FromImage($bmp);"
              "$g.CopyFromScreen($b.Location,[System.Drawing.Point]::Empty,$b.Size);"
              f"$bmp.Save('{tmp}');$g.Dispose();$bmp.Dispose()")
        candidates.append(["powershell", "-NoProfile", "-NonInteractive", "-Command", ps])
    elif IS_MAC:
        candidates.append(["screencapture", "-x", tmp])   # -x = ohne Ton
    else:
        if os.environ.get("WAYLAND_DISPLAY") 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. Linux: grim (Wayland) oder "
                "scrot/maim (X11) installieren. (Windows/macOS nutzen Bordmittel.)"}
    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(),
            "version": AGENT_VERSION, "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
        # use_tls kann bei Fehlschlag einmal auf ws:// fallen (RVS_TLS_FALLBACK),
        # danach wieder zurueck auf RVS_TLS (kein Sticky-Fallback).
        use_tls = RVS_TLS
        tls_fallback_tried = False
        connect_kwargs = {"max_size": 16 * 1024 * 1024,
                          "ping_interval": 20, "ping_timeout": 20}
        while True:
            proto = "wss" if use_tls else "ws"
            # Bei TLS + RVS_SNI: URI nutzt den HOSTNAMEN (Host-Header/SNI/Cert),
            # getaddrinfo mappt ihn auf die IP in RVS_HOST. Bei ws:// direkt die IP.
            uri_host = RVS_SNI if (use_tls and RVS_SNI) else RVS_HOST
            url = f"{proto}://{uri_host}:{RVS_PORT}?token={RVS_TOKEN}"
            fallback = False
            try:
                logger.info("Verbinde mit RVS %s://%s:%s%s …", proto, uri_host, RVS_PORT,
                            f" (TCP {RVS_HOST})" if (use_tls and RVS_SNI) else "")
                async with websockets.connect(url, **connect_kwargs) as ws:
                    self.ws = ws
                    backoff = 1
                    tls_fallback_tried = False
                    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)
                if use_tls and RVS_TLS_FALLBACK and not tls_fallback_tried:
                    logger.info("TLS fehlgeschlagen — Fallback auf ws://")
                    use_tls = False
                    tls_fallback_tried = True
                    fallback = True
            finally:
                self.ws = None
            if fallback:
                continue                        # sofort erneut mit ws://
            await asyncio.sleep(backoff)
            backoff = min(backoff * 2, 30)
            use_tls = RVS_TLS                    # kein Sticky-Fallback
            tls_fallback_tried = False


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


if __name__ == "__main__":
    main()