"""Prepare only exported Gentoo payloads and configure the dedicated guest network."""
import argparse
import hashlib
import ipaddress
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys

LABEL = 'HYDRA_OS'
BRIDGE = 'hydrabr0'
SUBNET = ipaddress.ip_network('10.203.0.0/24')
CONNECTION_UUID = 'c84844bd-2955-4a61-9463-5e540140f9a2'


class ImageError(RuntimeError):
    pass


def root_path(root, relative, follow_leaf=True):
    root = Path(root).resolve()
    if root == Path('/') or not root.is_dir():
        raise ImageError('Expected an isolated exported rootfs directory')
    relative = Path(relative)
    if relative.is_absolute() or '..' in relative.parts:
        raise ImageError('Payload path must be relative and remain inside the export')
    path = root / relative
    checked = path if follow_leaf else path.parent
    if not checked.resolve().is_relative_to(root):
        raise ImageError(f'Payload symlink escapes the export: {relative}')
    return path


def remove(root, relative):
    path = root_path(root, relative, follow_leaf=False)
    if path.is_symlink() or path.is_file():
        path.unlink()
    elif path.is_dir():
        shutil.rmtree(path)


def write(root, relative, content, mode=0o644):
    path = root_path(root, relative)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content)
    path.chmod(mode)
    return path


def sanitize_rootfs(root):
    """Erase builder identity/secrets in the exported copy, retaining Gentoo updates."""
    for relative in ('etc/portage/gnupg', 'etc/portage/gnupg.hydra-untrusted',
                     'var/lib/hydra-build', 'root', 'home', 'var/cache', 'var/tmp',
                     'tmp', 'run', 'var/log', 'var/lib/systemd/random-seed',
                     'var/lib/NetworkManager', 'etc/NetworkManager/system-connections',
                     'hydra-source', 'hydra-cache', 'hydra-output'):
        remove(root, relative)
    for name in root_path(root, 'etc/ssh').glob('ssh_host_*'):
        remove(root, str(name.relative_to(root)))
    for relative, mode in [('root', 0o700), ('home', 0o755), ('tmp', 0o1777),
                           ('var/tmp', 0o1777), ('run', 0o755), ('var/log', 0o755),
                           ('var/cache', 0o755), ('dev', 0o755), ('proc', 0o755), ('sys', 0o755)]:
        path = root_path(root, relative)
        path.mkdir(parents=True, exist_ok=True)
        path.chmod(mode)
    remove(root, 'etc/machine-id')
    write(root, 'etc/machine-id', '')
    remove(root, 'var/lib/dbus/machine-id')
    path = root_path(root, 'var/lib/dbus/machine-id', follow_leaf=False)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.symlink_to('/etc/machine-id')
    write(root, 'etc/hostname', 'hydra-os\n')
    write(root, 'etc/hosts', '127.0.0.1 localhost\n::1 localhost\n127.0.1.1 hydra-os\n')
    remove(root, 'etc/resolv.conf')
    root_path(root, 'etc/resolv.conf', follow_leaf=False).symlink_to('/run/NetworkManager/resolv.conf')


def copy_overlay(root, overlay):
    for source in sorted(Path(overlay).rglob('*')):
        destination = root_path(root, source.relative_to(overlay))
        if source.is_symlink():
            raise ImageError('Source overlay must contain ordinary directories/files')
        if source.is_dir():
            destination.mkdir(parents=True, exist_ok=True)
        elif source.is_file():
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(source, destination)
            destination.chmod(source.stat().st_mode & 0o777)


def kernel_version(root):
    modules = root_path(root, 'lib/modules')
    versions = sorted(path.name for path in modules.iterdir() if path.is_dir()) if modules.is_dir() else []
    if len(versions) != 1 or not re.fullmatch(r'[A-Za-z0-9_.+-]+', versions[0]):
        raise ImageError('Exactly one installed payload kernel is required')
    if not root_path(root, 'boot/vmlinuz-' + versions[0]).is_file():
        raise ImageError('Installed modules have no matching /boot/vmlinuz-VERSION')
    return versions[0]


def live_kernel_args():
    return [f'root=live:CDLABEL={LABEL}', 'rd.live.image', 'rd.live.overlay.overlayfs=1',
            'ro', 'console=tty0', 'console=ttyS0,115200n8']


def source_files(source):
    source = Path(source)
    result = {}
    for directory in ('image', 'hydra', 'bin', 'session', 'patches'):
        for path in sorted((source / directory).rglob('*')):
            if path.is_file() and '__pycache__' not in path.parts and path.suffix != '.pyc':
                result[str(path.relative_to(source))] = digest(path)
    return result


def source_identity(source):
    return hashlib.sha256(json.dumps(source_files(source), sort_keys=True).encode()).hexdigest()


def payload_manifest(root, source):
    packages = []
    database = root_path(root, 'var/db/pkg')
    for use in sorted(database.glob('*/*/USE')):
        package = use.parent
        built = package / 'BUILD_TIME'
        packages.append({'package': str(package.relative_to(database)), 'use': use.read_text().split(),
                         'build_time': built.read_text().strip() if built.is_file() else None})
    commit = os.environ.get('HYDRA_SOURCE_COMMIT')
    if not commit:
        git = subprocess.run(['git', '-C', str(source), 'rev-parse', 'HEAD'], text=True, capture_output=True)
        commit = git.stdout.strip() if git.returncode == 0 else None
    return {'format': 1, 'label': LABEL, 'kernel': kernel_version(root),
            'live_kernel_args': live_kernel_args(), 'source_commit': commit,
            'source_files': source_files(source), 'sources': json.loads((Path(source) / 'image/sources.json').read_text()),
            'packages': packages}


def artifacts_valid(directory, state):
    artifacts = state.get('artifacts')
    if not isinstance(artifacts, dict) or not artifacts:
        return False
    for name, expected in artifacts.items():
        if not isinstance(name, str) or Path(name).name != name or not isinstance(expected, str):
            return False
        path = Path(directory) / name
        try:
            if path.is_symlink() or not path.is_file() or digest(path) != expected:
                return False
        except OSError:
            return False
    return True


def configure_rootfs(root, source):
    root, source = Path(root), Path(source)
    # Refuse incomplete integrations before changing the exported copy.
    for name in ('bin/hydra', 'bin/hydra-install', 'hydra/session.py', 'hydra/installer.py'):
        if not (source / name).is_file():
            raise ImageError(f'Required committed runtime/installer asset is missing: {name}')
    version = kernel_version(root)
    sanitize_rootfs(root)
    copy_overlay(root, source / 'image/rootfs')
    write(root, 'etc/hydra/live-mode/dracut.conf', (source / 'image/dracut.conf').read_text())
    root_path(root, 'etc/hydra/live-mode/dracut.conf.d').mkdir()
    package = root_path(root, 'usr/lib/hydra/hydra')
    package.mkdir(parents=True, exist_ok=True)
    for module in sorted((source / 'hydra').glob('*.py')):
        shutil.copyfile(module, root_path(root, 'usr/lib/hydra/hydra/' + module.name))
    for name, module in (('hydra', 'cli'), ('hydra-install', 'installer')):
        write(root, 'usr/bin/' + name,
              '#!/usr/bin/python3\nimport sys\nsys.path.insert(0, "/usr/lib/hydra")\n'
              f'from hydra.{module} import main\nraise SystemExit(main())\n', 0o755)
    write(root, 'usr/bin/hydra-session', '#!/bin/sh\nexec /usr/bin/hydra session "$@"\n', 0o755)
    write(root, 'usr/bin/hydra-network', '#!/usr/bin/python3\nimport sys\nsys.path.insert(0, "/usr/lib/hydra")\n'
          'from hydra.image import network_main\nraise SystemExit(network_main())\n', 0o755)
    write(root, 'usr/share/hydra/build-manifest.json', json.dumps(payload_manifest(root, source), indent=2) + '\n')
    make = root_path(root, 'etc/portage/make.conf')
    if make.is_file():
        settings = make.read_text()
        settings = re.sub(r'(?<=[\s"\'])-(?:usersandbox|sandbox|network-sandbox|ipc-sandbox|pid-sandbox)(?=[\s"\']|$)', '', settings)
        make.write_text(settings)
    for path in root_path(root, 'etc/sudoers.d').iterdir():
        path.chmod(0o440)
    return version


def check_network_conflicts(routes, addresses):
    try:
        if not isinstance(routes, list) or not isinstance(addresses, list):
            raise ValueError('Expected JSON arrays')
        for route in routes:
            dst = route.get('dst')
            if dst in ('default', None):
                continue
            network = ipaddress.ip_network(dst, strict=False)
            if network.version != 4:
                raise ValueError('Unexpected non-IPv4 route')
            if network.overlaps(SUBNET) and not (route.get('dev') == BRIDGE and network.subnet_of(SUBNET)):
                raise ImageError(f'Guest subnet {SUBNET} conflicts with route {dst} on {route.get("dev")}')
        for interface in addresses:
            for addr in interface['addr_info']:
                if addr.get('family') != 'inet':
                    continue
                network = ipaddress.ip_network(f'{addr["local"]}/{addr["prefixlen"]}', strict=False)
                owned = interface['ifname'] == BRIDGE and addr['local'] == '10.203.0.1' and addr['prefixlen'] == 24
                if network.overlaps(SUBNET) and not owned:
                    raise ImageError(f'Guest subnet {SUBNET} conflicts with address on {interface["ifname"]}')
    except (ValueError, KeyError, TypeError, AttributeError) as error:
        raise ImageError('Cannot safely interpret IPv4 routes/addresses') from error


def configure_network(run=subprocess.run):
    def command(argv, **kwargs):
        result = run(argv, text=True, capture_output=True, **kwargs)
        if result.returncode:
            raise ImageError(f'Network command failed: {argv[0]}: {result.stderr.strip()}')
        return result.stdout
    routes = json.loads(command(['ip', '-json', '-4', 'route', 'show', 'table', 'all']))
    addresses = json.loads(command(['ip', '-json', '-4', 'address', 'show']))
    check_network_conflicts(routes, addresses)
    probe = run(['nmcli', '--get-values', 'connection.uuid,connection.type,connection.interface-name,ipv4.method,ipv4.addresses',
                 'connection', 'show', 'uuid', CONNECTION_UUID], text=True, capture_output=True)
    exists = probe.returncode == 0
    if exists and probe.stdout.strip().splitlines() != [CONNECTION_UUID, 'bridge', BRIDGE, 'manual', '10.203.0.1/24']:
        raise ImageError('Hydra bridge connection exists with unexpected configuration')
    if not exists and probe.returncode != 10:
        raise ImageError('Cannot query NetworkManager connections')
    links = json.loads(command(['ip', '-json', '-details', 'link', 'show']))
    for link in links:
        if link.get('ifname') == BRIDGE and (not exists or link.get('linkinfo', {}).get('info_kind') != 'bridge'):
            raise ImageError('Refusing to replace an existing interface named hydrabr0')
    if not exists:
        command(['nmcli', 'connection', 'add', 'type', 'bridge', 'con-name', 'hydra-guest',
                 'ifname', BRIDGE, 'connection.uuid', CONNECTION_UUID,
                 'connection.autoconnect', 'no', 'bridge.stp', 'no',
                 'ipv4.method', 'manual', 'ipv4.addresses', '10.203.0.1/24',
                 'ipv4.never-default', 'yes', 'ipv6.method', 'disabled'])
    # Activation is always gated by the fresh route/address check, including restarts.
    command(['nmcli', 'connection', 'up', 'uuid', CONNECTION_UUID])
    command(['sysctl', '-w', 'net.ipv4.ip_forward=1'])
    table = run(['nft', 'list', 'table', 'ip', 'hydra_nat'], text=True, capture_output=True)
    rules = ('delete table ip hydra_nat\n' if table.returncode == 0 else '') + '''table ip hydra_nat {
  chain postrouting {
    type nat hook postrouting priority srcnat; policy accept;
    ip saddr 10.203.0.0/24 oifname != "hydrabr0" masquerade
  }
}
'''
    command(['nft', '-f', '-'], input=rules)


def network_main():
    try:
        configure_network()
    except (ImageError, OSError, ValueError) as error:
        print(f'hydra-network: {error}', file=sys.stderr)
        return 1
    return 0


def digest(path):
    value = hashlib.sha256()
    with Path(path).open('rb') as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            value.update(block)
    return value.hexdigest()


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('step', choices=('configure', 'kernel', 'network'))
    parser.add_argument('root', nargs='?', type=Path)
    parser.add_argument('--source', type=Path)
    options = parser.parse_args(argv)
    try:
        if options.step == 'network':
            return network_main()
        if options.root is None:
            raise ImageError('An isolated rootfs directory is required')
        if options.step == 'kernel':
            print(kernel_version(options.root))
        else:
            if options.source is None:
                raise ImageError('Source checkout is required')
            configure_rootfs(options.root, options.source)
    except (ImageError, OSError, ValueError) as error:
        print(f'hydra-image: {error}', file=sys.stderr)
        return 1
    return 0


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