#!/usr/bin/python3
"""Restricted privileged helper for AnduinOS YubiKey Security Center."""

import fcntl
import json
import os
import pwd
import re
import subprocess
import sys
import tempfile

CONFIG_DIR = "/etc/anduinos-yubikey-manager"
GDM_MAPPING = CONFIG_DIR + "/u2f_mappings"
SUDO_MAPPING = CONFIG_DIR + "/sudo_u2f_mappings"
METADATA = CONFIG_DIR + "/enrollments.json"
PAM_CONFIG = {
    "gdm": ("/etc/pam.d/gdm-password", GDM_MAPPING),
    "sudo": ("/etc/pam.d/sudo", SUDO_MAPPING),
}
SUDOERS_DIR = "/etc/sudoers.d"
MANAGED_SUDOERS = SUDOERS_DIR + "/90-anduinos-passwordless-admin"
SUDO_STATE = "/var/lib/anduinos-passwordless-sudo/users"
USER_RE = re.compile(r"^[a-z_][a-z0-9_-]{0,31}$")
SERIAL_RE = re.compile(r"^(?:[0-9]{1,12}|usb-[A-Za-z0-9._-]{1,64})$")
CREDENTIAL_RE = re.compile(
    r"^[A-Za-z0-9_=/+.-]+(?:,[A-Za-z0-9_=/+.-]*){1,3}$"
)


def fail(message):
    print(message, file=sys.stderr)
    raise SystemExit(1)


def atomic_write(path, content, mode):
    directory = os.path.dirname(path)
    os.makedirs(directory, mode=0o755, exist_ok=True)
    fd, temporary = tempfile.mkstemp(prefix=".yubikey-", dir=directory)
    try:
        os.fchmod(fd, mode)
        with os.fdopen(fd, "w", encoding="utf-8") as output:
            output.write(content)
            output.flush()
            os.fsync(output.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def require_safe_optional_file(path):
    if os.path.islink(path) or (os.path.exists(path) and not os.path.isfile(path)):
        fail("A managed security policy path is unsafe.")


def read_metadata():
    try:
        with open(METADATA, encoding="utf-8") as source:
            value = json.load(source)
    except FileNotFoundError:
        value = {"version": 1, "enrollments": []}
    if not isinstance(value, dict) or not isinstance(value.get("enrollments"), list):
        fail("The enrollment metadata is invalid.")
    return value


def read_mappings(path):
    records = {}
    try:
        with open(path, encoding="utf-8") as source:
            lines = source.readlines()
    except FileNotFoundError:
        return records
    for raw in lines:
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        username, separator, credentials = line.partition(":")
        if not separator or not USER_RE.fullmatch(username):
            fail("The existing PAM mapping file is invalid.")
        records[username] = [item for item in credentials.split(":") if item]
    return records


def write_state(metadata, mappings, mapping_path):
    mapping_text = "".join(
        username + ":" + ":".join(credentials) + "\n"
        for username, credentials in sorted(mappings.items())
        if credentials
    )
    atomic_write(mapping_path, mapping_text, 0o600)
    atomic_write(METADATA, json.dumps(metadata, indent=2, sort_keys=True) + "\n", 0o644)


def ensure_pam(purpose, enabled):
    pam_file, mapping = PAM_CONFIG[purpose]
    marker = f"# Managed by anduinos-yubikey-manager ({purpose})"
    pam_line = (
        f"auth sufficient pam_u2f.so authfile={mapping} "
        "origin=pam://anduinos appid=pam://anduinos cue"
    )
    try:
        with open(pam_file, encoding="utf-8") as source:
            lines = source.read().splitlines()
    except FileNotFoundError:
        fail("GDM PAM configuration was not found.")
    filtered = []
    skip_next = False
    for line in lines:
        if line in (marker, "# Managed by anduinos-yubikey-manager"):
            skip_next = True
            continue
        if skip_next and "pam_u2f.so" in line:
            skip_next = False
            continue
        skip_next = False
        filtered.append(line)
    if enabled:
        position = next(
            (index for index, line in enumerate(filtered) if "common-auth" in line),
            len(filtered),
        )
        filtered[position:position] = [marker, pam_line]
    atomic_write(pam_file, "\n".join(filtered) + "\n", 0o644)


def validate_common(username, serial):
    if not USER_RE.fullmatch(username):
        fail("Invalid username.")
    try:
        pwd.getpwnam(username)
    except KeyError:
        fail("The selected user does not exist.")
    if not SERIAL_RE.fullmatch(serial):
        fail("Invalid YubiKey serial number.")


def validate_user(username):
    if not USER_RE.fullmatch(username):
        fail("Invalid username.")
    try:
        pwd.getpwnam(username)
    except KeyError:
        fail("The selected user does not exist.")


def has_usable_password(username):
    try:
        with open("/etc/shadow", encoding="utf-8") as source:
            for line in source:
                fields = line.rstrip("\n").split(":")
                if fields[0] == username:
                    value = fields[1]
                    return bool(value) and not value.startswith(("!", "*"))
    except OSError:
        return False
    return False


def sudo_enrolled(metadata, username):
    return any(
        item.get("username") == username
        and item.get("purpose", "gdm") == "sudo"
        for item in metadata["enrollments"]
    )


def read_managed_passwordless_users():
    require_safe_optional_file(MANAGED_SUDOERS)
    try:
        with open(MANAGED_SUDOERS, encoding="utf-8") as source:
            lines = source.read().splitlines()
    except FileNotFoundError:
        return []
    rule = re.compile(
        r"^\s*([a-z_][a-z0-9_-]{0,31})\s+ALL\s*=\s*"
        r"\(ALL(?::ALL)?\)\s+NOPASSWD\s*:\s*ALL\s*$"
    )
    users = []
    for line in lines:
        if not line.strip() or line.lstrip().startswith("#"):
            continue
        match = rule.fullmatch(line)
        if not match:
            fail("The managed sudo policy contains an unexpected rule.")
        if match.group(1) not in users:
            users.append(match.group(1))
    return users


def sudoers_dropin_paths():
    try:
        names = os.listdir(SUDOERS_DIR)
    except FileNotFoundError:
        return
    for name in sorted(names):
        # Match sudo's @includedir policy: editor backups and names containing
        # a dot are ignored and therefore cannot contribute effective rules.
        if "." in name or name.endswith("~"):
            continue
        path = os.path.join(SUDOERS_DIR, name)
        if os.path.islink(path) or not os.path.isfile(path):
            continue
        yield path


def read_effective_passwordless_users():
    users = set()
    rule = re.compile(
        r"^\s*([a-z_][a-z0-9_-]{0,31})\s+ALL\s*=\s*"
        r"\(ALL(?::ALL)?\)\s+NOPASSWD\s*:\s*ALL\s*$"
    )
    for path in sudoers_dropin_paths():
        with open(path, encoding="utf-8") as source:
            for line in source:
                match = rule.fullmatch(line.rstrip("\n"))
                if match:
                    users.add(match.group(1))
    return sorted(users)


def write_passwordless_state(users):
    atomic_write(
        SUDO_STATE,
        "".join(f"{username}\n" for username in sorted(set(users))),
        0o644,
    )


def repair_integrations(metadata):
    read_managed_passwordless_users()
    require_safe_optional_file(SUDO_STATE)
    validate_sudoers()
    entries = metadata["enrollments"]
    for purpose, (_, mapping_path) in PAM_CONFIG.items():
        mappings = {}
        for item in entries:
            if item.get("purpose", "gdm") != purpose:
                continue
            username = item.get("username", "")
            serial = item.get("serial", "")
            credential = item.get("credential", "")
            validate_common(username, serial)
            if not CREDENTIAL_RE.fullmatch(credential):
                fail("The enrollment metadata contains an invalid PAM credential.")
            user_credentials = mappings.setdefault(username, [])
            if credential not in user_credentials:
                user_credentials.append(credential)
        mapping_text = "".join(
            username + ":" + ":".join(credentials) + "\n"
            for username, credentials in sorted(mappings.items())
            if credentials
        )
        atomic_write(mapping_path, mapping_text, 0o600)
        ensure_pam(purpose, bool(mappings))
    passwordless_users = read_effective_passwordless_users()
    metadata["passwordless_sudo_users"] = passwordless_users
    write_passwordless_state(passwordless_users)
    atomic_write(
        METADATA,
        json.dumps(metadata, indent=2, sort_keys=True) + "\n",
        0o644,
    )


def validate_sudoers():
    result = subprocess.run(
        ["/usr/sbin/visudo", "-cf", "/etc/sudoers"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        check=False,
    )
    if result.returncode:
        fail(result.stderr.strip() or "visudo rejected the sudo configuration.")


def set_passwordless_sudo(username, enabled, metadata):
    validate_user(username)
    require_safe_optional_file(MANAGED_SUDOERS)
    require_safe_optional_file(SUDO_STATE)
    if not enabled and not has_usable_password(username) and not sudo_enrolled(metadata, username):
        fail(
            "This account has no usable password. Enroll at least one sudo YubiKey "
            "before disabling passwordless sudo."
        )
    validate_sudoers()
    backups = {}
    try:
        if os.path.exists(METADATA):
            with open(METADATA, encoding="utf-8") as source:
                backups[METADATA] = (
                    source.read(),
                    os.stat(METADATA).st_mode & 0o777,
                )
        else:
            backups[METADATA] = None
        if os.path.exists(SUDO_STATE):
            with open(SUDO_STATE, encoding="utf-8") as source:
                backups[SUDO_STATE] = (
                    source.read(),
                    os.stat(SUDO_STATE).st_mode & 0o777,
                )
        else:
            backups[SUDO_STATE] = None
        passwordless_users = set(read_managed_passwordless_users())
        if enabled:
            passwordless_users.add(username)
            if os.path.exists(MANAGED_SUDOERS):
                with open(MANAGED_SUDOERS, encoding="utf-8") as source:
                    backups[MANAGED_SUDOERS] = (
                        source.read(),
                        os.stat(MANAGED_SUDOERS).st_mode & 0o777,
                    )
            else:
                backups[MANAGED_SUDOERS] = None
            atomic_write(
                MANAGED_SUDOERS,
                "".join(
                    f"{item} ALL=(ALL:ALL) NOPASSWD: ALL\n"
                    for item in sorted(passwordless_users)
                ),
                0o440,
            )
        else:
            passwordless_users.discard(username)
            escaped = re.escape(username)
            rule = re.compile(
                rf"^\s*{escaped}\s+ALL\s*=\s*\(ALL(?::ALL)?\)\s+"
                r"NOPASSWD\s*:\s*ALL\s*$"
            )
            for path in sudoers_dropin_paths():
                with open(path, encoding="utf-8") as source:
                    original = source.read()
                lines = original.splitlines()
                filtered = [line for line in lines if not rule.fullmatch(line)]
                if filtered == lines:
                    continue
                backups[path] = (original, os.stat(path).st_mode & 0o777)
                atomic_write(
                    path,
                    "\n".join(filtered) + ("\n" if filtered else ""),
                    backups[path][1],
                )
        validate_sudoers()
        effective_users = read_effective_passwordless_users()
        write_passwordless_state(effective_users)
        metadata["passwordless_sudo_users"] = effective_users
        atomic_write(
            METADATA,
            json.dumps(metadata, indent=2, sort_keys=True) + "\n",
            0o644,
        )
    except BaseException:
        for path, backup in backups.items():
            if backup is None:
                if os.path.exists(path):
                    os.unlink(path)
            else:
                atomic_write(path, backup[0], backup[1])
        raise


def install_git():
    """Install the one package required by the Git-signing page.

    Keep the package name and apt-get arguments inside this privileged helper so
    an unprivileged caller cannot turn the polkit action into a general package
    installation primitive.
    """
    environment = os.environ.copy()
    environment["DEBIAN_FRONTEND"] = "noninteractive"
    result = subprocess.run(
        ["/usr/bin/apt-get", "install", "-y", "git"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        check=False,
        env=environment,
    )
    if result.returncode:
        fail(result.stdout.strip() or "apt-get could not install Git.")


def main():
    if os.geteuid() != 0:
        fail("This helper must run as root.")
    if len(sys.argv) == 2 and sys.argv[1] == "repair":
        os.makedirs(CONFIG_DIR, mode=0o755, exist_ok=True)
        with open(CONFIG_DIR + "/.lock", "w", encoding="utf-8") as lock:
            fcntl.flock(lock, fcntl.LOCK_EX)
            repair_integrations(read_metadata())
        return
    if len(sys.argv) == 2 and sys.argv[1] == "install-git":
        install_git()
        return
    if len(sys.argv) == 4 and sys.argv[1] == "passwordless-sudo":
        username, mode = sys.argv[2:4]
        if mode not in ("enable", "disable"):
            fail("Invalid sudo policy.")
        os.makedirs(CONFIG_DIR, mode=0o755, exist_ok=True)
        with open(CONFIG_DIR + "/.lock", "w", encoding="utf-8") as lock:
            fcntl.flock(lock, fcntl.LOCK_EX)
            set_passwordless_sudo(username, mode == "enable", read_metadata())
        return
    if len(sys.argv) not in (5, 6):
        fail(
            "Usage: helper repair | helper install-git | "
            "helper enroll PURPOSE USER SERIAL CREDENTIAL | "
            "helper remove PURPOSE USER SERIAL"
        )
    action, purpose, username, serial = sys.argv[1:5]
    if purpose not in PAM_CONFIG:
        fail("Invalid authentication purpose.")
    validate_common(username, serial)
    if action == "enroll":
        if len(sys.argv) != 6 or not CREDENTIAL_RE.fullmatch(sys.argv[5]):
            fail("Invalid PAM credential.")
        credential = sys.argv[5]
    elif action == "remove":
        if len(sys.argv) != 5:
            fail("Unexpected arguments.")
        credential = None
    else:
        fail("Unknown operation.")

    os.makedirs(CONFIG_DIR, mode=0o755, exist_ok=True)
    lock_path = CONFIG_DIR + "/.lock"
    with open(lock_path, "w", encoding="utf-8") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        metadata = read_metadata()
        mapping_path = PAM_CONFIG[purpose][1]
        mappings = read_mappings(mapping_path)
        entries = metadata["enrollments"]

        if action == "enroll":
            if not any(
                item.get("username") == username
                and item.get("serial") == serial
                and item.get("purpose", "gdm") == purpose
                for item in entries
            ):
                entries.append(
                    {
                        "username": username,
                        "serial": serial,
                        "purpose": purpose,
                        "credential": credential,
                    }
                )
            user_credentials = mappings.setdefault(username, [])
            if credential not in user_credentials:
                user_credentials.append(credential)
        else:
            removed = [
                item.get("credential")
                for item in entries
                if item.get("username") == username
                and item.get("serial") == serial
                and item.get("purpose", "gdm") == purpose
            ]
            metadata["enrollments"] = [
                item for item in entries
                if not (
                    item.get("username") == username
                    and item.get("serial") == serial
                    and item.get("purpose", "gdm") == purpose
                )
            ]
            mappings[username] = [
                item for item in mappings.get(username, []) if item not in removed
            ]

        write_state(metadata, mappings, mapping_path)
        enabled = any(
            item.get("purpose", "gdm") == purpose
            for item in metadata["enrollments"]
        )
        ensure_pam(purpose, enabled)


if __name__ == "__main__":
    main()
