#!/usr/bin/env python3
"""Test every pushed commit in a fresh exe.dev sandbox; post the result back.

From https://exe.dev/deploy/ci. Python 3 standard library only.

Works with push webhooks from GitHub, Forgejo, Gitea, and Codeberg, which
sign them the same way and share a commit status API. For each push to a
repository in REPOS:

  1. post a "pending" commit status that links to the job's log;
  2. download the commit's source through the forge integration;
  3. copy BASE_VM into a new sandbox (cp --sandbox) for this job alone;
  4. hand the sandbox the source through a one-time URL and run CI_SCRIPT
     (ci.sh) from the commit, streaming its output into the log;
  5. post "success" or "failure" and delete the sandbox.

The dispatcher holds no keys: the exe.dev API and the forge's API are
integrations (EXE_API, FORGE_API). The sandbox gets the source and nothing
else. Jobs not yet finished are kept in ~/ci/jobs and run again after a
restart.
"""
import base64, hashlib, hmac, html, json, os, queue, secrets, shlex, socket, threading, time
import urllib.error, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

ENV = os.environ
EXE_API = ENV["EXE_API"].rstrip("/")      # e.g. https://ci-exe.int.exe.xyz
FORGE_API = ENV["FORGE_API"].rstrip("/")  # e.g. https://ci-forge.int.exe.xyz (+/api/v1 on Forgejo, Gitea)
FORGE = ENV.get("FORGE", "github")         # github, or forgejo/gitea
REPOS = set(ENV["REPOS"].split())         # owner/repo [owner/repo ...]
BASE_VM = ENV["BASE_VM"]
CP_FLAGS = ENV.get("CP_FLAGS", "--cpu=2 --memory=4GB")
CI_SCRIPT = ENV.get("CI_SCRIPT", "ci.sh")
MAX_JOBS = int(ENV.get("MAX_JOBS", "4"))
MAX_MINUTES = int(ENV.get("MAX_MINUTES", "60"))
CONTEXT = ENV.get("CONTEXT", "exe.dev")
SECRET = ENV["WEBHOOK_SECRET"].encode()
HOST = socket.gethostname()
JOB_PREFIX = ENV.get("JOB_PREFIX", HOST + "-job-")
PUBLIC_URL = f"https://{HOST}.exe.xyz"                       # port 8000, public
LOG_PORT = int(ENV.get("LOG_PORT", "8080"))
LOG_URL = f"https://{HOST}.exe.xyz:{LOG_PORT}"                # private
DIR = os.path.expanduser("~/ci")
JOBS, LOGS, SRC = (os.path.join(DIR, d) for d in ("jobs", "logs", "src"))

jobs = queue.Queue()
sources = {}  # one-time token -> tarball path


def log(*a):
    print(time.strftime("%H:%M:%S"), *a, flush=True)


def call(method, url, body=None, accept="application/json", timeout=60):
    data = body if isinstance(body, (bytes, type(None))) else json.dumps(body).encode()
    req = urllib.request.Request(url, data=data, method=method, headers={
        "Accept": accept, "Content-Type": "application/json",
        "User-Agent": "exe-ci-dispatcher"})
    return urllib.request.urlopen(req, timeout=timeout)


def exe(cmd):
    """Run an exe.dev command over the HTTPS API; returns its JSON."""
    with call("POST", EXE_API + "/exec", cmd.encode()) as r:
        return json.loads(r.read() or b"null")


def status(job, state, description):
    try:
        call("POST", f"{FORGE_API}/repos/{job['repo']}/statuses/{job['sha']}", {
            "state": state, "context": CONTEXT, "description": description[:140],
            "target_url": f"{LOG_URL}/{job['id']}.log"}).close()
    except Exception as e:
        log(f"{job['id']}: status {state}: {e}")


def fetch_source(job, path):
    repo, sha = job["repo"], job["sha"]
    url = (f"{FORGE_API}/repos/{repo}/tarball/{sha}" if FORGE == "github"
           else f"{FORGE_API}/repos/{repo}/archive/{sha}.tar.gz")
    with call("GET", url, accept="*/*", timeout=300) as r, open(path, "wb") as f:
        while chunk := r.read(1 << 20):
            f.write(chunk)


def run_in(vm, script, out):
    """Run a bash script on vm, appending its output to out; returns the exit code."""
    b64 = base64.b64encode(script.encode()).decode()
    cmd = f"ssh {vm} 'echo {b64} | base64 -d > /tmp/ci-job.sh && bash /tmp/ci-job.sh'"
    with call("POST", EXE_API + "/exec", cmd.encode(), accept="application/x-ndjson",
              timeout=MAX_MINUTES * 60 + 120) as r:
        for line in r:
            if not line.strip():
                continue  # keepalive
            msg = json.loads(line)
            for k in ("stdout", "stderr"):
                if k in msg:
                    out.write(base64.b64decode(msg[k]))
                    out.flush()
            if "exit_code" in msg:
                return msg["exit_code"]
            if "error" in msg:
                raise RuntimeError(msg["error"])
    raise RuntimeError("the connection to the VM was cut off")


def run(job):
    name = JOB_PREFIX + job["id"]
    token = secrets.token_urlsafe(24)
    tarball = os.path.join(SRC, token + ".tar.gz")
    t0 = time.time()
    with open(os.path.join(LOGS, job["id"] + ".log"), "ab") as out:
        def note(s):
            out.write(f"==> {s}\n".encode())
            out.flush()
        try:
            note(f"{job['repo']} {job['ref']} {job['sha']}")
            fetch_source(job, tarball)
            sources[token] = tarball
            t1 = time.time()
            exe(f"cp {BASE_VM} {name} --sandbox {CP_FLAGS}")
            log(f"{job['id']}: started {name} for {job['repo']} {job['sha'][:7]}")
            note(f"copied {BASE_VM} to {name} in {time.time() - t1:.1f}s; running {CI_SCRIPT}")
            status(job, "pending", f"running on {name}")
            q = {k: shlex.quote(v) for k, v in job.items()}
            script = f"""set -eo pipefail
mkdir -p ~/src && cd ~/src
curl -fsS --retry 5 --retry-all-errors {PUBLIC_URL}/src/{token} | tar -xz --strip-components=1
export CI=true CI_REPO={q['repo']} CI_REF={q['ref']} CI_SHA={q['sha']}
test -f {CI_SCRIPT} || {{ echo "no {CI_SCRIPT} in this commit"; exit 78; }}
exec timeout {MAX_MINUTES}m bash {CI_SCRIPT}
"""
            for attempt in range(30):  # a new copy takes a moment to answer
                try:
                    code = run_in(name, script, out)
                    break
                except (urllib.error.URLError, RuntimeError) as e:
                    if attempt == 29 or token not in sources:
                        raise  # the job had started: don't run it twice
                    time.sleep(1)
            took = f"{time.time() - t0:.0f}s"
            if code == 0:
                state, desc = "success", f"passed in {took}"
            elif code == 78:
                state, desc = "error", f"no {CI_SCRIPT} in this commit"
            elif code == 124:
                state, desc = "failure", f"timed out after {MAX_MINUTES} minutes"
            else:
                state, desc = "failure", f"failed (exit {code}) in {took}"
        except Exception as e:
            state, desc = "error", f"CI error: {e}"
        note(desc)
        sources.pop(token, None)
        if os.path.exists(tarball):
            os.remove(tarball)
        try:
            exe(f"rm {name}")
        except Exception:
            pass  # never created, or the next start deletes it
    status(job, state, desc)
    os.remove(os.path.join(JOBS, job["id"] + ".json"))
    log(f"{job['id']}: {state}: {desc}; deleted {name}")


def worker():
    while True:
        job = jobs.get()
        try:
            run(job)
        except Exception as e:
            log(f"{job['id']}: {e!r}")


def enqueue(repo, ref, sha):
    job = {"id": time.strftime("%Y%m%d-%H%M%S-") + secrets.token_hex(2),
           "repo": repo, "ref": ref, "sha": sha}
    with open(os.path.join(JOBS, job["id"] + ".json"), "w") as f:
        json.dump(job, f)
    status(job, "pending", "queued")
    jobs.put(job)
    log(f"{job['id']}: queued {repo} {ref} {sha[:7]}")


class Public(BaseHTTPRequestHandler):
    """Port 8000, public: signed webhooks in, one-time source downloads out."""

    def do_POST(self):
        body = self.rfile.read(int(self.headers.get("Content-Length", 0)))
        want = "sha256=" + hmac.new(SECRET, body, hashlib.sha256).hexdigest()
        if not hmac.compare_digest(want, self.headers.get("X-Hub-Signature-256", "")):
            return self.reply(401)
        if self.headers.get("X-GitHub-Event") == "push":
            ev = json.loads(body)
            repo, sha = ev["repository"]["full_name"], ev.get("after", "")
            if repo in REPOS and not ev.get("deleted") and sha.strip("0"):
                enqueue(repo, ev["ref"], sha)
        self.reply(204)

    def do_GET(self):
        path = sources.pop(self.path.removeprefix("/src/"), None)
        if not path:
            return self.reply(404)
        self.send_response(200)
        self.send_header("Content-Type", "application/gzip")
        self.send_header("Content-Length", str(os.path.getsize(path)))
        self.end_headers()
        with open(path, "rb") as f:
            self.wfile.write(f.read())

    def reply(self, code):
        self.send_response(code)
        self.end_headers()

    def log_message(self, *a):
        pass


class Logs(BaseHTTPRequestHandler):
    """Port 8080, private behind exe.dev's proxy: the job logs."""

    def do_GET(self):
        name = self.path.strip("/")
        if not name:
            items = sorted(os.listdir(LOGS), reverse=True)[:200]
            body = "<!doctype html><title>CI logs</title><pre>" + "".join(
                f'<a href="/{html.escape(n)}">{html.escape(n)}</a>\n' for n in items) + "</pre>"
            return self.send(body.encode(), "text/html; charset=utf-8")
        path = os.path.join(LOGS, os.path.basename(name))
        if not os.path.isfile(path):
            self.send_response(404)
            return self.end_headers()
        with open(path, "rb") as f:
            self.send(f.read(), "text/plain; charset=utf-8")

    def send(self, body, ctype):
        self.send_response(200)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *a):
        pass


if __name__ == "__main__":
    for d in (JOBS, LOGS, SRC):
        os.makedirs(d, exist_ok=True)
    # A restart interrupted these: delete their sandboxes, run them again.
    for vm in exe("ls --json")["vms"]:
        if vm["vm_name"].startswith(JOB_PREFIX):
            exe(f"rm {vm['vm_name']}")
            log(f"deleted leftover {vm['vm_name']}")
    for f in sorted(os.listdir(SRC)):
        os.remove(os.path.join(SRC, f))
    for f in sorted(os.listdir(JOBS)):
        with open(os.path.join(JOBS, f)) as fh:
            job = json.load(fh)
        with open(os.path.join(LOGS, job["id"] + ".log"), "a") as out:
            out.write("==> the dispatcher restarted; running the job again\n")
        jobs.put(job)
    for _ in range(MAX_JOBS):
        threading.Thread(target=worker, daemon=True).start()
    threading.Thread(target=ThreadingHTTPServer(("", LOG_PORT), Logs).serve_forever,
                     daemon=True).start()
    log(f"testing pushes to {' '.join(sorted(REPOS))} on copies of {BASE_VM}")
    ThreadingHTTPServer(("", 8000), Public).serve_forever()