#!/usr/bin/env python3
"""Age-only SOPS recovery to a new private file; no plaintext process output."""
from __future__ import annotations

import json
import fcntl
import os
from pathlib import Path
import re
import shutil
import signal
import stat
import subprocess
import sys
import tempfile
import uuid

ROOT = Path(__file__).resolve().parents[2]
RECIPIENT = re.compile(r'age1[023456789acdefghjklmnpqrstuvwxyz]{58}\Z')


class Refused(Exception):
    pass


def public_recipient():
    lines = (ROOT / 'keys/age.pub').read_text().splitlines()
    values = [line.strip() for line in lines if line.strip() and not line.startswith('#')]
    if len(values) != 1 or not RECIPIENT.fullmatch(values[0]):
        raise Refused('registered public recipient is invalid')
    result = subprocess.run(['age', '-r', values[0]], input=b'', stdout=subprocess.DEVNULL,
                            stderr=subprocess.DEVNULL, timeout=10)
    if result.returncode:
        raise Refused('registered public recipient is invalid')
    return values[0]


def unique_object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise Refused('duplicate encrypted document key')
        result[key] = value
    return result


def parse_document(raw, suffix):
    if suffix == '.json':
        data = json.loads(raw, object_pairs_hook=unique_object)
        kind = 'json'
    elif suffix in {'.yaml', '.yml'}:
        import yaml
        class UniqueLoader(yaml.SafeLoader):
            def construct_mapping(self, node, deep=False):
                self.flatten_mapping(node)
                return unique_object((self.construct_object(key, deep=deep),
                                      self.construct_object(value, deep=deep))
                                     for key, value in node.value)
        try:
            data = yaml.load(raw, Loader=UniqueLoader)
        except yaml.YAMLError:
            raise Refused('invalid encrypted YAML input') from None
        kind = 'yaml'
    else:
        raise Refused('age recovery supports SOPS JSON/YAML input')
    return data, kind


def age_document(source, recipient, *, raw=None):
    # Inspect ciphertext metadata before giving a credential to any subprocess.
    # This repository's admitted SOPS profile is age-only, not arbitrary cloud
    # KMS, key services, SSH identities or executable age plugins.
    if raw is None:
        with source.open("rb") as encrypted:
            raw = encrypted.read(8 * 1024 * 1024 + 1)
    if len(raw) > 8 * 1024 * 1024:
        raise Refused('encrypted input is too large')
    data, kind = parse_document(raw, source.suffix)
    metadata = data.get('sops') if isinstance(data, dict) else None
    if not isinstance(metadata, dict):
        raise Refused('SOPS metadata is missing')
    known = {'age', 'kms', 'gcp_kms', 'azure_kv', 'hc_vault', 'hckms', 'pgp',
             'key_groups', 'shamir_threshold', 'lastmodified', 'mac', 'version',
             'unencrypted_suffix', 'encrypted_suffix', 'unencrypted_regex',
             'encrypted_regex', 'unencrypted_comment_regex', 'encrypted_comment_regex',
             'mac_only_encrypted'}
    if set(metadata) - known or metadata.get('key_groups'):
        raise Refused('unsupported SOPS custody profile')
    if any(metadata.get(key) for key in ('kms', 'gcp_kms', 'azure_kv', 'hc_vault', 'hckms', 'pgp')):
        raise Refused('recovery requires age-only metadata')
    recipients = metadata.get('age')
    if not isinstance(recipients, list) or not recipients:
        raise Refused('age recipient metadata is missing')
    if any(not isinstance(row, dict) or not RECIPIENT.fullmatch(str(row.get('recipient', '')))
           for row in recipients):
        raise Refused('unsupported age recipient')
    if recipient not in [row['recipient'] for row in recipients]:
        raise Refused('document does not include the registered recipient')
    return raw, kind


def private_key_fd():
    inherited = os.environ.get('NETKINGDOM_CUSTODY_AGE_KEY_FD', '')
    name = os.environ.get('NETKINGDOM_CUSTODY_AGE_KEY_FILE', '')
    if inherited:
        if name or not re.fullmatch(r'[1-9][0-9]{0,8}',inherited) or int(inherited)<3:
            raise Refused('ambiguous custody descriptor')
        fd=os.dup(int(inherited))
        try:
            info=os.fstat(fd)
            required=fcntl.F_SEAL_WRITE|fcntl.F_SEAL_GROW|fcntl.F_SEAL_SHRINK|fcntl.F_SEAL_SEAL
            if (not stat.S_ISREG(info.st_mode) or info.st_uid!=os.getuid() or info.st_mode&0o077
                    or info.st_nlink!=0 or not 0<info.st_size<=16384
                    or fcntl.fcntl(fd,fcntl.F_GET_SEALS)!=required):
                raise Refused('custody descriptor must be anonymous, private and sealed')
            # No named file exists to survive parent SIGKILL. Existing callers
            # own the returned duplicate and verify the expected age recipient.
            return fd
        except BaseException:os.close(fd);raise

    if not name:
        raise Refused('explicit temporary custody key required')
    path = Path(os.path.abspath(name))
    if any(p.is_symlink() for p in (path, *path.parents)):
        raise Refused('custody key path must not contain symlinks')
    fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
    info = os.fstat(fd)
    if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or info.st_mode & 0o077 or info.st_nlink != 1 or info.st_size > 16384:
        os.close(fd)
        raise Refused('custody key must be an owned private regular file')
    return fd


def decrypt(recipient):
    source_name = os.environ.get('NK_SOPS_FILE', '')
    output_name = os.environ.get('NK_SOPS_OUTPUT', '')
    if not source_name or not output_name:
        raise Refused('supply FILE and OUTPUT; plaintext stdout is disabled')
    source = Path(source_name).resolve(strict=True)
    output = Path(output_name).absolute()
    parent = output.parent.resolve(strict=True)
    if any((p / '.git').is_file() or (p / '.git/HEAD').is_file()
           or ((p / 'HEAD').is_file() and (p / 'objects').is_dir() and (p / 'refs').is_dir())
           for p in [parent, *parent.parents]):
        raise Refused('plaintext output must be outside Git worktrees')
    raw, kind = age_document(source, recipient)
    binary = shutil.which('sops')
    if binary is None:
        raise Refused('sops executable is unavailable')
    directory = os.open(parent, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
    temporary = '.netkingdom-sops-' + uuid.uuid4().hex
    keyfd = None
    try:
        info = os.fstat(directory)
        if info.st_uid != os.getuid() or info.st_mode & 0o077:
            raise Refused('OUTPUT requires an owned private parent directory')
        try:
            os.stat(output.name, dir_fd=directory, follow_symlinks=False)
        except FileNotFoundError:
            pass
        else:
            raise Refused('OUTPUT already exists')
        keyfd = private_key_fd()
        derived = subprocess.run(['age-keygen', '-y', f'/proc/self/fd/{keyfd}'],
                                 stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
                                 pass_fds=(keyfd,), timeout=10)
        if derived.returncode or derived.stdout.strip() != recipient.encode():
            raise Refused('custody key does not match the registered recipient')
        fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW,
                     0o600, dir_fd=directory)
        with os.fdopen(fd, 'wb') as plain, tempfile.TemporaryDirectory(prefix='netkingdom-sops-home-') as isolated:
            # HOME/XDG isolation matters: SOPS loads ambient keys in addition to
            # an explicit key file. No inline keys, SSH keys, key commands,
            # proxies or remote keyservice configuration are inherited.
            env = {'PATH': os.environ.get('PATH', os.defpath), 'HOME': isolated,
                   'XDG_CONFIG_HOME': isolated, 'GNUPGHOME': isolated,
                   'SOPS_AGE_KEY_FILE': f'/proc/self/fd/{keyfd}',
                   'AWS_EC2_METADATA_DISABLED': 'true'}
            result = subprocess.run([binary, '--decrypt', '--input-type', kind,
                                     '--output-type', kind, '/dev/stdin'],
                                    input=raw, stdout=plain, stderr=subprocess.DEVNULL,
                                    env=env, cwd=isolated, pass_fds=(keyfd,), timeout=30)
            if result.returncode:
                raise Refused('SOPS decryption failed; no plaintext published')
            plain.flush()
            os.fsync(plain.fileno())
        # Atomic no-overwrite publication within the private directory.
        os.link(temporary, output.name, src_dir_fd=directory, dst_dir_fd=directory,
                follow_symlinks=False)
        os.unlink(temporary, dir_fd=directory)
        os.fsync(directory)
    finally:
        if keyfd is not None:
            os.close(keyfd)
        try:
            os.unlink(temporary, dir_fd=directory)
        except FileNotFoundError:
            pass
        os.close(directory)
    print('Decrypted to the requested private file. Remove it after the approved operation.')


def interrupted(_signum, _frame):
    raise Refused('operation interrupted')


def main():
    for sig in (signal.SIGTERM, signal.SIGHUP, signal.SIGINT):
        signal.signal(sig, interrupted)
    try:
        if len(sys.argv) != 2 or sys.argv[1] not in {'setup', 'decrypt'}:
            raise Refused('unsupported custody operation')
        recipient = public_recipient()
        if sys.argv[1] == 'setup':
            print('Registered public recipient validated. No private key was read, copied or generated.')
        else:
            decrypt(recipient)
        return 0
    except Refused as exc:
        print('ERROR: ' + str(exc), file=sys.stderr)
    except (OSError, ValueError, TypeError, ImportError, subprocess.SubprocessError):
        # Parser, filesystem and tool errors can include user data. Never echo.
        print('ERROR: custody operation failed; input and tool diagnostics suppressed.', file=sys.stderr)
    return 1


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