AkurAI-EC2 / libexec/akurai-ec2-dns.py

Latest change d8aa29121486f4b02d0016888d94fa66cac36ffe - Checkpoint WIP + rustfmt + AkurAI Build CI/CD pipeline by Ólafur Búi Ólafsson

#!/usr/bin/env python3
"""Manage AkurAI EC2 DNS zones, with optional legacy 1984 Hosting access."""

from __future__ import annotations

import argparse
import datetime
import html
import http.cookiejar
import json
import os
import re
import socket
import ssl
import subprocess
import sys
import urllib.error
import urllib.parse
import urllib.request
import tomllib
from dataclasses import asdict, dataclass

BASE = "https://1984.hosting"


class DnsError(RuntimeError):
    pass


@dataclass(frozen=True)
class Record:
    id: str
    fqdn: str
    type: str
    ttl: int | None
    value: str


class Client:
    def __init__(self) -> None:
        self.jar = http.cookiejar.CookieJar()
        self.opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(self.jar))
        self.csrf = ""

    def request(self, path: str, data: dict[str, str] | None = None, referer: str | None = None) -> str:
        body = urllib.parse.urlencode(data).encode() if data is not None else None
        headers = {"User-Agent": "akurai-ec2-dns/1"}
        if referer:
            headers["Referer"] = referer
        if self.csrf:
            headers["X-CSRFToken"] = self.csrf
        try:
            with self.opener.open(urllib.request.Request(BASE + path, body, headers), timeout=20) as response:
                return response.read().decode("utf-8", "replace")
        except urllib.error.URLError as exc:
            raise DnsError(f"1984 Hosting request failed: {exc.reason}") from exc

    def login(self) -> None:
        username = os.environ.get("One984HOSTING_Username", "")
        password = os.environ.get("One984HOSTING_Password", "")
        session = os.environ.get("One984HOSTING_SESSIONID_COOKIE", "")
        csrf_cookie = os.environ.get("One984HOSTING_CSRFTOKEN_COOKIE", "")
        if session and csrf_cookie:
            for raw in (session, csrf_cookie):
                name, _, value = raw.partition("=")
                self.jar.set_cookie(http.cookiejar.Cookie(0, name, value, None, False, "1984.hosting", True,
                    False, "/", True, False, None, False, None, None, {}))
            self.csrf = csrf_cookie.partition("=")[2]
            auth = json.loads(self.request("/api/auth/"))
            if auth.get("ok"):
                return
            raise DnsError("cached 1984 Hosting session is invalid")
        if not username or not password:
            raise DnsError("1984 Hosting credentials unavailable; inject PassVault values as "
                           "One984HOSTING_Username and One984HOSTING_Password")
        self.request("/accounts/login/")
        self.csrf = next((cookie.value for cookie in self.jar if cookie.name == "csrftoken"), "")
        otp = ""
        secret = os.environ.get("One984HOSTING_TOTP_Secret", "")
        if secret:
            try:
                otp = subprocess.run(["oathtool", "--base32", "--totp", secret], check=True,
                                     text=True, capture_output=True).stdout.strip()
            except (FileNotFoundError, subprocess.CalledProcessError) as exc:
                raise DnsError("oathtool is required for the configured 1984 Hosting TOTP") from exc
        result = json.loads(self.request("/api/auth/", {
            "username": username,
            "password": password,
            "otpkey": otp,
        }, BASE + "/accounts/login/"))
        if not result.get("loggedin"):
            raise DnsError("1984 Hosting login failed")
        self.csrf = next((cookie.value for cookie in self.jar if cookie.name == "csrftoken"), self.csrf)

    def zone_id(self, zone: str) -> str:
        page = self.request("/domains")
        match = re.search(r'href=["\']/domains/(zone/\d+)["\'][^>]*>[^<]*' + re.escape(zone), page, re.I)
        if not match:
            match = re.search(r'(zone/\d+)[^\n]{0,800}' + re.escape(zone), page, re.I)
        if not match:
            raise DnsError(f"managed DNS zone not found: {zone}")
        return match.group(1)

    def records(self, zone: str) -> list[Record]:
        zid = self.zone_id(zone)
        page = self.request(f"/domains/{zid}")
        rows = re.findall(r'<tr\b[^>]*>.*?</tr>', page, re.I | re.S)
        records: list[Record] = []
        for row in rows:
            entry = re.search(r'entry_(\d+)', row)
            if not entry:
                continue
            text = html.unescape(re.sub(r'<[^>]+>', ' ', row))
            cells = [re.sub(r'\s+', ' ', html.unescape(re.sub(r'<[^>]+>', ' ', cell))).strip()
                     for cell in re.findall(r'<td\b[^>]*>(.*?)</td>', row, re.I | re.S)]
            joined = " | ".join(cells) if cells else re.sub(r'\s+', ' ', text).strip()
            rtype_match = re.search(r'\b(A|AAAA|CAA|CNAME|MX|NS|SRV|TXT)\b', joined, re.I)
            ttl_match = re.search(r'\b(\d{2,6})\b', joined)
            if not rtype_match:
                continue
            rtype = rtype_match.group(1).upper()
            host = next((cell for cell in cells if cell == "@" or cell.endswith(zone) or
                         re.fullmatch(r"[A-Za-z0-9_*.-]+", cell or "") and cell.upper() != rtype), "@")
            fqdn = zone if host == "@" else host if host.endswith(zone) else f"{host}.{zone}"
            value = cells[-1] if cells else joined
            records.append(Record(entry.group(1), fqdn.rstrip('.'), rtype,
                                  int(ttl_match.group(1)) if ttl_match else None, value))
        return records

    def add(self, zone: str, fqdn: str, rtype: str, value: str, ttl: int) -> None:
        self.zone_id(zone)
        host = "@" if fqdn.rstrip(".") == zone else fqdn.removesuffix("." + zone)
        payload = {"entry": "new", "type": rtype, "ttl": str(ttl), "zone": zone,
                   "host": host, "rdata": f'"{value}"' if rtype == "TXT" else value}
        result = json.loads(self.request("/domains/entry/", payload, BASE + "/domains"))
        if result.get("haserrors") or result.get("auth") is False:
            raise DnsError(f"1984 Hosting rejected {rtype} record for {fqdn}")

    def delete(self, zone: str, record_id: str) -> None:
        zid = self.zone_id(zone)
        result = json.loads(self.request("/domains/delentry/", {"entry": record_id},
                                         BASE + f"/domains/{zid}"))
        if not result.get("ok"):
            raise DnsError(f"1984 Hosting failed to delete entry {record_id}")


class Ec2Client:
    def __init__(self) -> None:
        self.cli = os.environ.get("AKURAI_EC2_CLI", "akurai-ec2")

    @staticmethod
    def _path(zone: str) -> str:
        if not re.fullmatch(r"[A-Za-z0-9.-]+", zone):
            raise DnsError(f"invalid zone: {zone}")
        return f"/etc/akurai-dns/zones/{zone}.toml"

    def _run(self, remote: str, content: str | None = None) -> str:
        result = subprocess.run([self.cli, "ssh", remote], input=content, text=True,
                                capture_output=True)
        if result.returncode:
            raise DnsError(result.stderr.strip() or "EC2 DNS operation failed")
        return result.stdout

    def _load(self, zone: str) -> tuple[str, dict]:
        text = self._run(f"sudo cat {self._path(zone)}")
        try:
            return text, tomllib.loads(text)
        except tomllib.TOMLDecodeError as exc:
            raise DnsError(f"invalid EC2 zone file for {zone}: {exc}") from exc

    def records(self, zone: str) -> list[Record]:
        _, data = self._load(zone)
        ttl = int(data.get("zone", {}).get("ttl", 900))
        records = []
        for index, item in enumerate(data.get("record", [])):
            host = str(item.get("name", "@")).rstrip(".")
            fqdn = zone if host == "@" else host if host.endswith(zone) else f"{host}.{zone}"
            records.append(Record(str(index), fqdn, str(item.get("type", "")).upper(),
                                  int(item.get("ttl", ttl)), str(item.get("value", ""))))
        return records

    @staticmethod
    def _record_blocks(text: str) -> list[re.Match[str]]:
        return list(re.finditer(r"(?ms)^\[\[record\]\]\n.*?(?=^\[\[record\]\]|\Z)", text))

    @staticmethod
    def _bump_serial(text: str) -> str:
        match = re.search(r"(?m)^serial\s*=\s*(\d+)\s*$", text)
        if not match:
            raise DnsError("zone SOA serial is missing")
        serial = max(int(match.group(1)) + 1, int(datetime.datetime.now(datetime.UTC).strftime("%Y%m%d01")))
        return text[:match.start(1)] + str(serial) + text[match.end(1):]

    def _save(self, zone: str, text: str) -> None:
        path = self._path(zone)
        try:
            tomllib.loads(text)
        except tomllib.TOMLDecodeError as exc:
            raise DnsError(f"refusing to save invalid EC2 zone file for {zone}: {exc}") from exc
        remote = (
            f"set -euo pipefail; tmp=$(mktemp); backup={path}.bak-$(date -u +%Y%m%d-%H%M%S); "
            "cat >$tmp; "
            f"sudo cp {path} $backup; sudo install -m 0644 $tmp {path}; rm -f $tmp; "
            "if ! sudo systemctl reload akurai-dns.service || "
            "! systemctl is-active --quiet akurai-dns.service; then "
            f"sudo cp $backup {path}; sudo systemctl restart akurai-dns.service; exit 1; fi"
        )
        self._run(remote, text)

    def delete(self, zone: str, record_id: str) -> None:
        text, _ = self._load(zone)
        blocks = self._record_blocks(text)
        index = int(record_id)
        if index < 0 or index >= len(blocks):
            raise DnsError(f"EC2 DNS record index no longer exists: {record_id}")
        block = blocks[index]
        updated = text[:block.start()] + text[block.end():]
        self._save(zone, self._bump_serial(updated).rstrip() + "\n")

    def add(self, zone: str, fqdn: str, rtype: str, value: str, ttl: int) -> None:
        text, data = self._load(zone)
        host = "@" if fqdn.rstrip(".") == zone else fqdn.removesuffix("." + zone)
        zone_ttl = int(data.get("zone", {}).get("ttl", 900))
        lines = ["[[record]]", f'name = {json.dumps(host)}', f'type = {json.dumps(rtype)}']
        if ttl != zone_ttl:
            lines.append(f"ttl = {ttl}")
        lines.append(f'value = {json.dumps(value)}')
        updated = self._bump_serial(text).rstrip() + "\n\n" + "\n".join(lines) + "\n"
        self._save(zone, updated)

def infer_zone(fqdn: str, explicit: str | None) -> str:
    if explicit:
        return explicit.rstrip(".")
    labels = fqdn.rstrip(".").split(".")
    if len(labels) < 2:
        raise DnsError("a fully qualified domain name is required")
    return ".".join(labels[-2:])


def public_values(fqdn: str, rtype: str) -> list[str]:
    url = "https://dns.google/resolve?" + urllib.parse.urlencode({"name": fqdn, "type": rtype})
    try:
        with urllib.request.urlopen(url, timeout=15, context=ssl.create_default_context()) as response:
            payload = json.load(response)
        return [str(answer["data"]).strip('"') for answer in payload.get("Answer", [])
                if int(answer.get("type", 0)) > 0]
    except (urllib.error.URLError, ValueError, KeyError, socket.timeout) as exc:
        raise DnsError(f"public DNS verification failed: {exc}") from exc


def parser() -> argparse.ArgumentParser:
    root = argparse.ArgumentParser(prog="akurai-ec2 dns")
    sub = root.add_subparsers(dest="action", required=True)
    for name in ("list", "get"):
        p = sub.add_parser(name)
        p.add_argument("name")
        p.add_argument("--zone")
        p.add_argument("--type")
        p.add_argument("--provider", choices=("ec2", "1984"), default=os.environ.get("AKURAI_DNS_PROVIDER", "ec2"))
    p = sub.add_parser("upsert")
    p.add_argument("name"); p.add_argument("type"); p.add_argument("value")
    p.add_argument("--zone"); p.add_argument("--ttl", type=int, default=900); p.add_argument("--yes", action="store_true")
    p.add_argument("--provider", choices=("ec2", "1984"), default=os.environ.get("AKURAI_DNS_PROVIDER", "ec2"))
    p = sub.add_parser("delete")
    p.add_argument("name"); p.add_argument("--zone"); p.add_argument("--type", required=True)
    p.add_argument("--value"); p.add_argument("--yes", action="store_true")
    p.add_argument("--provider", choices=("ec2", "1984"), default=os.environ.get("AKURAI_DNS_PROVIDER", "ec2"))
    p = sub.add_parser("verify")
    p.add_argument("name"); p.add_argument("--type", default="A"); p.add_argument("--absent", action="store_true")
    return root


def main() -> int:
    args = parser().parse_args()
    fqdn = args.name.rstrip(".")
    if args.action == "verify":
        values = public_values(fqdn, args.type.upper())
        ok = not values if args.absent else bool(values)
        print(json.dumps({"name": fqdn, "type": args.type.upper(), "values": values,
                          "state": "absent" if not values else "present", "ok": ok}, separators=(",", ":")))
        return 0 if ok else 1

    if args.action in ("delete", "upsert") and not args.yes:
        raise DnsError(f"{args.action} requires --yes")
    zone = infer_zone(fqdn, args.zone)
    client = Ec2Client() if args.provider == "ec2" else Client()
    if isinstance(client, Client):
        client.login()
    records = client.records(zone)
    selected = [record for record in records if (args.action == "list" or record.fqdn == fqdn) and
                (not getattr(args, "type", None) or record.type == args.type.upper())]
    if args.action in ("list", "get"):
        print(json.dumps({"provider": args.provider, "zone": zone,
                          "records": [asdict(record) for record in selected]}, separators=(",", ":")))
        return 0
    # The mutation gate runs before authentication so an accidental command
    # cannot even open a privileged provider session.
    if args.action == "delete":
        if args.value is not None:
            selected = [record for record in selected if args.value in record.value]
        if len(selected) != 1:
            raise DnsError(f"refusing deletion: expected exactly one matching record, found {len(selected)}")
        client.delete(zone, selected[0].id)
        remaining = [record for record in client.records(zone)
                     if record.fqdn == selected[0].fqdn and record.type == selected[0].type]
        if remaining:
            raise DnsError("matching record still present after deletion")
        print(json.dumps({"deleted": asdict(selected[0]), "ok": True}, separators=(",", ":")))
        return 0
    if args.action == "upsert":
        ordered = reversed(selected) if isinstance(client, Ec2Client) else selected
        for record in ordered:
            client.delete(zone, record.id)
        client.add(zone, fqdn, args.type.upper(), args.value, args.ttl)
        print(json.dumps({"name": fqdn, "type": args.type.upper(), "value": args.value,
                          "ttl": args.ttl, "replaced": len(selected), "ok": True}, separators=(",", ":")))
        return 0
    raise DnsError("unsupported DNS action")


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except DnsError as exc:
        print(json.dumps({"ok": False, "error": str(exc)}, separators=(",", ":")), file=sys.stderr)
        raise SystemExit(1)