#!/usr/bin/python3
"""Explicit, fail-closed installer for an official nginx / GeoIP2 pair."""
import argparse
import contextlib
import functools
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile
import urllib.request


class Error(RuntimeError):
    pass


def run(*args, cwd=None, env=None):
    result = subprocess.run(args, cwd=cwd, env=env or {**os.environ, 'LC_ALL': 'C'},
                            text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if result.returncode:
        raise Error(f"Command failed: {' '.join(map(str, args))}\n{result.stderr}{result.stdout}")
    return result.stdout


def compare(a, b):
    if a == b:
        return 0
    return -1 if subprocess.run(['dpkg', '--compare-versions', a, 'lt', b],
                               check=False).returncode == 0 else 1


def release(html, legacy=False):
    title = 'Legacy versions' if legacy else 'Stable version'
    section = re.search(r'<h4>\s*' + title + r'\s*</h4>(.*?)(?=<h4>|$)', html, re.S)
    versions = set(re.findall(r'/download/nginx-(\d+\.\d+\.\d+)\.tar\.gz["\']',
                             section.group(1) if section else ''))
    if not versions:
        raise Error(f'Cannot identify {title} on nginx.org; refusing fallback.')
    return max(versions, key=functools.cmp_to_key(compare))


NGINX = 'https://nginx.org/packages/debian'
OURS = 'https://mmdb.netcube.bg/apt'
NGINX_FPR = '8540A6F18833A80E9C1653A42FD21310B49F6B46'
OUR_FPR = '3B432A371C6C86A8D1B8C3CC4539E8A767F07A71'


def primary_fingerprints(text):
    result, primary = set(), False
    for line in text.splitlines():
        fields = line.split(':')
        if fields[0] in ('pub', 'sub'):
            primary = fields[0] == 'pub'
        elif fields[0] == 'fpr' and primary:
            result.add(fields[9])
            primary = False
    return result


def validate_fingerprints(text, expected):
    if primary_fingerprints(text) != {expected}:
        raise Error('Exported keyring does not contain ONLY the pinned primary fingerprint.')


def verify_origin(policy, version, origin, suite):
    active, sources = False, []
    for line in policy.splitlines():
        header = re.match(r'^\s+(?:\*\*\*\s+)?(\S+)\s+(-?\d+)\s*$', line)
        if header:
            active = header[1] == version
        elif active and re.search(r'\bPackages\s*$', line):
            sources.append(line.split())
    if not sources or any(len(s) != 5 or s[1].rstrip('/') != origin or
                          s[2].split('/')[0] != suite or s[3] != 'amd64'
                          for s in sources):
        raise Error(f'{version}: absent or ambiguous package origin; expected {origin} {suite}.')


def records(text):
    result = []
    for paragraph in re.split(r'\n\s*\n', text.strip()):
        record, previous = {}, None
        for line in paragraph.splitlines():
            if line.startswith(' ') and previous:
                record[previous] += ' ' + line.strip()
            elif ': ' in line:
                previous, value = line.split(': ', 1)
                record[previous] = value
        if record:
            result.append(record)
    return result


def package_record(text, package, version):
    found = records(text)
    if not found or any(r.get('Package') != package or r.get('Version') != version or
                        r.get('Architecture') != 'amd64' or
                        not re.fullmatch('[0-9a-f]{64}', r.get('SHA256', '')) for r in found):
        raise Error(f'Invalid exact package metadata: {package}={version}')
    if len({r['SHA256'] for r in found}) != 1:
        raise Error('Conflicting package hashes for the same version.')
    return found[0]


def module_record(text, version, nginx_version):
    record = package_record(text, 'nginx-module-geoip2', version)
    dependencies = [d.strip() for d in record.get('Depends', '').split(',')]
    nginx_deps = [d for d in dependencies if re.search(r'\bnginx\b', d)]
    if nginx_deps != [f'nginx (= {nginx_version})']:
        raise Error('Module lacks the exact official nginx dependency.')
    return record


MARKER = '# Managed by nginx-geoip2-installer'
SOURCE_NAME = 'nginx-geoip2-installer.list'


def platform(text, arch):
    values = dict(line.split('=', 1) for line in text.splitlines() if '=' in line)
    values = {k: v.strip('"\'') for k, v in values.items()}
    if values.get('ID') != 'debian' or values.get('VERSION_CODENAME') not in ('bookworm', 'trixie') or arch != 'amd64':
        raise Error('Only Debian bookworm/trixie amd64 is supported.')
    return values['VERSION_CODENAME']


def migration(installed):
    return any(p == 'nginx-common' or p.startswith(('nginx-full', 'nginx-light', 'nginx-core', 'nginx-extras', 'libnginx-mod-')) for p in installed) or ('nginx' in installed and not re.fullmatch(r'\d+\.\d+\.\d+-\d+~(?:bookworm|trixie)', installed['nginx']))


def authorize_migration(needed, allowed, yes):
    if not needed:
        return
    if not allowed:
        raise Error('Debian/unknown nginx detected. Back up configuration; --allow-migrate is required (even with --yes). No packages have been removed.')
    if not yes and input('Migrate nginx vendor? Type MIGRATE to confirm: ') != 'MIGRATE':
        raise Error('Migration not confirmed.')


def validate_plan(text, versions, installed):
    if re.search(r'^Remv\s', text, re.M):
        raise Error('APT proposes removals. Prepare migration manually; no implicit removals are allowed.')
    for package, version in versions.items():
        if package in installed and compare(version, installed[package]) < 0:
            raise Error(f'Downgrade refused: {package}')
    for line in text.splitlines():
        if not line.startswith('Inst '):
            continue
        match = re.match(r'Inst (\S+)(?: \[[^]]+\])? \((\S+)', line)
        if not match:
            raise Error('Unrecognized APT simulation output.')
        package, version = match.groups()
        package = package.removesuffix(':amd64')
        if package in versions:
            if version != versions[package]:
                raise Error('APT selected a different nginx/module version.')
        elif package in installed or package.startswith(('nginx', 'libnginx')):
            raise Error(f'Unrelated installed-package change refused: {package}')


def check_sources(root=Path('/')):
    source_dir = root / 'etc/apt/sources.list.d'
    paths = [root / 'etc/apt/sources.list', *source_dir.glob('*.list'), *source_dir.glob('*.sources')]
    for path in paths:
        if not path.exists():
            continue
        text = path.read_text()
        if path.name == SOURCE_NAME:
            if path.is_symlink() or not text.startswith(MARKER + '\n'):
                raise Error(f'Unmanaged source would be overwritten: {path}')
            continue
        active = '\n'.join(line for line in text.splitlines() if not line.lstrip().startswith('#'))
        if re.search(r'(nginx\.org|mmdb\.netcube\.bg)', active, re.I):
            raise Error(f'Conflicting unmanaged source: {path}. Archive/disable it manually, then retry.')


def managed_write(path, data):
    import time
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.is_symlink():
        raise Error(f'Refusing symlink: {path}')
    if path.exists():
        if path.read_bytes() == data:
            return
        shutil.copy2(path, str(path) + f'.backup-{time.time_ns()}')
    with tempfile.NamedTemporaryFile(dir=path.parent, delete=False) as f:
        temporary = Path(f.name)
        f.write(data)
    temporary.chmod(0o644)
    temporary.replace(path)


@contextlib.contextmanager
def no_service_actions(root=Path('/')):
    # Debian package service helpers honor policy-rc.d. Never alter nginx config.
    path = root / 'usr/sbin/policy-rc.d'
    backup = path.with_name('policy-rc.d.nginx-geoip2-installer-backup')
    if backup.exists() or backup.is_symlink():
        raise Error(f'Stale service-policy backup: {backup}; restore it manually before retrying.')
    existed = path.exists() or path.is_symlink()
    if existed:
        path.rename(backup)
    try:
        path.write_text('#!/bin/sh\nexit 101\n')
        path.chmod(0o755)
        yield
    finally:
        path.unlink(missing_ok=True)
        if existed:
            backup.rename(path)


STRICT = ['-o', 'APT::Get::AllowUnauthenticated=false',
          '-o', 'Acquire::AllowInsecureRepositories=false',
          '-o', 'Acquire::AllowDowngradeToInsecureRepositories=false',
          '-o', 'Acquire::https::Verify-Peer=true',
          '-o', 'Acquire::https::Verify-Host=true',
          '-o', 'APT::Update::Error-Mode=any']


def fetch(url):
    with urllib.request.urlopen(url, timeout=60) as response:
        if not response.url.startswith('https://'):
            raise Error('HTTPS downgrade refused.')
        data = response.read(8 * 1024 * 1024 + 1)
        if len(data) > 8 * 1024 * 1024:
            raise Error('Unexpectedly large key/release document.')
        return data


def verified_key(work, name, url, fingerprint):
    home = work / (name + '-gnupg')
    home.mkdir(mode=0o700)
    raw = work / (name + '.download')
    raw.write_bytes(fetch(url))
    gpg = ['gpg', '--batch', '--no-options', '--homedir', str(home)]
    shown = run(*gpg, '--with-colons', '--show-keys', str(raw))
    if fingerprint not in primary_fingerprints(shown):
        raise Error(f'{name}: pinned primary key missing from downloaded bundle.')
    run(*gpg, '--import', str(raw))
    dest = work / (name + '.gpg')
    run(*gpg, '--output', str(dest), '--export', fingerprint)
    validate_fingerprints(run(*gpg, '--with-colons', '--show-keys', str(dest)), fingerprint)
    dest.chmod(0o644)
    return dest


def source_text(suite, nginx_key, our_key):
    return (MARKER + '\n' +
            f'deb [arch=amd64 signed-by={nginx_key}] {NGINX} {suite} nginx\n' +
            f'deb [arch=amd64 signed-by={our_key}] {OURS} {suite} main\n')


def prepare_pair(work, suite, upstream):
    keys = {
        'nginx-geoip2-nginx.gpg': verified_key(work, 'nginx', 'https://nginx.org/keys/nginx_signing.key', NGINX_FPR),
        'nginx-geoip2-ip-signal.gpg': verified_key(work, 'ip-signal', 'https://mmdb.netcube.bg/apt/ip-signal-archive-keyring.gpg', OUR_FPR),
    }
    sources = work / 'pair.list'
    sources.write_text(source_text(suite, *keys.values()))
    for name in ('lists/partial', 'archives/partial', 'empty'):
        (work / name).mkdir(parents=True)
    # APT_CONFIG is read before system apt.conf.d: disable host hooks for this
    # isolated, temporary metadata/download operation, including in dry-run.
    config = work / 'apt.conf'
    config.write_text('Dir::Etc::parts "-";\nDir::Etc::main "-";\n')
    env = {**os.environ, 'LC_ALL': 'C', 'APT_CONFIG': str(config)}
    opts = STRICT + ['-o', f'Dir::Etc::sourcelist={sources}',
                    '-o', f'Dir::Etc::sourceparts={work / "empty"}',
                    '-o', f'Dir::State::lists={work / "lists"}',
                    '-o', f'Dir::Cache::archives={work / "archives"}',
                    '-o', 'Dir::Cache::pkgcache=', '-o', 'Dir::Cache::srcpkgcache=',
                    '-o', 'APT::Architecture=amd64', '-o', 'APT::Architectures::=amd64']
    run('apt-get', *opts, 'update', env=env)
    # Select full Debian revision using dpkg ordering, never lexical sorting.
    nginx_policy = run('apt-cache', *opts, 'policy', 'nginx', env=env)
    candidates = re.findall(r'^\s+(?:\*\*\*\s+)?(' + re.escape(upstream) + r'-[^\s]+~' + suite + r')\s+-?\d+\s*$', nginx_policy, re.M)
    if not candidates:
        raise Error(f'Official nginx {upstream} for {suite} is not available.')
    nginx_version = max(candidates, key=functools.cmp_to_key(compare))
    versions = {'nginx': nginx_version, 'nginx-module-geoip2': f'{upstream}+gitcc3f1cb1534e-1~{suite}'}
    metadata = {}
    for package, version in versions.items():
        policy = nginx_policy if package == 'nginx' else run('apt-cache', *opts, 'policy', package, env=env)
        verify_origin(policy, version, NGINX if package == 'nginx' else OURS, suite)
        text = run('apt-cache', *opts, 'show', '--no-all-versions', f'{package}={version}', env=env)
        metadata[package] = package_record(text, package, version) if package == 'nginx' else module_record(text, version, nginx_version)
    # Do not download/install nginx unless BOTH records and dependency verified.
    downloads = work / 'downloads'
    downloads.mkdir()
    run('apt-get', *opts, 'download', *[f'{p}={v}' for p, v in versions.items()], cwd=downloads, env=env)
    files = list(downloads.glob('*.deb'))
    if len(files) != 2:
        raise Error('Expected exactly two downloaded packages.')
    seen = set()
    for path in files:
        package = run('dpkg-deb', '-f', str(path), 'Package').strip()
        version = run('dpkg-deb', '-f', str(path), 'Version').strip()
        arch = run('dpkg-deb', '-f', str(path), 'Architecture').strip()
        if package not in versions or package in seen or version != versions[package] or arch != 'amd64' or hashlib.sha256(path.read_bytes()).hexdigest() != metadata[package]['SHA256']:
            raise Error('Downloaded package does not match authenticated APT metadata.')
        seen.add(package)
    return versions, files, keys


def installed_packages():
    text = run('dpkg-query', '-W', '-f=${Package}\t${Version}\t${db:Status-Status}\n')
    return {p: v for line in text.splitlines() if len(parts := line.split('\t')) == 3
            for p, v, status in [parts] if status == 'installed'}


def configure_repositories(suite, keys, root=Path('/')):
    check_sources(root)
    target = root / 'usr/share/keyrings'
    for name, source in keys.items():
        managed_write(target / name, source.read_bytes())
    managed_write(root / 'etc/apt/sources.list.d' / SOURCE_NAME,
                  source_text(suite, '/usr/share/keyrings/nginx-geoip2-nginx.gpg',
                              '/usr/share/keyrings/nginx-geoip2-ip-signal.gpg').encode())


def verify_system_pair(versions, files, suite):
    # APT can prefer a repository copy over a local .deb of the same version.
    # Verify global origins and hashes again after persistent apt update.
    for path in files:
        package = run('dpkg-deb', '-f', str(path), 'Package').strip()
        version = versions[package]
        verify_origin(run('apt-cache', 'policy', package), version,
                      NGINX if package == 'nginx' else OURS, suite)
        record = package_record(run('apt-cache', 'show', '--no-all-versions',
                                    f'{package}={version}'), package, version)
        if record['SHA256'] != hashlib.sha256(path.read_bytes()).hexdigest():
            raise Error('Repository package changed after verification; refusing install.')


def execute(args, root=Path('/')):
    suite = platform((root / 'etc/os-release').read_text(), run('dpkg', '--print-architecture').strip())
    installed = installed_packages()
    authorize_migration(migration(installed), args.allow_migrate, args.yes)
    check_sources(root)
    upstream = release(fetch('https://nginx.org/en/download.html').decode(), args.legacy)
    with tempfile.TemporaryDirectory(prefix='nginx-geoip2-') as temporary:
        work = Path(temporary)
        work.chmod(0o755)
        versions, files, keys = prepare_pair(work, suite, upstream)
        # Transaction-local freeze of every already installed unrelated package.
        # No persistent/global APT pin and no inadvertent dependency upgrades.
        (work / 'empty-preferences').mkdir()
        preferences = work / 'freeze.pref'
        preferences.write_text(''.join(f'Package: {p}\nPin: version {v}\nPin-Priority: 1001\n\n'
                                       for p, v in installed.items() if p not in versions))
        opts = STRICT + ['-o', f'Dir::Etc::preferences={preferences}',
                         '-o', f'Dir::Etc::preferencesparts={work / "empty-preferences"}',
                         '-o', 'Dpkg::Options::=--force-confold',
                         '--no-remove', '--no-install-recommends']
        command = ['apt-get', *opts, 'install', *map(str, files)]
        plan = run(command[0], '--simulate', *command[1:])
        validate_plan(plan, versions, installed)
        print(json.dumps({'suite': suite, 'channel': 'legacy' if args.legacy else 'stable', 'packages': versions}, indent=2))
        print(plan)
        if args.dry_run:
            print('DRY RUN: verified pair; no persistent repository/package/service changes. Temporary downloads are removed.')
            return
        if not args.yes and input('Configure signed repositories and install this exact pair? Type INSTALL: ') != 'INSTALL':
            raise Error('Installation not confirmed.')
        configure_repositories(suite, keys, root)
        run('apt-get', *STRICT, 'update')
        verify_system_pair(versions, files, suite)
        # Refresh plan after index update and reject concurrent dpkg changes.
        if installed_packages() != installed:
            raise Error('Installed packages changed during planning; retry.')
        validate_plan(run(command[0], '--simulate', *command[1:]), versions, installed)
        with no_service_actions(root):
            print(run(command[0], '--yes', *command[1:], env={**os.environ, 'LC_ALL': 'C', 'DEBIAN_FRONTEND': 'noninteractive', 'NEEDRESTART_MODE': 'l'}))
        final = installed_packages()
        if any(final.get(p) != v for p, v in versions.items()):
            raise Error('Post-install exact-version verification failed.')
        print(run('/usr/sbin/nginx', '-t'))
        print('Exact pair installed. No service start/reload/restart requested; existing workers may still run the old binary. Review module load_module configuration manually, run nginx -t, then schedule your own restart. No module-load configuration was changed.')


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--legacy', action='store_true', help='Latest legacy release instead of latest stable on nginx.org')
    parser.add_argument('--allow-migrate', action='store_true', help='Acknowledge Debian/unknown nginx vendor migration; never permits removals/downgrades')
    parser.add_argument('--yes', action='store_true', help='Confirm prompts (migration still requires --allow-migrate)')
    parser.add_argument('--dry-run', action='store_true', help='Verify signed metadata, download exact pair and simulate; no persistent changes')
    args = parser.parse_args(argv)
    if os.geteuid() != 0 and not args.dry_run:
        raise Error('Run as root (or use --dry-run/--help).')
    execute(args)


if __name__ == '__main__':
    try:
        main()
    except (Error, OSError, ValueError, EOFError, KeyboardInterrupt) as exc:
        print(f'nginx-geoip2-install: {exc}', file=sys.stderr)
        sys.exit(1)
