#!/usr/bin/env python3
"""Reconcile a dedicated unprivileged Hermes SSH account on one Atlas node."""

from __future__ import annotations

import argparse
import base64
import errno
import os
import re
import shlex
import stat
import struct
from contextlib import suppress
from pathlib import Path

from node_account_audit import audit_membership, audit_privilege_policies
from node_account_io import (
    FileSnapshot,
    HardeningError,
    account_lock,
    assert_unchanged,
    atomic_write,
    backup_once,
    open_private_temporary,
    read_mounted_secret,
    read_regular,
)

HOST_ETC = Path("/host-etc")
HOST_HOME = Path("/host-home")
HOST_K3S = Path("/host-k3s")
HOST_KUBELET = Path("/host-kubelet")
HOST_RUN_K3S = Path("/host-run-k3s")
HOST_RUN_CONTAINERD = Path("/host-run-containerd")
HOST_POLKIT_SHARE = Path("/host-polkit-share")
ACCOUNT = "hermes-agent"
ACCOUNT_UID = 1200
ACCOUNT_GID = 1200
ACCOUNT_HOME = "/home/hermes-agent"
ACCOUNT_SHELL = "/bin/bash"
HOST_ROOT_UID = 0
HOST_ROOT_GID = 0
LEGACY_ACCOUNTS = ("atlas", "oceanus")
MAX_ACCOUNT_FILE = 2 * 1024 * 1024
MAX_AUTHORIZED_KEYS = 1024 * 1024
ACL_VERSION = 2
ACL_UNDEFINED_ID = 0xFFFFFFFF
ACL_USER_OBJ = 0x01
ACL_USER = 0x02
ACL_GROUP_OBJ = 0x04
ACL_GROUP = 0x08
ACL_MASK = 0x10
ACL_OTHER = 0x20
ACL_HEADER = struct.Struct("<I")
ACL_ENTRY = struct.Struct("<HHI")
ACL_XATTR = "system.posix_acl_access"


def _read_regular(
    path: Path, maximum: int = MAX_ACCOUNT_FILE
) -> tuple[bytes, os.stat_result]:
    snapshot = read_regular(path, maximum)
    return snapshot.value, path.stat(follow_symlinks=False)


def _records(value: bytes, fields: int, name: str) -> list[list[str]]:
    try:
        text = value.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise HardeningError(f"{name} is not UTF-8") from exc
    if not text.endswith("\n"):
        raise HardeningError(f"{name} is missing its final newline")
    records = []
    names = set()
    for line in text.splitlines():
        parts = line.split(":")
        if len(parts) != fields or not parts[0] or parts[0] in names:
            raise HardeningError(f"{name} has an invalid record")
        names.add(parts[0])
        records.append(parts)
    return records


def _encode(records: list[list[str]]) -> bytes:
    return ("\n".join(":".join(record) for record in records) + "\n").encode()


def _expected_records() -> dict[str, list[str]]:
    return {
        "passwd": [
            ACCOUNT,
            "x",
            str(ACCOUNT_UID),
            str(ACCOUNT_GID),
            "Hermes Agent",
            ACCOUNT_HOME,
            ACCOUNT_SHELL,
        ],
        "group": [ACCOUNT, "x", str(ACCOUNT_GID), ""],
        "shadow": [ACCOUNT, "!", "1", "0", "99999", "7", "", "", ""],
        "gshadow": [ACCOUNT, "!", "", ""],
    }


def _reconcile_record(
    records: list[list[str]], expected: list[str], *, identity_index: int
) -> list[list[str]]:
    for record in records:
        same_name = record[0] == expected[0]
        same_identity = record[identity_index] == expected[identity_index]
        if same_name or same_identity:
            if record != expected:
                raise HardeningError("dedicated Hermes account identity conflicts")
            return records
    return [*records, expected]


def _reconcile_databases() -> None:
    expected = _expected_records()
    definitions = (
        ("passwd", 7, 2),
        ("group", 4, 2),
        ("shadow", 9, 0),
        ("gshadow", 4, 0),
    )
    planned: list[tuple[Path, FileSnapshot, bytes]] = []
    account_present = False
    parsed: dict[str, list[list[str]]] = {}
    for name, fields, identity_index in definitions:
        path = HOST_ETC / name
        snapshot = read_regular(path, MAX_ACCOUNT_FILE)
        records = _records(snapshot.value, fields, name)
        parsed[name] = records
        updated = _reconcile_record(
            records, expected[name], identity_index=identity_index
        )
        if name == "passwd":
            account_present = expected[name] in records
        planned.append((path, snapshot, _encode(updated)))
    audit_membership(ACCOUNT, parsed["group"], parsed["gshadow"])
    audit_privilege_policies(
        ACCOUNT, ACCOUNT_UID, HOST_ROOT_UID, HOST_ETC, HOST_POLKIT_SHARE
    )
    try:
        (HOST_HOME / ACCOUNT).lstat()
    except FileNotFoundError:
        pass
    else:
        if not account_present:
            raise HardeningError("dedicated Hermes account home already exists")
    for path, snapshot, _updated in planned:
        assert_unchanged(path, snapshot, MAX_ACCOUNT_FILE)
    for path, snapshot, _updated in planned:
        backup_once(path, snapshot, MAX_ACCOUNT_FILE)
    for path, snapshot, _updated in planned:
        assert_unchanged(path, snapshot, MAX_ACCOUNT_FILE)
    written: list[tuple[Path, FileSnapshot, bytes]] = []
    try:
        for path, snapshot, updated in planned:
            assert_unchanged(path, snapshot, MAX_ACCOUNT_FILE)
            atomic_write(path, updated, snapshot)
            written.append((path, snapshot, updated))
        for path, _snapshot, _updated in planned:
            name = path.name
            fields = next(item[1] for item in definitions if item[0] == name)
            records = _records(_read_regular(path)[0], fields, name)
            if expected[name] not in records:
                raise HardeningError("dedicated Hermes account validation failed")
    except Exception as error:
        for path, snapshot, updated in reversed(written):
            current = read_regular(path, MAX_ACCOUNT_FILE)
            if (
                current.value != updated
                or current.mode != snapshot.mode
                or current.uid != snapshot.uid
                or current.gid != snapshot.gid
                or current.xattrs != snapshot.xattrs
            ):
                raise HardeningError(
                    f"concurrent host account change prevents rollback: {path.name}"
                ) from error
            atomic_write(path, snapshot.value, snapshot)
        raise


SUPPORTED_KEY_TYPES = {"ssh-ed25519", "ecdsa-sha2-nistp256", "ssh-rsa"}
KEY_TYPE_RE = re.compile(r"(?:sk-)?(?:ssh|ecdsa)-[A-Za-z0-9@._+-]+\Z")


def _key_identity(line: bytes) -> tuple[str, bytes] | None:
    """Return OpenSSH key type/blob, ignoring options and comments."""
    stripped = line.strip(b"\r\n")
    if not stripped or stripped.lstrip().startswith(b"#"):
        return None
    try:
        fields = shlex.split(stripped.decode("ascii"), posix=True)
    except (UnicodeDecodeError, ValueError) as exc:
        raise HardeningError("authorized key line is malformed") from exc
    indexes = [
        index for index, field in enumerate(fields) if KEY_TYPE_RE.fullmatch(field)
    ]
    if len(indexes) != 1 or indexes[0] + 1 >= len(fields):
        raise HardeningError("authorized key line has ambiguous key material")
    key_type = fields[indexes[0]]
    try:
        blob = base64.b64decode(fields[indexes[0] + 1], validate=True)
    except ValueError as exc:
        raise HardeningError("authorized key payload is invalid") from exc
    encoded_type = key_type.encode("ascii")
    if len(blob) < 4 or int.from_bytes(blob[:4], "big") != len(encoded_type):
        raise HardeningError("authorized key blob is malformed")
    if blob[4 : 4 + len(encoded_type)] != encoded_type:
        raise HardeningError("authorized key type does not match its blob")
    return key_type, blob


def _validated_public_key(path: Path) -> tuple[bytes, tuple[str, bytes]]:
    # Vault CSI projects the key as a symlink chain inside its own mount.
    value = read_mounted_secret(path, path.parent, 16 * 1024).value
    line = value.strip()
    if b"\n" in line or b"\r" in line:
        raise HardeningError("Hermes public key must contain one line")
    identity = _key_identity(line)
    fields = line.split()
    if (
        identity is None
        or identity[0] not in SUPPORTED_KEY_TYPES
        or len(fields) not in {2, 3}
        or fields[0].decode("ascii") != identity[0]
    ):
        raise HardeningError("Hermes public key format is unsupported")
    return line, identity


def _directory(path: Path, *, mode: int, uid: int, gid: int) -> None:
    with suppress(FileExistsError):
        path.mkdir(mode=mode)
    metadata = path.lstat()
    if not stat.S_ISDIR(metadata.st_mode) or stat.S_ISLNK(metadata.st_mode):
        raise HardeningError(f"unsafe account directory: {path.name}")
    if metadata.st_uid not in {0, uid} or metadata.st_gid not in {0, gid}:
        raise HardeningError(f"account directory ownership conflicts: {path.name}")
    os.chown(path, uid, gid)
    path.chmod(mode)


def _without_key(value: bytes, identity: tuple[str, bytes]) -> bytes:
    kept = []
    for line in value.splitlines(keepends=True):
        try:
            line_identity = _key_identity(line)
        except HardeningError:
            # authorized_keys is shared with human operators on the legacy
            # accounts.  An unrelated line that this reconciler cannot prove
            # valid is not authority to delete it or block all reconciliation.
            kept.append(line)
            continue
        if line_identity != identity:
            kept.append(line)
    return b"".join(kept)


def _move_key(public_key: Path) -> None:
    key, identity = _validated_public_key(public_key)
    target = HOST_HOME / ACCOUNT / ".ssh" / "authorized_keys"
    legacy_plan: list[tuple[Path, FileSnapshot, bytes]] = []
    for legacy in LEGACY_ACCOUNTS:
        authorized = HOST_HOME / legacy / ".ssh" / "authorized_keys"
        try:
            snapshot = read_regular(authorized, MAX_AUTHORIZED_KEYS)
        except FileNotFoundError:
            continue
        updated = _without_key(snapshot.value, identity)
        if updated != snapshot.value:
            legacy_plan.append((authorized, snapshot, updated))
    if legacy_plan and target.exists():
        target_snapshot = read_regular(target, MAX_AUTHORIZED_KEYS)
        target_without_key = _without_key(target_snapshot.value, identity)
        if target_without_key != target_snapshot.value:
            backup_once(target, target_snapshot, MAX_AUTHORIZED_KEYS)
            assert_unchanged(target, target_snapshot, MAX_AUTHORIZED_KEYS)
            atomic_write(target, target_without_key, target_snapshot)
            verified_target = read_regular(target, MAX_AUTHORIZED_KEYS).value
            if _without_key(verified_target, identity) != verified_target:
                raise HardeningError("preexisting Hermes authorization removal failed")
    for authorized, snapshot, _updated in legacy_plan:
        assert_unchanged(authorized, snapshot, MAX_AUTHORIZED_KEYS)
        backup_once(authorized, snapshot, MAX_AUTHORIZED_KEYS)
    for authorized, snapshot, updated in legacy_plan:
        assert_unchanged(authorized, snapshot, MAX_AUTHORIZED_KEYS)
        atomic_write(authorized, updated, snapshot)
    for legacy in LEGACY_ACCOUNTS:
        authorized = HOST_HOME / legacy / ".ssh" / "authorized_keys"
        try:
            value = read_regular(authorized, MAX_AUTHORIZED_KEYS).value
        except FileNotFoundError:
            continue
        if _without_key(value, identity) != value:
            raise HardeningError("legacy Hermes authorization key removal failed")

    # Installation happens only after every legacy authorization is absent. A
    # crash can therefore remove access temporarily, but can never duplicate
    # the machine credential across privileged and unprivileged accounts.
    home = HOST_HOME / ACCOUNT
    ssh = home / ".ssh"
    _directory(home, mode=0o700, uid=ACCOUNT_UID, gid=ACCOUNT_GID)
    _directory(ssh, mode=0o700, uid=ACCOUNT_UID, gid=ACCOUNT_GID)
    target = ssh / "authorized_keys"
    if target.exists():
        current = read_regular(target, MAX_AUTHORIZED_KEYS)
        if current.value != key + b"\n":
            backup_once(target, current, MAX_AUTHORIZED_KEYS)
    else:
        current = FileSnapshot(
            value=b"",
            device=0,
            inode=0,
            mode=0o600,
            uid=ACCOUNT_UID,
            gid=ACCOUNT_GID,
            size=0,
            mtime_ns=0,
            ctime_ns=0,
            xattrs=(),
        )
    atomic_write(target, key + b"\n", current)
    if _read_regular(target, MAX_AUTHORIZED_KEYS)[0] != key + b"\n":
        raise HardeningError("dedicated Hermes authorized key validation failed")


def _decode_acl(value: bytes, mode: int) -> list[tuple[int, int, int]]:
    """Decode a bounded POSIX ACL or derive one from ordinary mode bits."""
    if not value:
        group = (mode >> 3) & 0o7
        return [
            (ACL_USER_OBJ, (mode >> 6) & 0o7, ACL_UNDEFINED_ID),
            (ACL_GROUP_OBJ, group, ACL_UNDEFINED_ID),
            (ACL_MASK, group, ACL_UNDEFINED_ID),
            (ACL_OTHER, mode & 0o7, ACL_UNDEFINED_ID),
        ]
    if len(value) < ACL_HEADER.size or (len(value) - ACL_HEADER.size) % ACL_ENTRY.size:
        raise HardeningError("sensitive directory ACL is malformed")
    if ACL_HEADER.unpack_from(value)[0] != ACL_VERSION:
        raise HardeningError("sensitive directory ACL version is unsupported")
    entries = [
        ACL_ENTRY.unpack_from(value, offset)
        for offset in range(ACL_HEADER.size, len(value), ACL_ENTRY.size)
    ]
    if any(permission > 0o7 for _tag, permission, _identifier in entries):
        raise HardeningError("sensitive directory ACL permission is malformed")
    required = {ACL_USER_OBJ, ACL_GROUP_OBJ, ACL_OTHER}
    if not required <= {tag for tag, _permission, _identifier in entries}:
        raise HardeningError("sensitive directory ACL is incomplete")
    return entries


def _encode_acl(entries: list[tuple[int, int, int]]) -> bytes:
    return ACL_HEADER.pack(ACL_VERSION) + b"".join(
        ACL_ENTRY.pack(*entry) for entry in entries
    )


def _acl_with_deny(value: bytes, mode: int) -> bytes:
    entries = _decode_acl(value, mode)
    entries = [
        entry
        for entry in entries
        if not (entry[0] == ACL_USER and entry[2] == ACCOUNT_UID)
    ]
    entries.append((ACL_USER, 0, ACCOUNT_UID))
    if not any(tag == ACL_MASK for tag, _permission, _identifier in entries):
        entries.append((ACL_MASK, (mode >> 3) & 0o7, ACL_UNDEFINED_ID))
    order = {
        ACL_USER_OBJ: 0,
        ACL_USER: 1,
        ACL_GROUP_OBJ: 2,
        ACL_GROUP: 3,
        ACL_MASK: 4,
        ACL_OTHER: 5,
    }
    entries.sort(key=lambda entry: (order.get(entry[0], 99), entry[2]))
    return _encode_acl(entries)


def _read_acl(path: Path) -> bytes:
    try:
        return os.getxattr(path, ACL_XATTR, follow_symlinks=False)
    except OSError as exc:
        if exc.errno in {errno.ENODATA, getattr(errno, "ENOATTR", errno.ENODATA)}:
            return b""
        raise


def _validate_acl_backup(value: bytes) -> None:
    if value[:1] not in {b"A", b"N"}:
        raise HardeningError("sensitive directory ACL backup is malformed")
    if value[:1] == b"A":
        if len(value) == 1:
            raise HardeningError("sensitive directory ACL backup is malformed")
        _decode_acl(value[1:], 0)
    elif value != b"N":
        raise HardeningError("sensitive directory ACL backup is malformed")


def _acl_backup_once(path: Path, value: bytes, backup_name: str | None = None) -> None:
    """Persist the original ACL as a private, durable host recovery file."""
    root = HOST_ETC / "hermes-node-boundary"
    _directory(root, mode=0o700, uid=HOST_ROOT_UID, gid=HOST_ROOT_GID)
    backup = root / f"{backup_name or path.name}.acl"
    encoded = b"A" + value if value else b"N"
    try:
        existing, _metadata = _read_regular(backup, 64 * 1024)
    except FileNotFoundError:
        pass
    else:
        _validate_acl_backup(existing)
        return

    temporary, descriptor = open_private_temporary(
        root / f"{path.name}.acl", 0o600
    )
    try:
        if os.write(descriptor, encoded) != len(encoded):
            raise HardeningError("short sensitive directory ACL backup write")
        os.fchmod(descriptor, 0o600)
        os.fchown(descriptor, HOST_ROOT_UID, HOST_ROOT_GID)
        os.fsync(descriptor)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise
    finally:
        os.close(descriptor)
    try:
        os.link(temporary, backup, follow_symlinks=False)
    except FileExistsError:
        # Another reconciler won the race; never overwrite the first backup.
        pass
    finally:
        temporary.unlink(missing_ok=True)
    stored, _metadata = _read_regular(backup, 64 * 1024)
    _validate_acl_backup(stored)
    directory = os.open(root, os.O_RDONLY)
    try:
        os.fsync(directory)
    finally:
        os.close(directory)


def _deny_sensitive_root(path: Path, backup_name: str | None = None) -> None:
    metadata = path.lstat()
    if not stat.S_ISDIR(metadata.st_mode) or stat.S_ISLNK(metadata.st_mode):
        raise HardeningError(f"unsafe sensitive directory: {path.name}")
    if metadata.st_uid != HOST_ROOT_UID:
        raise HardeningError(f"sensitive directory is not root-owned: {path.name}")
    current = _read_acl(path)
    _acl_backup_once(path, current, backup_name)
    updated = _acl_with_deny(current, stat.S_IMODE(metadata.st_mode))
    try:
        os.setxattr(path, ACL_XATTR, updated, follow_symlinks=False)
        verified = _decode_acl(_read_acl(path), stat.S_IMODE(metadata.st_mode))
        if (ACL_USER, 0, ACCOUNT_UID) not in verified:
            raise HardeningError("sensitive directory ACL validation failed")
    except Exception:
        if current:
            os.setxattr(path, ACL_XATTR, current, follow_symlinks=False)
        else:
            try:
                os.removexattr(path, ACL_XATTR, follow_symlinks=False)
            except OSError as exc:
                if exc.errno not in {
                    errno.ENODATA,
                    getattr(errno, "ENOATTR", errno.ENODATA),
                }:
                    raise
        raise


def _deny_sensitive_roots() -> None:
    for path, backup_name in (
        (HOST_K3S, "var-lib-rancher-k3s"),
        (HOST_KUBELET, "var-lib-kubelet"),
        (HOST_RUN_K3S, "run-k3s"),
        (HOST_RUN_CONTAINERD, "run-containerd"),
    ):
        _deny_sensitive_root(path, backup_name)


def reconcile(public_key: Path) -> None:
    """Create the locked account and move only the Hermes authorization key."""
    with account_lock(HOST_ETC, expected_uid=HOST_ROOT_UID):
        _reconcile_databases()
        _deny_sensitive_roots()
        _move_key(public_key)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--public-key-file", type=Path, required=True)
    args = parser.parse_args()
    reconcile(args.public_key_file)
    print("Dedicated Hermes node account reconciled.", flush=True)
    return 0


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