#!/usr/bin/python3 -I
"""Control one digest-bound leoclean no-send container by verified full ID."""

from __future__ import annotations

import argparse
import json
import os
import re
import stat
import subprocess
import sys
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any

DOCKER_BINARY = "/usr/bin/docker"
DOCKER_HOST = "unix:///var/run/docker.sock"
DOCKER_CONFIG = "/var/lib/livingip/leoclean-nosend/docker-config"
CONTAINER_NAME = "livingip-leoclean-nosend"
CIDFILE = Path("/run/leoclean-gcp-nosend/container.cid")
IMAGE_REPOSITORY = "europe-west6-docker.pkg.dev/teleo-501523/teleo/leoclean-nosend-staging"
SAFE_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
MAX_DOCKER_OUTPUT_BYTES = 1024 * 1024
HEX_64 = re.compile(r"^[0-9a-f]{64}$")
HEX_PARTIAL = re.compile(r"^[0-9a-f]{0,63}$")
IMAGE_REFERENCE = re.compile(rf"^{re.escape(IMAGE_REPOSITORY)}@sha256:[0-9a-f]{{64}}$")
IMAGE_CONFIG_DIGEST = re.compile(r"^sha256:[0-9a-f]{64}$")

Runner = Callable[[list[str]], subprocess.CompletedProcess[bytes]]


class ControlError(RuntimeError):
    """A bounded service-control failure."""


def _require(condition: bool, message: str) -> None:
    if not condition:
        raise ControlError(message)


def _strict_json(raw: bytes) -> Any:
    _require(len(raw) <= MAX_DOCKER_OUTPUT_BYTES, "docker output exceeded the safety limit")

    def reject_duplicates(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
        value: dict[str, Any] = {}
        for key, item in pairs:
            if key in value:
                raise ValueError("duplicate key")
            value[key] = item
        return value

    try:
        return json.loads(
            raw.decode("utf-8", "strict"),
            object_pairs_hook=reject_duplicates,
            parse_constant=lambda _value: (_ for _ in ()).throw(ValueError("constant")),
        )
    except (UnicodeError, ValueError, json.JSONDecodeError) as exc:
        raise ControlError("docker inspection was invalid") from exc


def _default_runner(argv: list[str]) -> subprocess.CompletedProcess[bytes]:
    environment = {"PATH": SAFE_PATH}
    return subprocess.run(
        argv,
        check=False,
        stdin=subprocess.DEVNULL,
        capture_output=True,
        timeout=45,
        env=environment,
    )


def _docker(arguments: list[str], runner: Runner) -> subprocess.CompletedProcess[bytes]:
    argv = [
        DOCKER_BINARY,
        f"--host={DOCKER_HOST}",
        f"--config={DOCKER_CONFIG}",
        *arguments,
    ]
    try:
        completed = runner(argv)
    except (OSError, subprocess.SubprocessError) as exc:
        raise ControlError("docker command failed") from exc
    _require(
        len(completed.stdout) <= MAX_DOCKER_OUTPUT_BYTES and len(completed.stderr) <= MAX_DOCKER_OUTPUT_BYTES,
        "docker output exceeded the safety limit",
    )
    return completed


def _inspect(reference: str, runner: Runner) -> dict[str, Any] | None:
    completed = _docker(["container", "inspect", reference], runner)
    if completed.returncode != 0:
        try:
            error = completed.stderr.decode("utf-8", "strict").strip()
        except UnicodeError as exc:
            raise ControlError("container absence could not be proven") from exc
        absence = re.compile(
            rf"^(?:Error response from daemon: |Error: )?No such (?:container|object): {re.escape(reference)}$"
        )
        _require(absence.fullmatch(error) is not None, "container absence could not be proven")
        return None
    value = _strict_json(completed.stdout)
    _require(
        isinstance(value, list) and len(value) == 1 and isinstance(value[0], dict),
        "container inspection was not exact",
    )
    return value[0]


def _validate_container(
    value: dict[str, Any],
    *,
    container_id: str,
    release_sha256: str,
    image_reference: str,
    image_config_digest: str,
) -> bool:
    config = value.get("Config")
    labels = config.get("Labels") if isinstance(config, dict) else None
    state = value.get("State")
    _require(
        value.get("Id") == container_id
        and value.get("Name") == f"/{CONTAINER_NAME}"
        and value.get("Image") == image_config_digest
        and isinstance(config, dict)
        and config.get("Image") == image_reference
        and isinstance(labels, dict)
        and labels.get("livingip.release.sha256") == release_sha256
        and labels.get("livingip.container.role") == "leoclean-gcp-nosend"
        and isinstance(state, dict)
        and isinstance(state.get("Running"), bool),
        "container identity did not match the installed release",
    )
    return bool(state["Running"])


def _cidfile_parent_descriptor(cidfile: Path) -> int:
    flags = os.O_RDONLY
    if hasattr(os, "O_DIRECTORY"):
        flags |= os.O_DIRECTORY
    if hasattr(os, "O_NOFOLLOW"):
        flags |= os.O_NOFOLLOW
    try:
        descriptor = os.open(cidfile.parent, flags)
        observed = os.fstat(descriptor)
        named = os.stat(cidfile.parent, follow_symlinks=False)
    except OSError as exc:
        raise ControlError("container ID directory was unsafe") from exc
    if not (
        stat.S_ISDIR(observed.st_mode)
        and (observed.st_dev, observed.st_ino) == (named.st_dev, named.st_ino)
        and observed.st_uid == os.geteuid()
        and observed.st_mode & 0o077 == 0
    ):
        os.close(descriptor)
        raise ControlError("container ID directory was unsafe")
    return descriptor


def _read_cidfile(cidfile: Path) -> tuple[str | None, tuple[int, int]] | None:
    parent_descriptor = _cidfile_parent_descriptor(cidfile)
    descriptor: int | None = None
    try:
        flags = os.O_RDONLY
        if hasattr(os, "O_NOFOLLOW"):
            flags |= os.O_NOFOLLOW
        try:
            descriptor = os.open(cidfile.name, flags, dir_fd=parent_descriptor)
        except FileNotFoundError:
            return None
        observed = os.fstat(descriptor)
        named = os.stat(cidfile.name, dir_fd=parent_descriptor, follow_symlinks=False)
        _require(
            stat.S_ISREG(observed.st_mode)
            and (observed.st_dev, observed.st_ino) == (named.st_dev, named.st_ino)
            and observed.st_uid == os.geteuid()
            and stat.S_IMODE(observed.st_mode) == 0o600
            and observed.st_nlink == 1
            and observed.st_size <= 65,
            "container ID file was unsafe",
        )
        raw = os.read(descriptor, 66)
        try:
            container_id = raw.decode("ascii", "strict").strip()
        except UnicodeError as exc:
            raise ControlError("container ID file was invalid") from exc
        if HEX_64.fullmatch(container_id) is not None:
            return container_id, (observed.st_dev, observed.st_ino)
        _require(HEX_PARTIAL.fullmatch(container_id) is not None, "container ID file was invalid")
        return None, (observed.st_dev, observed.st_ino)
    except OSError as exc:
        raise ControlError("container ID file could not be read") from exc
    finally:
        if descriptor is not None:
            os.close(descriptor)
        os.close(parent_descriptor)


def _unlink_cidfile(cidfile: Path, identity: tuple[int, int]) -> None:
    parent_descriptor = _cidfile_parent_descriptor(cidfile)
    try:
        observed = os.stat(cidfile.name, dir_fd=parent_descriptor, follow_symlinks=False)
        _require(
            stat.S_ISREG(observed.st_mode) and (observed.st_dev, observed.st_ino) == identity,
            "container ID file changed during cleanup",
        )
        os.unlink(cidfile.name, dir_fd=parent_descriptor)
        os.fsync(parent_descriptor)
    except OSError as exc:
        raise ControlError("container ID file cleanup failed") from exc
    finally:
        os.close(parent_descriptor)


def _wait_absent(container_id: str, runner: Runner, sleeper: Callable[[float], None]) -> None:
    for _attempt in range(60):
        if _inspect(container_id, runner) is None:
            return
        sleeper(0.25)
    raise ControlError("container did not stop within the bounded wait")


def _require_name_unbound(runner: Runner) -> None:
    _require(_inspect(CONTAINER_NAME, runner) is None, "container name was rebound after the trusted container exited")


def _stop_exact(container_id: str, runner: Runner, sleeper: Callable[[float], None]) -> None:
    completed = _docker(["container", "stop", "--time=30", container_id], runner)
    _require(completed.returncode == 0, "container stop failed")
    _wait_absent(container_id, runner, sleeper)


def _remove_created_exact(
    container_id: str,
    *,
    release_sha256: str,
    image_reference: str,
    image_config_digest: str,
    runner: Runner,
    sleeper: Callable[[float], None],
) -> None:
    for _attempt in range(4):
        completed = _docker(["container", "rm", container_id], runner)
        if completed.returncode == 0:
            _wait_absent(container_id, runner, sleeper)
            return
        observed = _inspect(container_id, runner)
        if observed is None:
            return
        running = _validate_container(
            observed,
            container_id=container_id,
            release_sha256=release_sha256,
            image_reference=image_reference,
            image_config_digest=image_config_digest,
        )
        if running:
            _stop_exact(container_id, runner, sleeper)
            return
        sleeper(0.25)
    raise ControlError("created container removal failed")


def _complete_cidfile(
    cidfile: Path,
    identity: tuple[int, int],
    *,
    sleeper: Callable[[float], None],
) -> tuple[str, tuple[int, int]]:
    for _attempt in range(40):
        sleeper(0.25)
        record = _read_cidfile(cidfile)
        _require(record is not None and record[1] == identity, "container ID file changed during startup")
        if record[0] is not None:
            return record[0], record[1]
    raise ControlError("container ID file did not become complete within the bounded wait")


def control_container(
    action: str,
    release_sha256: str,
    image_reference: str,
    image_config_digest: str,
    *,
    runner: Runner = _default_runner,
    sleeper: Callable[[float], None] = time.sleep,
    cidfile: Path = CIDFILE,
) -> dict[str, object]:
    _require(action in {"prestart", "stop"}, "service-control action was invalid")
    _require(HEX_64.fullmatch(release_sha256) is not None, "release hash was invalid")
    _require(IMAGE_REFERENCE.fullmatch(image_reference) is not None, "image reference was invalid")
    _require(IMAGE_CONFIG_DIGEST.fullmatch(image_config_digest) is not None, "image config digest was invalid")

    cid = _read_cidfile(cidfile)
    by_name = _inspect(CONTAINER_NAME, runner)
    if cid is None:
        _require(by_name is None, "container name is occupied without a trusted ID file")
        return {"action": action, "status": "pass", "container": "absent"}

    container_id, cid_identity = cid
    if container_id is None:
        container_id, cid_identity = _complete_cidfile(
            cidfile,
            cid_identity,
            sleeper=sleeper,
        )
        by_name = _inspect(CONTAINER_NAME, runner)
    by_id = _inspect(container_id, runner)
    if by_id is None:
        _require(by_name is None, "container name was rebound after the trusted container exited")
        _require_name_unbound(runner)
        _unlink_cidfile(cidfile, cid_identity)
        return {"action": action, "status": "pass", "container": "absent"}

    running = _validate_container(
        by_id,
        container_id=container_id,
        release_sha256=release_sha256,
        image_reference=image_reference,
        image_config_digest=image_config_digest,
    )
    _require(
        by_name is not None and by_name.get("Id") == container_id,
        "container name and trusted ID did not agree",
    )
    if running:
        _stop_exact(container_id, runner, sleeper)
    else:
        _require(
            isinstance(by_id.get("State"), dict) and by_id["State"].get("Status") == "created",
            "non-running container was not safely removable",
        )
        _remove_created_exact(
            container_id,
            release_sha256=release_sha256,
            image_reference=image_reference,
            image_config_digest=image_config_digest,
            runner=runner,
            sleeper=sleeper,
        )
    _require_name_unbound(runner)
    _unlink_cidfile(cidfile, cid_identity)
    return {"action": action, "status": "pass", "container": "stopped"}


class _Parser(argparse.ArgumentParser):
    def error(self, _message: str) -> None:
        raise ControlError("invalid arguments")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = _Parser(description=__doc__, add_help=False)
    parser.add_argument("action", choices=("prestart", "stop"))
    parser.add_argument("release_sha256")
    parser.add_argument("image_reference")
    parser.add_argument("image_config_digest")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    try:
        args = parse_args(argv)
        result = control_container(
            args.action,
            args.release_sha256,
            args.image_reference,
            args.image_config_digest,
        )
        sys.stdout.write(json.dumps(result, separators=(",", ":"), sort_keys=True) + "\n")
        return 0
    except ControlError as exc:
        sys.stderr.write(
            json.dumps({"error": str(exc), "status": "fail"}, separators=(",", ":"), sort_keys=True) + "\n"
        )
        return 65
    except Exception:
        sys.stderr.write('{"error":"service control failed","status":"fail"}\n')
        return 65


if __name__ == "__main__":
    raise SystemExit(main())
