#!/usr/bin/env python3
"""Capability gate for vps-deployer's dedicated privileged SSH key."""
from __future__ import annotations

import argparse
import base64
import json
import os
from pathlib import PurePosixPath
import re
import shutil
import sys
import tempfile
import unicodedata

# Bump when the client/gate wire protocol changes incompatibly.
GATE_PROTOCOL = 1

SAFE_NAME = re.compile(r"^[a-z][a-z0-9-]{0,62}$")
SAFE_USER = re.compile(r"^[a-z_][a-z0-9_-]{0,31}$")
SAFE_RELEASE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
SAFE_UNIT = re.compile(r"^vps-deployer-[a-z][a-z0-9-]{0,62}\.(service|timer)$")
SAFE_PATH = re.compile(r"^/[A-Za-z0-9._+@:/-]+$")


def _under(path: str, root: str) -> bool:
    candidate, parent = PurePosixPath(path), PurePosixPath(root)
    return candidate.is_absolute() and ".." not in candidate.parts and (candidate == parent or parent in candidate.parents)


def _deployment_path(path: str, managed_root: str) -> bool:
    if not _under(path, managed_root):
        return False
    relative = PurePosixPath(path).relative_to(PurePosixPath(managed_root))
    return not relative.parts or (SAFE_NAME.fullmatch(relative.parts[0]) is not None)


def _owned_file(path: str, managed_root: str, storage: set[str]) -> bool:
    return (_deployment_path(path, managed_root) or path in storage or
            re.fullmatch(r"/etc/vps-deployer/[a-z][a-z0-9-]{0,62}\.(env|state)", path) is not None or
            re.fullmatch(r"/etc/systemd/system/vps-deployer-[a-z][a-z0-9-]{0,62}\.(service|timer)", path) is not None or
            re.fullmatch(r"/etc/nginx/sites-(available|enabled)/[a-z][a-z0-9-]{0,62}", path) is not None)


def _upload_path(path: str) -> bool:
    return re.fullmatch(r"/tmp/vps-deployer-[A-Za-z0-9._-]+\.tar\.gz", path) is not None


def _unit_path(path: str) -> bool:
    return re.fullmatch(r"/etc/systemd/system/vps-deployer-[a-z][a-z0-9-]{0,62}\.(service|timer)", path) is not None


def validate_unit(target: str, content: str, managed_root: str, storage: set[str]) -> None:
    match = re.fullmatch(r"/etc/systemd/system/vps-deployer-([a-z][a-z0-9-]{0,62})\.(service|timer)", target)
    if not match:
        raise ValueError("invalid managed unit target")
    name, kind = match.groups()
    if kind == "timer":
        timer = re.fullmatch(
            rf"\[Unit\]\nDescription=vps-deployer timer {re.escape(name)}\n\n"
            rf"\[Timer\]\nOnCalendar=([^\n%]+)\nPersistent=(true|false)\n"
            rf"(?:RandomizedDelaySec=([0-9]+)\n)?Unit=vps-deployer-{re.escape(name)}\.service\n\n"
            rf"\[Install\]\nWantedBy=timers\.target\n", content)
        if not timer or (timer.group(3) is not None and int(timer.group(3)) > 86400):
            raise ValueError("invalid managed timer content")
        return
    service = re.fullmatch(
        rf"\[Unit\]\nDescription=vps-deployer service {re.escape(name)}\n"
        rf"After=network-online\.target\nWants=network-online\.target\n\n"
        rf"\[Service\]\nType=(simple|oneshot)\nUser=([a-z_][a-z0-9_-]{{0,31}})\nGroup=([^\n]+)\n"
        rf"WorkingDirectory=([^\n%]+)\nEnvironmentFile=/etc/vps-deployer/{re.escape(name)}\.env\n"
        rf"ExecStart=([^\n%]+)\n(Restart=(always|on-failure|no)\n)?"
        rf"NoNewPrivileges=true\nPrivateTmp=true\nProtectHome=true\nProtectSystem=strict\n"
        rf"((?:ReadWritePaths=[^\n]+\n)*)\n{{1,3}}"
        rf"((?:\[Install\]\nWantedBy=multi-user\.target\n)?)", content)
    if not service:
        raise ValueError("invalid managed service content")
    service_type, user, group, working, command = service.group(1, 2, 3, 4, 5)
    restart, writable, install = service.group(7), service.group(8), service.group(9)
    if user == "root" or group != user:
        raise ValueError("managed services must use a matching non-root identity")
    base = f"{managed_root}/{name}/current"
    if not _under(working, base) or not SAFE_PATH.fullmatch(working):
        raise ValueError("managed service working directory is outside its release")
    executable = command.split(" ", 1)[0].strip("'")
    if not _under(executable, working) or not SAFE_PATH.fullmatch(executable):
        raise ValueError("managed service command is outside its working directory")
    paths = [line.removeprefix("ReadWritePaths=") for line in writable.splitlines()]
    if any(path not in storage for path in paths):
        raise ValueError("managed service writable path is outside policy")
    if service_type == "simple" and (restart is None or not install):
        raise ValueError("persistent service lifecycle is incomplete")
    if service_type == "oneshot" and (restart is not None or install):
        raise ValueError("timer service lifecycle is invalid")


def build_argv(request: dict[str, object], managed_root: str, storage: set[str]) -> list[str]:
    operation = request.get("operation")
    args = request.get("arguments")
    if not isinstance(operation, str) or not isinstance(args, list) or not all(isinstance(x, str) for x in args):
        raise ValueError("invalid operation request")
    if any(any(unicodedata.category(char)[0] == "C" or unicodedata.category(char) in {"Zl", "Zp"}
               for char in value) for value in args):
        raise ValueError("invalid operation argument")

    if operation == "read-file" and len(args) == 1 and _owned_file(args[0], managed_root, storage):
        return ["cat", args[0]]
    if operation == "read-executable" and len(args) == 1 and re.fullmatch(r"/usr/local/bin/[A-Za-z0-9][A-Za-z0-9._+-]{0,63}", args[0]):
        return ["cat", args[0]]
    if operation == "path-exists" and len(args) == 2 and args[0] == "-e" and _owned_file(args[1], managed_root, storage):
        return ["test", "-e", args[1]]
    if (operation == "expect-path" and len(args) == 2 and args[0] == "-e" and
            args[1] != "/" and SAFE_PATH.fullmatch(args[1]) and ".." not in PurePosixPath(args[1]).parts):
        return ["test", "-e", args[1]]
    if operation == "read-link" and len(args) == 1 and _deployment_path(args[0], managed_root):
        return ["readlink", args[0]]
    if operation == "ensure-directories" and args and args[0] == "-p" and all(
            _deployment_path(path, managed_root) or path == "/etc/vps-deployer" for path in args[1:]):
        return ["mkdir", "-p", *args[1:]]
    if operation == "create-user" and args[:4] == ["--system", "--no-create-home", "--shell", "/usr/sbin/nologin"] and len(args) == 5 and SAFE_USER.fullmatch(args[4]) and args[4] != "root":
        return ["useradd", *args]
    if operation == "install":
        if len(args) >= 8 and args[:7] == ["-d", "-o", "root", "-g", "root", "-m", "0755"]:
            paths = args[7:]
            if paths and all(_deployment_path(path, managed_root) or path == "/etc/vps-deployer" for path in paths):
                return ["install", *args]
        if len(args) == 8 and args[:1] == ["-d"] and args[1:7:2] == ["-o", "-g", "-m"]:
            owner, group, mode, path = args[2], args[4], args[6], args[7]
            if SAFE_USER.fullmatch(owner) and group == owner and mode == "0750" and path in storage:
                return ["install", *args]
        raise ValueError("install operation is outside policy")
    if operation == "write-file" and len(args) == 4:
        target, owner, group, mode = args
        if (not _unit_path(target) and _owned_file(target, managed_root, storage) and SAFE_USER.fullmatch(owner) and
                SAFE_USER.fullmatch(group) and mode in {"0640", "0644"}):
            return ["__write_file__", target, owner, group, mode]
    if operation == "write-unit" and len(args) == 1 and _unit_path(args[0]):
        return ["__write_unit__", args[0]]
    if operation == "write-executable" and len(args) == 1 and re.fullmatch(r"/usr/local/bin/[A-Za-z0-9][A-Za-z0-9._+-]{0,63}", args[0]):
        return ["__write_executable__", args[0]]
    if operation == "remove-executable" and len(args) == 2 and args[0] == "-f" and re.fullmatch(r"/usr/local/bin/[A-Za-z0-9][A-Za-z0-9._+-]{0,63}", args[1]):
        return ["rm", "-f", args[1]]
    if operation == "hold-lock" and len(args) == 1 and SAFE_NAME.fullmatch(args[0]):
        return ["__hold_lock__", args[0]]
    if operation == "service-control":
        if args in (["daemon-reload"], ["reset-failed"], ["reload", "nginx"]):
            return ["systemctl", *args]
        actions = {"enable", "disable", "restart", "start", "stop", "status", "is-active"}
        units = [value for value in args[1:] if not value.startswith("-")]
        if args and args[0] in actions and units and all(SAFE_UNIT.fullmatch(unit) for unit in units):
            return ["systemctl", *args]
        if (len(args) == 6 and args[0] == "show" and SAFE_UNIT.fullmatch(args[1]) and
                args[2:] == ["--no-pager", "--property=Result", "--property=ExecMainStatus",
                             "--property=ExecMainExitTimestamp"]):
            return ["systemctl", *args]
    if operation == "validate-nginx" and args == ["-t"]:
        return ["nginx", "-t"]
    if operation == "set-link" and len(args) == 3 and args[0] in {"-s", "-sfn"}:
        source, target = args[1:]
        if (_owned_file(target, managed_root, storage) and
                (_owned_file(source, managed_root, storage) or re.fullmatch(r"releases/[A-Za-z0-9._-]+", source))):
            return ["ln", *args]
    if operation == "list-releases" and len(args) == 8 and _deployment_path(args[0], managed_root) and args[1:] == ["-mindepth", "1", "-maxdepth", "1", "-type", "d", "-print"]:
        return ["find", *args]
    if operation == "extract-release" and len(args) == 4 and args[0] == "-xzf" and args[2] == "-C" and _upload_path(args[1]) and _deployment_path(args[3], managed_root):
        return ["tar", "--no-same-owner", "--no-same-permissions", "-xzf", args[1], "-C", args[3]]
    if operation == "remove-paths" and args and args[0] in {"-f", "-rf"} and all(_owned_file(path, managed_root, storage) or _upload_path(path) for path in args[1:]):
        if args[0] != "-rf" or all("/releases/" in path and _deployment_path(path, managed_root) for path in args[1:]):
            return ["rm", *args]
    if operation == "set-ownership" and len(args) == 3 and args[0] == "-R" and re.fullmatch(r"root:[a-z_][a-z0-9_-]{0,31}", args[1]) and _deployment_path(args[2], managed_root):
        return ["chown", *args]
    if operation == "set-mode" and len(args) in {2, 3}:
        recursive = args[0] == "-R"
        mode_index = 1 if recursive else 0
        path_index = 2 if recursive else 1
        allowed_mode = args[mode_index] == "u=rwX,g=rX,o=" if recursive else args[mode_index] == "0750"
        if allowed_mode and _deployment_path(args[path_index], managed_root):
            return ["chmod", *args]
    if operation == "service-logs" and len(args) >= 5 and args[:1] == ["-u"] and SAFE_UNIT.fullmatch(args[1]) and args[2:4] == ["--no-pager", "-n"] and args[4].isdigit():
        tail = args[5:]
        if tail[:1] == ["--since"] and len(tail) >= 2 and tail[1] and not tail[1].startswith("-"):
            tail = tail[2:]
        if tail in ([], ["--follow"]):
            return ["journalctl", *args]
    raise ValueError("operation is not allowed")


def atomic_write(target: str, owner: str, group: str, mode: str, content) -> None:
    import grp
    import pwd
    parent = str(PurePosixPath(target).parent)
    descriptor, temporary = tempfile.mkstemp(prefix=".vps-deployer-", dir=parent)
    try:
        with os.fdopen(descriptor, "wb", closefd=True) as stream:
            shutil.copyfileobj(content, stream)
            stream.flush()
            os.fsync(stream.fileno())
            os.fchown(stream.fileno(), pwd.getpwnam(owner).pw_uid, grp.getgrnam(group).gr_gid)
            os.fchmod(stream.fileno(), int(mode, 8))
        os.replace(temporary, target)
    finally:
        try:
            os.unlink(temporary)
        except FileNotFoundError:
            pass


def write_file(argv: list[str]) -> None:
    _, target, owner, group, mode = argv
    atomic_write(target, owner, group, mode, sys.stdin.buffer)


def write_unit(argv: list[str], managed_root: str, storage: set[str]) -> None:
    import io
    _, target = argv
    raw = sys.stdin.buffer.read(64 * 1024 + 1)
    if len(raw) > 64 * 1024:
        raise ValueError("managed unit is too large")
    try:
        content = raw.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise ValueError("managed unit is not UTF-8") from exc
    validate_unit(target, content, managed_root, storage)
    atomic_write(target, "root", "root", "0644", io.BytesIO(raw))


def write_executable(argv: list[str], managed_root: str) -> None:
    import io
    _, target = argv
    raw = sys.stdin.buffer.read(16 * 1024 + 1)
    if len(raw) > 16 * 1024:
        raise ValueError("managed executable is too large")
    try:
        content = raw.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise ValueError("managed executable is not UTF-8") from exc
    root = re.escape(managed_root)
    pattern = (r"#!/bin/sh\nset -eu\nset -a\n\. /etc/vps-deployer/([a-z][a-z0-9-]{0,62})\.env\nset \+a\n"
               + r'exec "' + root + r'/\1/current/[A-Za-z0-9._+/-]+" "\$@"\n')
    if not re.fullmatch(pattern, content) or ".." in content:
        raise ValueError("managed executable content is invalid")
    atomic_write(target, "root", "root", "0755", io.BytesIO(raw))


def hold_lock(deployment: str) -> None:
    import fcntl
    path = f"/run/lock/vps-deployer-{deployment}.lock"
    descriptor = os.open(path, os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600)
    try:
        try:
            fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            raise ValueError("deployment is already locked") from exc
        sys.stdout.write("LOCKED\n")
        sys.stdout.flush()
        while sys.stdin.buffer.read(64 * 1024):
            pass
    finally:
        os.close(descriptor)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--managed-root", default="/srv/vps-deployer")
    parser.add_argument("--storage", action="append", default=[])
    policy = parser.parse_args()
    if not PurePosixPath(policy.managed_root).is_absolute() or ".." in PurePosixPath(policy.managed_root).parts:
        raise ValueError("invalid managed root policy")
    storage = set(policy.storage)
    if any(not PurePosixPath(path).is_absolute() or ".." in PurePosixPath(path).parts or path == "/" for path in storage):
        raise ValueError("invalid storage policy")
    original = os.environ.get("SSH_ORIGINAL_COMMAND", "")
    match = re.fullmatch(r"vps-deployer-op ([A-Za-z0-9_=-]+)", original)
    if not match:
        raise ValueError("interactive access is disabled")
    request = json.loads(base64.urlsafe_b64decode(match.group(1)))
    if request.get("operation") == "gate-version":
        sys.stdout.write(f"vps-deployer-gate {GATE_PROTOCOL}\n")
        return
    argv = build_argv(request, policy.managed_root, storage)
    if argv[0] == "__write_file__":
        write_file(argv)
        return
    if argv[0] == "__write_unit__":
        write_unit(argv, policy.managed_root, storage)
        return
    if argv[0] == "__write_executable__":
        write_executable(argv, policy.managed_root)
        return
    if argv[0] == "__hold_lock__":
        hold_lock(argv[1])
        return
    executable = shutil.which(argv[0])
    if not executable:
        raise ValueError("executable unavailable")
    os.execv(executable, argv)


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(f"vps-deployer gate: {exc}", file=sys.stderr)
        raise SystemExit(126)
