#!/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 stat
import struct
from pathlib import Path

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")
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"


class HardeningError(RuntimeError):
    """Raised before an unsafe or ambiguous host-account change."""


def _read_regular(path: Path, maximum: int = MAX_ACCOUNT_FILE) -> tuple[bytes, os.stat_result]:
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(path, flags)
    try:
        metadata = os.fstat(descriptor)
        if not stat.S_ISREG(metadata.st_mode) or metadata.st_size > maximum:
            raise HardeningError(f"unsafe regular file: {path.name}")
        value = os.read(descriptor, maximum + 1)
    finally:
        os.close(descriptor)
    if len(value) != metadata.st_size:
        raise HardeningError(f"short file read: {path.name}")
    return value, metadata


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 _backup_once(path: Path, value: bytes, metadata: os.stat_result) -> Path:
    backup = path.with_name(path.name + ".hermes-boundary-backup")
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(backup, flags, stat.S_IMODE(metadata.st_mode))
    except FileExistsError:
        _read_regular(backup)
        return backup
    try:
        if os.write(descriptor, value) != len(value):
            raise HardeningError(f"short backup write: {path.name}")
        os.fchown(descriptor, metadata.st_uid, metadata.st_gid)
        os.fsync(descriptor)
    finally:
        os.close(descriptor)
    return backup


def _atomic_write(
    path: Path, value: bytes, *, mode: int, uid: int, gid: int
) -> None:
    temporary = path.with_name(f".{path.name}.hermes-{os.getpid()}")
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(temporary, flags, mode)
    try:
        if os.write(descriptor, value) != len(value):
            raise HardeningError(f"short atomic write: {path.name}")
        os.fchmod(descriptor, mode)
        os.fchown(descriptor, uid, gid)
        os.fsync(descriptor)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise
    finally:
        os.close(descriptor)
    os.replace(temporary, path)
    directory = os.open(path.parent, os.O_RDONLY)
    try:
        os.fsync(directory)
    finally:
        os.close(directory)


def _reconcile_databases() -> None:
    expected = _expected_records()
    definitions = (
        ("passwd", 7, 2),
        ("group", 4, 2),
        ("shadow", 9, 0),
        ("gshadow", 4, 0),
    )
    planned = []
    account_present = False
    for name, fields, identity_index in definitions:
        path = HOST_ETC / name
        value, metadata = _read_regular(path)
        records = _records(value, fields, name)
        updated = _reconcile_record(
            records, expected[name], identity_index=identity_index
        )
        if name == "passwd":
            account_present = expected[name] in records
        planned.append((path, value, metadata, _encode(updated)))
    try:
        (HOST_HOME / ACCOUNT).lstat()
    except FileNotFoundError:
        pass
    else:
        if not account_present:
            raise HardeningError("dedicated Hermes account home already exists")
    for path, value, metadata, _updated in planned:
        _backup_once(path, value, metadata)
    try:
        for path, _old, metadata, updated in planned:
            _atomic_write(
                path,
                updated,
                mode=stat.S_IMODE(metadata.st_mode),
                uid=metadata.st_uid,
                gid=metadata.st_gid,
            )
        for path, _old, _metadata, _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:
        for path, old, metadata, _updated in planned:
            _atomic_write(
                path,
                old,
                mode=stat.S_IMODE(metadata.st_mode),
                uid=metadata.st_uid,
                gid=metadata.st_gid,
            )
        raise


def _validated_public_key(path: Path) -> bytes:
    value, _ = _read_regular(path, 16 * 1024)
    line = value.strip()
    if b"\n" in line or b"\r" in line:
        raise HardeningError("Hermes public key must contain one line")
    fields = line.split()
    if len(fields) not in {2, 3} or fields[0] not in {
        b"ssh-ed25519",
        b"ecdsa-sha2-nistp256",
        b"ssh-rsa",
    }:
        raise HardeningError("Hermes public key format is unsupported")
    try:
        base64.b64decode(fields[1], validate=True)
    except ValueError as exc:
        raise HardeningError("Hermes public key payload is invalid") from exc
    return line


def _directory(path: Path, *, mode: int, uid: int, gid: int) -> None:
    try:
        path.mkdir(mode=mode)
    except FileExistsError:
        pass
    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, key: bytes) -> bytes:
    return b"".join(
        line
        for line in value.splitlines(keepends=True)
        if line.strip(b"\r\n") != key
    )


def _move_key(public_key: Path) -> None:
    key = _validated_public_key(public_key)
    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, metadata = _read_regular(target, MAX_AUTHORIZED_KEYS)
        if current != key + b"\n":
            _backup_once(target, current, metadata)
    _atomic_write(target, key + b"\n", mode=0o600, uid=ACCOUNT_UID, gid=ACCOUNT_GID)
    for legacy in LEGACY_ACCOUNTS:
        authorized = HOST_HOME / legacy / ".ssh" / "authorized_keys"
        try:
            value, metadata = _read_regular(authorized, MAX_AUTHORIZED_KEYS)
        except FileNotFoundError:
            continue
        updated = _without_key(value, key)
        if updated == value:
            continue
        _backup_once(authorized, value, metadata)
        _atomic_write(
            authorized,
            updated,
            mode=stat.S_IMODE(metadata.st_mode),
            uid=metadata.st_uid,
            gid=metadata.st_gid,
        )
    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) -> 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"{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 = root / f".{path.name}.acl.hermes-{os.getpid()}"
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(temporary, flags, 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) -> 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)
    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 in (HOST_K3S, HOST_KUBELET, HOST_RUN_K3S, HOST_RUN_CONTAINERD):
        _deny_sensitive_root(path)


def reconcile(public_key: Path) -> None:
    """Create the locked account and move only the Hermes authorization key."""
    _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())