#!/usr/bin/env python3
"""
Heron VPN – plattformuebergreifende Tray-App (PyQt6), GlobalProtect.

Baut das VPN rootlos auf (openconnect --protocol=gp + ocproxy als SOCKS5) und
startet einen dedizierten Browser, der per PAC nur die internen Hosts ueber das
VPN schickt und alles andere direkt. Start/Stop ueber das Tray-Icon.

Eigene Ports/Profil/Bundle -> laeuft parallel zur Wuerth-App.

Auth: GlobalProtect-SAML (Duo). Der SAML-Login laeuft als eigener Cocoa/pyobjc-
Prozess (heron_gp_saml.py --auth-only, via CLT-Python) und liefert nur den
prelogin-cookie; diese App startet damit openconnect (+HIP-Report) ueber ocproxy.

Abhaengigkeiten:
  openconnect + ocproxy (macOS: brew install openconnect ocproxy),
  CLT-Python mit pyobjc (pyobjc-framework-WebKit) fuer heron_gp_saml.py.
"""

from __future__ import annotations

import glob as _glob
import json
import os
import platform
import shutil
import signal
import socket
import subprocess
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from PyQt6.QtCore import QProcess, QProcessEnvironment, QTimer, Qt
from PyQt6.QtGui import (
    QAction, QColor, QIcon, QImage, QPainter, QPixmap, qAlpha, qRgba,
)
from PyQt6.QtWidgets import (
    QApplication, QSystemTrayIcon, QMenu, QMessageBox,
)

# ===== CONFIG =================================================================
APP_NAME      = "Heron VPN"
PROTOCOL      = "gp"             # GlobalProtect (Palo Alto)
SERVER        = "ssl.heron.at"
PORT          = 11090            # SOCKS5 (ocproxy) – andere Ports als Wuerth
PACPORT       = 11091            # lokaler PAC-HTTP-Server
# Hosts, die ueber das VPN laufen (Domain-Globs und/oder IP-Globs).
INTERNAL_GLOBS = ["*.heron.at", "172.16.*"]
PROFILE       = os.path.expanduser("~/.heron-edge")
START_URL     = "https://robotunits.com"
LOGFILE       = os.path.expanduser("~/.heron-vpn.log")
PAC_NAME      = "heron.pac"
CLT_PYTHON    = "/Library/Developer/CommandLineTools/usr/bin/python3"  # hat pyobjc

# Lesezeichen, die ins frische Profil gesaet werden (Name, URL).
BOOKMARKS: list[tuple[str, str]] = [
    ("Storybook", "http://172.16.80.30:30060/"),
    ("Heron", "https://heron.at"),
]
# =============================================================================


# ---------- Plattform-Abstraktion --------------------------------------------
def _system() -> str:
    return platform.system()  # 'Darwin' | 'Linux' | 'Windows'


def browser_argv() -> list[str]:
    s = _system()
    if s == "Darwin":
        return ["/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge"]
    if s == "Windows":
        for p in (
            r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe",
            r"C:\Program Files\Microsoft\Edge\Application\msedge.exe",
        ):
            if os.path.exists(p):
                return [p]
        return ["msedge.exe"]
    for name in ("microsoft-edge", "microsoft-edge-stable", "msedge"):
        path = shutil.which(name)
        if path:
            return [path]
    return ["microsoft-edge"]


def seed_profile():
    """Saet beim ersten Start Lesezeichen + sichtbare Leiste ins Profil."""
    default = os.path.join(PROFILE, "Default")
    bm_path = os.path.join(default, "Bookmarks")
    if not os.path.exists(bm_path):
        os.makedirs(default, exist_ok=True)
        children = [
            {"date_added": "0", "id": str(100 + i), "name": name,
             "type": "url", "url": url}
            for i, (name, url) in enumerate(BOOKMARKS)
        ]
        data = {
            "checksum": "",
            "roots": {
                "bookmark_bar": {"children": children, "date_added": "0",
                                 "id": "1", "name": "Lesezeichenleiste",
                                 "type": "folder"},
                "other":  {"children": [], "date_added": "0", "id": "2",
                           "name": "Weitere Lesezeichen", "type": "folder"},
                "synced": {"children": [], "date_added": "0", "id": "3",
                           "name": "Mobile Lesezeichen", "type": "folder"},
            },
            "version": 1,
        }
        with open(bm_path, "w") as f:
            json.dump(data, f, indent=2)

    pref_path = os.path.join(default, "Preferences")
    if not os.path.exists(pref_path):
        os.makedirs(default, exist_ok=True)
        with open(pref_path, "w") as f:
            json.dump({"bookmark_bar": {"show_on_all_tabs": True}}, f)


def _tool(name: str) -> str:
    path = shutil.which(name)
    if path:
        return path
    if _system() == "Darwin":
        cand = f"/opt/homebrew/bin/{name}"
        if os.path.exists(cand):
            return cand
    return name


def _saml_script() -> str:
    """Pfad zu heron_gp_saml.py (im Bundle unter Resources, sonst neben uns)."""
    base = os.environ.get(
        "RESOURCEPATH", os.path.dirname(os.path.abspath(__file__)))
    return os.path.join(base, "heron_gp_saml.py")


def _hipreport() -> str | None:
    for p in ["/opt/homebrew/opt/openconnect/libexec/openconnect/hipreport.sh",
              *_glob.glob("/opt/homebrew/Cellar/openconnect/*/libexec/openconnect/hipreport.sh"),
              "/usr/local/opt/openconnect/libexec/openconnect/hipreport.sh"]:
        if os.path.exists(p):
            return p
    return None


def _resolver_rules() -> str:
    return " , ".join(f"MAP {g} ~NOTFOUND" for g in INTERNAL_GLOBS)


# ---------- PAC-Server (in-process) ------------------------------------------
_PAC_COND = " || ".join(f'shExpMatch(host, "{g}")' for g in INTERNAL_GLOBS)
PAC_BODY = (
    "function FindProxyForURL(url, host) {\n"
    f"  if ({_PAC_COND}) {{\n"
    f'    return "SOCKS5 127.0.0.1:{PORT}";\n'
    "  }\n"
    '  return "DIRECT";\n'
    "}\n"
)


class _PacHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        body = PAC_BODY.encode()
        self.send_response(200)
        self.send_header("Content-Type", "application/x-ns-proxy-autoconfig")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *_):
        pass


def start_pac_server() -> ThreadingHTTPServer:
    srv = ThreadingHTTPServer(("127.0.0.1", PACPORT), _PacHandler)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    return srv


# ---------- Icon (Reiher-Silhouette in Zustandsfarbe) ------------------------
_ICON_CACHE: dict[str, QIcon] = {}
_MASK: QImage | None = None


def _heron_mask() -> QImage:
    global _MASK
    if _MASK is None:
        base = os.environ.get(
            "RESOURCEPATH", os.path.dirname(os.path.abspath(__file__)))
        img = QImage(os.path.join(base, "heron_mask.png"))
        _MASK = (img.scaled(44, 44, Qt.AspectRatioMode.KeepAspectRatio,
                            Qt.TransformationMode.SmoothTransformation)
                 .convertToFormat(QImage.Format.Format_ARGB32)
                 if not img.isNull() else QImage())
    return _MASK


def heron_icon(color: str) -> QIcon:
    if color in _ICON_CACHE:
        return _ICON_CACHE[color]
    m = _heron_mask()
    if m.isNull():                       # Fallback: einfacher Punkt
        pm = QPixmap(22, 22)
        pm.fill(Qt.GlobalColor.transparent)
        p = QPainter(pm)
        p.setRenderHint(QPainter.RenderHint.Antialiasing)
        p.setBrush(QColor(color))
        p.setPen(Qt.PenStyle.NoPen)
        p.drawEllipse(4, 4, 14, 14)
        p.end()
        return QIcon(pm)
    tgt = QColor(color)
    o = QImage(m)
    for y in range(o.height()):
        for x in range(o.width()):
            a = qAlpha(o.pixel(x, y))
            if a:
                o.setPixel(x, y, qRgba(tgt.red(), tgt.green(), tgt.blue(), a))
    icon = QIcon(QPixmap.fromImage(o))
    _ICON_CACHE[color] = icon
    return icon


# ---------- App --------------------------------------------------------------
class VpnTray:
    def __init__(self, app: QApplication):
        self.app = app
        self.oc: subprocess.Popen | None = None
        self.browser: subprocess.Popen | None = None
        self.auth_proc: QProcess | None = None
        self.state = "disconnected"

        self.pac = start_pac_server()

        self.tray = QSystemTrayIcon()
        self.tray.setIcon(heron_icon("#9a9a9a"))
        self.tray.setToolTip(f"{APP_NAME} – getrennt")

        self.menu = QMenu()
        self.act_status = QAction("Getrennt")
        self.act_status.setEnabled(False)
        self.act_connect = QAction("Verbinden")
        self.act_disconnect = QAction("Trennen")
        self.act_open = QAction("Edge oeffnen")
        self.act_quit = QAction("Beenden")

        self.act_connect.triggered.connect(self.connect)
        self.act_disconnect.triggered.connect(self.disconnect)
        self.act_open.triggered.connect(self.launch_browser)
        self.act_quit.triggered.connect(self.quit)

        for a in (self.act_status, None, self.act_connect, self.act_disconnect,
                  self.act_open, None, self.act_quit):
            if a is None:
                self.menu.addSeparator()
            else:
                self.menu.addAction(a)

        self.tray.setContextMenu(self.menu)
        self.tray.show()
        self._set_state("disconnected")

        self.timer = QTimer()
        self.timer.timeout.connect(self._tick)
        self.timer.start(700)

    def _set_state(self, state: str):
        self.state = state
        color, label = {
            "disconnected": ("#9a9a9a", "Getrennt"),
            "connecting":   ("#e6a817", "Verbinde…"),
            "connected":    ("#2e7d32", "Verbunden"),
            "error":        ("#cc3333", "Fehler"),
        }[state]
        self.tray.setIcon(heron_icon(color))
        self.tray.setToolTip(f"{APP_NAME} – {label}")
        self.act_status.setText(label)
        self.act_connect.setEnabled(state in ("disconnected", "error"))
        self.act_disconnect.setEnabled(state in ("connecting", "connected"))
        self.act_open.setEnabled(state == "connected")

    def connect(self):
        if self.auth_proc is not None:
            return
        self._set_state("connecting")
        # SAML-Login als eigener pyobjc/WKWebView-Prozess -> liefert Cookie als JSON.
        self.auth_proc = QProcess()
        # Bereinigte Umgebung: sonst erbt das CLT-Python die PYTHONHOME/PATH des
        # py2app-Bundles und findet sein pyobjc nicht (Subprozess stirbt sofort).
        env = QProcessEnvironment.systemEnvironment()
        for v in ("PYTHONHOME", "PYTHONPATH", "PYTHONEXECUTABLE",
                  "PYTHONNOUSERSITE", "PYTHONDONTWRITEBYTECODE"):
            env.remove(v)
        self.auth_proc.setProcessEnvironment(env)
        self.auth_proc.finished.connect(self._on_auth_done)
        self.auth_proc.start(CLT_PYTHON, [_saml_script(), "--auth-only"])
        if not self.auth_proc.waitForStarted(3000):
            self.auth_proc = None
            self._set_state("error")
            QMessageBox.critical(None, APP_NAME,
                                 "SAML-Login (CLT-Python/pyobjc) nicht startbar.")

    def _on_auth_done(self, code, status):
        out = bytes(self.auth_proc.readAllStandardOutput()).decode("utf-8", "replace")
        self.auth_proc = None
        auth = None
        for line in out.strip().splitlines():
            line = line.strip()
            if line.startswith("{"):
                try:
                    auth = json.loads(line)
                except Exception:
                    pass
        if not auth or not auth.get("prelogin-cookie"):
            self._set_state("error")
            return
        self._start_openconnect(auth.get("saml-username", ""),
                                auth["prelogin-cookie"])

    def _start_openconnect(self, user: str, cookie: str):
        # Reste eines frueheren Laufs wegraeumen, sonst Portkonflikt auf SOCKS.
        if _system() in ("Darwin", "Linux"):
            subprocess.run(["pkill", "-f", f"ocproxy -D {PORT}"],
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        cmd = [
            _tool("openconnect"), f"--protocol={PROTOCOL}",
            "--user", user, "--usergroup=gateway:prelogin-cookie",
            "--passwd-on-stdin",
            "--script", f"{_tool('ocproxy')} -D {PORT}", "--script-tun",
        ]
        hip = _hipreport()
        if hip:
            cmd += ["--csd-wrapper", hip]   # HIP-Report, sonst trennt das Gateway
        cmd.append(SERVER)
        try:
            self._log = open(LOGFILE, "wb")
            self.oc = subprocess.Popen(
                cmd, stdin=subprocess.PIPE,
                stdout=self._log, stderr=subprocess.STDOUT,
            )
            self.oc.stdin.write((cookie + "\n").encode())
            self.oc.stdin.flush()
        except FileNotFoundError:
            self._set_state("error")
            QMessageBox.critical(None, APP_NAME, "openconnect nicht gefunden.")
            return
        # bleibt "connecting"; _tick erkennt den offenen SOCKS-Port

    def disconnect(self):
        if self.auth_proc is not None:
            try:
                self.auth_proc.kill()
            except Exception:
                pass
            self.auth_proc = None
        self._kill_browser()
        if self.oc and self.oc.poll() is None:
            try:
                self.oc.send_signal(signal.SIGINT)
                self.oc.wait(timeout=4)
            except Exception:
                try:
                    self.oc.terminate()
                except Exception:
                    pass
        self.oc = None
        self._set_state("disconnected")

    def _kill_browser(self):
        if _system() in ("Darwin", "Linux"):
            subprocess.run(
                ["pkill", "-f", f"user-data-dir={PROFILE}"],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            )
        elif self.browser and self.browser.poll() is None:
            self.browser.terminate()
        self.browser = None

    def launch_browser(self):
        seed_profile()
        argv = browser_argv() + [
            f"--user-data-dir={PROFILE}",
            f"--proxy-pac-url=http://127.0.0.1:{PACPORT}/{PAC_NAME}",
            f"--host-resolver-rules={_resolver_rules()}",
            "--no-first-run", "--no-default-browser-check",
            "--disable-features=msSmartScreenProtection",
            "--test-type",
            START_URL,
        ]
        self.browser = subprocess.Popen(argv)

    def _tick(self):
        if self.state == "connecting":
            if self.oc and self.oc.poll() is not None:
                self._set_state("error")
                return
            if _port_open(PORT):
                self._set_state("connected")
                self.launch_browser()
        elif self.state == "connected":
            if self.oc and self.oc.poll() is not None:
                self.disconnect()

    def quit(self):
        self.disconnect()
        self.pac.shutdown()
        self.app.quit()


def _port_open(port: int) -> bool:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(0.3)
        return s.connect_ex(("127.0.0.1", port)) == 0


def main():
    app = QApplication([])
    app.setApplicationName("HeronVPN")
    app.setQuitOnLastWindowClosed(False)
    _tray = VpnTray(app)  # noqa: F841
    app.exec()


if __name__ == "__main__":
    main()
