#!/usr/bin/env python3
"""
PoC: GlobalProtect SAML-POST (Duo) auf macOS rootlos.

Holt den GP-Prelogin, fuehrt das Duo-SAML in einem eingebetteten WKWebView durch,
greift den `prelogin-cookie`/`saml-username` aus den HTTP-Antwort-Headern ab
(was QtWebEngine nicht kann, das native WKWebView ueber den Navigation-Response-
Delegate aber schon) und startet damit openconnect rootlos ueber ocproxy.

Start:  /Library/Developer/CommandLineTools/usr/bin/python3 ~/apps/heron_gp_saml.py
"""
from __future__ import annotations

import base64
import json
import os
import re
import shutil
import subprocess
import sys
import urllib.request

import objc
from Cocoa import (
    NSApplication, NSObject, NSWindow, NSURL, NSMakeRect,
    NSBackingStoreBuffered,
)
from Foundation import NSHTTPURLResponse, NSHTTPCookie
from WebKit import WKWebView, WKWebViewConfiguration

SERVER = "ssl.heron.at"
PORT = 11090
GP_UA = "PAN GlobalProtect"   # GP liefert den prelogin-cookie nur dem Client-UA
NSWindowStyleMaskTitled = 1
NSWindowStyleMaskClosable = 2
NSWindowStyleMaskResizable = 8
WKNavigationResponsePolicyAllow = 1

captured: dict[str, str] = {}


def _tool(n: str) -> str:
    return shutil.which(n) or (
        f"/opt/homebrew/bin/{n}" if os.path.exists(f"/opt/homebrew/bin/{n}") else n
    )


def _hipreport() -> str | None:
    import glob
    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 prelogin() -> tuple[str, str]:
    """Liefert (saml_html, sessid)."""
    # Gateway-Prelogin direkt (nicht das Portal) -> liefert den Gateway-Cookie.
    url = (f"https://{SERVER}/ssl-vpn/prelogin.esp"
           "?tmp=tmp&clientVer=4100&clientos=Mac")
    req = urllib.request.Request(url, method="POST", data=b"",
                                 headers={"User-Agent": GP_UA})
    with urllib.request.urlopen(req, timeout=20) as r:
        xml = r.read().decode("utf-8", "replace")
        sessid = ""
        for c in r.headers.get_all("Set-Cookie") or []:
            m = re.search(r"SESSID=([^;]+)", c)
            if m:
                sessid = m.group(1)
    m = re.search(r"<saml-request>(.*?)</saml-request>", xml, re.S)
    if not m:
        raise SystemExit("Kein <saml-request> im Prelogin.")
    return base64.b64decode(m.group(1)).decode("utf-8", "replace"), sessid


class Delegate(NSObject):
    def webView_decidePolicyForNavigationAction_decisionHandler_(
            self, webView, navAction, handler):
        u = str(navAction.request().URL().absoluteString())
        print("NAV ->", u[:120], file=sys.stderr)
        if u.startswith("globalprotectcallback:"):
            payload = u[len("globalprotectcallback:"):]
            try:
                dec = base64.b64decode(payload).decode("utf-8", "replace")
            except Exception:
                dec = payload
            print("CALLBACK-Payload:", dec[:500], file=sys.stderr)
            for key in ("prelogin-cookie", "saml-username",
                        "portal-userauthcookie"):
                m = re.search(rf"{key}[=:>]+\s*([^<&\"']+)", dec)
                if m:
                    captured[key] = m.group(1).strip()
            handler(0)  # Cancel
            if "prelogin-cookie" in captured:
                NSApplication.sharedApplication().stop_(None)
            return
        handler(1)  # Allow

    def webView_decidePolicyForNavigationResponse_decisionHandler_(
            self, webView, navResponse, handler):
        resp = navResponse.response()
        if resp.isKindOfClass_(NSHTTPURLResponse):
            for k, v in resp.allHeaderFields().items():
                lk = str(k).lower()
                if lk in ("prelogin-cookie", "saml-username",
                          "portal-userauthcookie", "saml-auth-status"):
                    captured[lk] = str(v)
        handler(WKNavigationResponsePolicyAllow)
        if "prelogin-cookie" in captured and "saml-username" in captured:
            NSApplication.sharedApplication().stop_(None)

    def webView_didFinishNavigation_(self, webView, nav):
        # Nur die SAML-Weiterleitungsform (id=myform) anstossen, sonst wuerde
        # auch das Duo-Login-Formular dauernd neu abgesendet (Flackern).
        webView.evaluateJavaScript_completionHandler_(
            "var f=document.getElementById('myform'); if(f){f.submit();}", None)


def run_saml(html: str, sessid: str):
    app = NSApplication.sharedApplication()
    app.setActivationPolicy_(0)   # Regular -> Fenster bekommt sauber den Fokus
    rect = NSMakeRect(0, 0, 540, 720)
    style = (NSWindowStyleMaskTitled | NSWindowStyleMaskClosable
             | NSWindowStyleMaskResizable)
    win = NSWindow.alloc().initWithContentRect_styleMask_backing_defer_(
        rect, style, NSBackingStoreBuffered, False)
    win.setTitle_("Heron SSO-Login")

    cfg = WKWebViewConfiguration.alloc().init()
    wv = WKWebView.alloc().initWithFrame_configuration_(rect, cfg)
    wv.setCustomUserAgent_(GP_UA)
    delegate = Delegate.alloc().init()
    wv.setNavigationDelegate_(delegate)
    win.setContentView_(wv)
    win.makeKeyAndOrderFront_(None)

    base = NSURL.URLWithString_(f"https://{SERVER}/")

    def load(_=None):
        wv.loadHTMLString_baseURL_(html, base)

    if sessid:
        store = cfg.websiteDataStore().httpCookieStore()
        cookie = NSHTTPCookie.cookieWithProperties_({
            "Name": "SESSID", "Value": sessid, "Domain": SERVER,
            "Path": "/", "Secure": "TRUE",
        })
        store.setCookie_completionHandler_(cookie, load)
    else:
        load()

    app.activateIgnoringOtherApps_(True)
    app.run()


def main():
    auth_only = "--auth-only" in sys.argv
    html, sessid = prelogin()
    run_saml(html, sessid)

    if "prelogin-cookie" not in captured:
        print(f"Kein prelogin-cookie erhalten. Erfasst: {list(captured)}",
              file=sys.stderr)
        raise SystemExit(2)

    if auth_only:
        # Nur den Cookie als JSON auf stdout -> die Tray-App liest ihn.
        print(json.dumps({
            "prelogin-cookie": captured["prelogin-cookie"],
            "saml-username": captured.get("saml-username", ""),
        }))
        return

    print("Erfasst:", {k: v[:20] + "…" for k, v in captured.items()},
          file=sys.stderr)
    cmd = [
        _tool("openconnect"), "--protocol=gp",
        "--user", captured.get("saml-username", ""),
        "--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)
    print("Starte:", " ".join(cmd), file=sys.stderr)
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE)
    p.stdin.write((captured["prelogin-cookie"] + "\n").encode())
    p.stdin.flush()
    p.wait()


if __name__ == "__main__":
    main()