#!/usr/bin/env python3
"""Guardian agent: keeps this computer's browser policy in line with Guardian.

Runs as a root systemd service (``guardian-agent``) so it can write the
browsers' managed-policy files. It only ever does one thing: write the Chrome
policy Guardian sends to the fixed ``zz-guardian.json`` paths below (or remove
them when the policy is empty). It accepts no commands and no other paths.

It keeps itself up to date: when Guardian reports a newer agent version, it
fetches the code over its signed channel, checks it compiles, replaces this
file and restarts into it.

    guardian_agent.py enroll --server https://guardian.example --token <one-time token>
    guardian_agent.py run
    guardian_agent.py check      # DNS and internet connectivity, right now

Standard library only, so it runs on any stock Linux desktop with python3.
"""

import argparse
import hashlib
import hmac
import json
import logging
import os
import platform
import re
import shutil
import socket
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request

VERSION = "4"
CONFIG_PATH = "/etc/guardian-agent/config.json"
STATE_PATH = "/var/lib/guardian-agent/state.json"
POLL_WAIT = 25
AGENT_PATH = os.path.abspath(__file__)
UPDATE_RETRY = 3600         # after a failed self-update, wait before trying again

# Chrome, Chromium and Brave read every file in these dirs; Guardian owns
# zz-guardian.json and leaves any other policy file alone. When two files set
# the same policy the one whose name sorts last wins, hence the "zz-".
POLICY_DIRS = [
    "/etc/opt/chrome/policies/managed",
    "/etc/chromium/policies/managed",
    "/etc/brave/policies/managed",
]
POLICY_PATHS = [d + "/zz-guardian.json" for d in POLICY_DIRS]
# Agent v1 wrote guardian.json, which lost conflicts to e.g. policy.json.
LEGACY_PATHS = [d + "/guardian.json" for d in POLICY_DIRS]
_POLICY_NAME = re.compile(r"^[A-Za-z][A-Za-z0-9]*$")
BROWSERS = ["google-chrome", "chromium", "chromium-browser", "brave-browser"]

# Network health, checked every NET_CHECK_EVERY seconds and sent with polls.
NET_CHECK_EVERY = 300
NET_TIMEOUT = 5
DNS_PROBES = ["google.com", "cloudflare.com"]
# The internet counts as reachable if a TCP connect works to an address the DNS
# probes returned or, so a DNS outage isn't reported as offline, to one of these
# raw IPs (often blocked on filtered networks, hence only a fallback).
INTERNET_PROBES = [("1.1.1.1", 443), ("8.8.8.8", 443), ("9.9.9.9", 443)]

log = logging.getLogger("guardian-agent")


def sign(secret, *parts):
    return hmac.new(secret.encode(), "\n".join(parts).encode(), hashlib.sha256).hexdigest()


def sha256(data):
    return hashlib.sha256(data).hexdigest()


def write_json_atomic(path, obj, mode):
    os.makedirs(os.path.dirname(path), mode=0o755, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=os.path.dirname(path), prefix=".guardian-")
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(obj, f, indent=2, sort_keys=True)
            f.write("\n")
        os.chmod(tmp, mode)
        os.replace(tmp, path)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def load_json(path, default):
    try:
        with open(path) as f:
            return json.load(f)
    except FileNotFoundError:
        return default


def system_info():
    info = {"hostname": socket.gethostname(), "kernel": platform.release(), "agent": VERSION}
    try:
        with open("/etc/os-release") as f:
            for line in f:
                if line.startswith("PRETTY_NAME="):
                    info["os"] = line.split("=", 1)[1].strip().strip('"')
    except OSError:
        pass
    info["browsers"] = [b for b in BROWSERS if shutil.which(b)]
    info["network"] = network_status()
    return info


# ---- network health --------------------------------------------------------
def _timed(fn, timeout):
    """Run fn in a thread (getaddrinfo has no timeout); returns (result, error, ms)."""
    out = {}

    def target():
        try:
            out["result"] = fn()
        except Exception as e:
            out["error"] = str(e) or type(e).__name__
    started = time.monotonic()
    t = threading.Thread(target=target, daemon=True)
    t.start()
    t.join(timeout)
    ms = int((time.monotonic() - started) * 1000)
    if t.is_alive():
        return None, "timed out after %ss" % timeout, ms
    return out.get("result"), out.get("error"), ms


def resolvers():
    try:
        with open("/etc/resolv.conf") as f:
            return [line.split()[1] for line in f
                    if line.startswith("nameserver") and len(line.split()) > 1]
    except OSError:
        return []


def check_dns(host, timeout=NET_TIMEOUT):
    addrs, err, ms = _timed(lambda: sorted({a[4][0] for a in socket.getaddrinfo(
        host, 443, proto=socket.IPPROTO_TCP)}), timeout)
    return {"host": host, "ok": not err, "ms": ms, "addrs": (addrs or [])[:4],
            **({"error": err} if err else {})}


def check_tcp(ip, port, timeout=NET_TIMEOUT):
    _, err, ms = _timed(lambda: socket.create_connection((ip, port), timeout).close(),
                        timeout + 1)
    return {"target": "%s:%d" % (ip, port), "ok": not err, "ms": ms,
            **({"error": err} if err else {})}


def check_network():
    """Can this computer resolve names and reach the internet? Never raises."""
    dns = [check_dns(h) for h in DNS_PROBES]
    internet = []
    for target in [(d["addrs"][0], 443) for d in dns if d["addrs"]] + INTERNET_PROBES:
        internet.append(check_tcp(*target))
        if internet[-1]["ok"]:
            break
    dns_ok, internet_ok = any(d["ok"] for d in dns), any(i["ok"] for i in internet)
    return {"status": "ok" if dns_ok and internet_ok
            else "dns_failing" if internet_ok else "offline",
            "dns_ok": dns_ok, "internet_ok": internet_ok, "resolvers": resolvers(),
            "dns": dns, "internet": internet, "checked_at": int(time.time())}


_net = {"status": None, "at": 0.0}


def network_status():
    """The last network check, re-run when older than NET_CHECK_EVERY."""
    if time.monotonic() - _net["at"] >= NET_CHECK_EVERY or not _net["status"]:
        status = check_network()
        if _net["status"] and _net["status"]["status"] != status["status"]:
            log.warning("network: %s -> %s", _net["status"]["status"], status["status"])
        _net.update(status=status, at=time.monotonic())
    return _net["status"]


def post(server, path, body, agent_id=None, secret=None, timeout=30):
    """POST JSON; with a secret the request is signed and the reply verified."""
    data = json.dumps(body).encode()
    headers = {"Content-Type": "application/json", "User-Agent": "guardian-agent/" + VERSION}
    sig = None
    if secret:
        ts = str(int(time.time()))
        sig = sign(secret, "POST", path, ts, sha256(data))
        headers.update({"X-Agent-Id": agent_id, "X-Agent-Ts": ts, "X-Agent-Sig": sig})
    req = urllib.request.Request(server.rstrip("/") + path, data=data, headers=headers,
                                 method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as r:
        raw = r.read()
        if secret and not hmac.compare_digest(r.headers.get("X-Guardian-Sig", ""),
                                              sign(secret, "resp", sig, sha256(raw))):
            raise RuntimeError("reply signature does not match; ignoring it")
    return json.loads(raw)


# ---- the one job ---------------------------------------------------------
def apply_chrome_policy(policy):
    if not isinstance(policy, dict) or not all(_POLICY_NAME.match(k) for k in policy):
        raise ValueError("not a Chrome policy object")
    written, removed = [], []
    for path in POLICY_PATHS:
        if policy:
            write_json_atomic(path, policy, 0o644)
            written.append(path)
        elif os.path.exists(path):
            os.unlink(path)
            removed.append(path)
    return {"written": written, "removed": removed}


def migrate_legacy_files():
    """Rename v1's guardian.json to zz-guardian.json so it wins conflicts."""
    for old, new in zip(LEGACY_PATHS, POLICY_PATHS):
        if os.path.exists(old):
            if os.path.exists(new):
                os.unlink(old)
            else:
                os.replace(old, new)
            log.info("moved %s to %s", old, new)


def run_job(job):
    result = {"job_id": job.get("id")}
    try:
        if job.get("kind") != "chrome_policy":
            raise ValueError("unsupported job kind %r" % job.get("kind"))
        result.update(apply_chrome_policy((job.get("payload") or {}).get("policy")))
        result["ok"] = True
    except Exception as e:  # report every failure back to Guardian
        result.update(ok=False, error="%s: %s" % (type(e).__name__, e))
    return result


# ---- self-update -----------------------------------------------------------
def self_update(cfg, want):
    """Replace this file with Guardian's agent ``want``; returns True when done."""
    reply = post(cfg["server"], "/api/agent/code", {}, cfg["agent_id"], cfg["secret"])
    code = reply.get("code") or ""
    m = re.search(r'^VERSION = "([^"]+)"', code, re.M)
    if reply.get("version") != want or not m or m.group(1) != want:
        raise RuntimeError("server sent agent %r, expected %r" % (m and m.group(1), want))
    compile(code, AGENT_PATH, "exec")  # never swap in code that cannot run
    fd, tmp = tempfile.mkstemp(dir=os.path.dirname(AGENT_PATH), prefix=".guardian-")
    try:
        with os.fdopen(fd, "w") as f:
            f.write(code)
        os.chmod(tmp, 0o755)
        os.replace(tmp, AGENT_PATH)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise
    log.info("updated agent v%s -> v%s", VERSION, want)
    return True


def maybe_update(cfg, reply, state):
    """Self-update when Guardian ships another version; True means restart."""
    want = reply.get("agent_version")
    if not want or want == VERSION:
        return False
    failed = state.get("update_failed") or {}
    if failed.get("version") == want and time.time() - failed.get("at", 0) < UPDATE_RETRY:
        return False
    try:
        return self_update(cfg, want)
    except Exception as e:
        log.warning("self-update to v%s failed: %s", want, e)
        state["update_failed"] = {"version": want, "at": time.time()}
        write_json_atomic(STATE_PATH, state, 0o600)
        return False


def restart():
    os.execv(sys.executable, [sys.executable, AGENT_PATH] + sys.argv[1:])


# ---- commands --------------------------------------------------------------
def cmd_enroll(args):
    reply = post(args.server, "/api/agent/enroll", {"token": args.token, "info": system_info()})
    os.makedirs(os.path.dirname(CONFIG_PATH), mode=0o700, exist_ok=True)
    write_json_atomic(CONFIG_PATH, {"server": args.server.rstrip("/"),
                                    "agent_id": reply["agent_id"],
                                    "secret": reply["secret"]}, 0o600)
    print("Enrolled as '%s'." % reply["name"])


def poll_once(cfg, state, wait=POLL_WAIT):
    """One long-poll: report finished jobs, apply new ones. Mutates and saves state."""
    reply = post(cfg["server"], "/api/agent/poll",
                 {"info": system_info(), "results": state["outbox"], "wait": wait},
                 cfg["agent_id"], cfg["secret"], timeout=wait + 20)
    acked = set(reply.get("acked") or [])
    state["outbox"] = [r for r in state["outbox"] if r["job_id"] not in acked]
    for job in reply.get("jobs") or []:
        if job.get("id") in state["done"]:
            continue  # applied already; its result is in the outbox
        result = run_job(job)
        log.info("job %s: %s", job.get("id"), "ok" if result["ok"] else result["error"])
        state["done"] = (state["done"] + [job["id"]])[-200:]
        state["outbox"].append(result)
    write_json_atomic(STATE_PATH, state, 0o600)
    return reply


def cmd_check(_args):
    status = check_network()
    print(json.dumps(status, indent=2))
    sys.exit(0 if status["status"] == "ok" else 1)


def cmd_run(_args):
    cfg = load_json(CONFIG_PATH, None)
    if not cfg:
        sys.exit("not enrolled: run '%s enroll' first" % sys.argv[0])
    state = load_json(STATE_PATH, {"done": [], "outbox": []})
    migrate_legacy_files()
    backoff = 1
    log.info("agent v%s polling %s as %s", VERSION, cfg["server"], cfg["agent_id"])
    while True:
        try:
            reply = poll_once(cfg, state)
            backoff = 1
            if maybe_update(cfg, reply, state):
                restart()
        except urllib.error.HTTPError as e:
            log.warning("server said %s %s", e.code, e.read()[:200].decode(errors="replace"))
            time.sleep(min(backoff, 300))
            backoff *= 2
        except Exception as e:
            log.warning("poll failed: %s", e)
            time.sleep(min(backoff, 60))
            backoff *= 2


def main():
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = p.add_subparsers(dest="cmd", required=True)
    e = sub.add_parser("enroll", help="register this computer with Guardian")
    e.add_argument("--server", required=True)
    e.add_argument("--token", required=True)
    sub.add_parser("run", help="poll Guardian and apply policy (the service)")
    sub.add_parser("check", help="check DNS and internet connectivity and print it")
    args = p.parse_args()
    {"enroll": cmd_enroll, "run": cmd_run, "check": cmd_check}[args.cmd](args)


if __name__ == "__main__":
    main()