#!/usr/bin/env python3
"""Boot/install/reboot acceptance on a fresh, isolated UEFI/KVM VM only.

No existing disk option, host sudo, SSH, host TTY, or GUI is used. Artifacts and
logs remain in a new private scratch directory even after failure.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import json
import os
from pathlib import Path
import re
import secrets
import select
import shlex
import shutil
import socket
import stat
import subprocess
import sys
import tempfile
import time

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from hydra.qmp import QMPClient, QMPError


class SmokeError(RuntimeError):
    pass


def reserve_disk(path):
    try:
        fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
        os.close(fd)
    except OSError as exc:
        raise SmokeError('Test disk must be a newly created regular file') from exc


def owned_file(path, scratch):
    path, scratch = Path(path), Path(scratch)
    try:
        info = path.lstat()
    except OSError as exc:
        raise SmokeError('Missing disposable drive') from exc
    if (not stat.S_ISREG(info.st_mode) or info.st_nlink != 1
            or path.parent.resolve() != scratch.resolve()
            or info.st_uid != os.getuid()):
        raise SmokeError('Drive must be an owned, unlinked regular file directly in scratch')


def qemu_path(path):
    # QEMU keyval arguments use commas as separators. Refuse ambiguous paths.
    if any(c in str(path) for c in ',\n\r\x00'):
        raise SmokeError('QEMU file paths cannot contain comma or control characters')
    return str(path)


def qemu_argv(qemu, scratch, disk, code, variables, iso, memory):
    scratch = Path(scratch)
    owned_file(disk, scratch)
    owned_file(variables, scratch)
    argv = [qemu, '-machine', 'q35,accel=kvm', '-cpu', 'host', '-smp', '2',
            '-m', str(memory), '-display', 'none', '-monitor', 'none',
            '-no-reboot', '-nic', 'user,model=virtio-net-pci',
            '-drive', f'if=pflash,format=raw,readonly=on,file={qemu_path(code)}',
            '-drive', f'if=pflash,format=raw,file={qemu_path(variables)}',
            '-drive', f'if=virtio,format=qcow2,file={qemu_path(disk)}',
            '-chardev', f'socket,id=console,path={qemu_path(scratch / "serial.sock")},server=on,wait=off',
            '-serial', 'chardev:console', '-qmp', f'unix:{qemu_path(scratch / "qmp.sock")},server=on,wait=off']
    if iso is not None:
        if not Path(iso).is_file():
            raise SmokeError('Trusted ISO must exist as a regular file')
        argv += ['-drive', f'file={qemu_path(iso)},media=cdrom,format=raw,readonly=on',
                 '-boot', 'order=d,strict=on']
    else:
        argv += ['-boot', 'order=c,strict=on']
    return argv


def frame_command(command, token):
    if not re.fullmatch(r'[a-zA-Z0-9]+', token):
        raise SmokeError('Invalid framing token')
    encoded = base64.b64encode(command.encode()).decode()
    # Literal control bytes appear only in output, never in the echoed command.
    return (f"printf '\\036{token}:BEGIN\\037'; "
            f"sh -c \"$(printf %s {shlex.quote(encoded)} | base64 -d)\"; "
            f"rc=$?; printf '\\036{token}:END:%s\\037' \"$rc\"")


def parse_frame(data, token):
    begin = b'\x1e' + token.encode() + b':BEGIN\x1f'
    start = data.find(begin)
    if start < 0:
        return None
    end = re.search(rb'\x1e' + token.encode() + rb':END:(\d+)\x1f', data[start + len(begin):])
    if end is None:
        return None
    output = data[start + len(begin):start + len(begin) + end.start()]
    return int(end.group(1)), output


class Console:
    def __init__(self, connection, logfile, passwords):
        self.sock = connection
        self.log = open(logfile, 'ab', buffering=0)
        self.passwords = [p.encode() for p in passwords if p]
        self.pending_log = b''
        self.buffer = b''
        self.closed = False

    def record(self, data, final=False):
        self.pending_log += data
        for password in self.passwords:
            self.pending_log = self.pending_log.replace(password, b'[REDACTED]')
        keep = 0
        if not final:
            for password in self.passwords:
                for length in range(1, len(password)):
                    if self.pending_log.endswith(password[:length]):
                        keep = max(keep, length)
        if keep:
            self.log.write(self.pending_log[:-keep])
            self.pending_log = self.pending_log[-keep:]
        else:
            self.log.write(self.pending_log)
            self.pending_log = b''

    def close(self):
        if not self.closed:
            self.record(b'', final=True)
            self.log.close()
            self.sock.close()
            self.closed = True

    def send(self, data):
        self.sock.sendall(data.encode() if isinstance(data, str) else data)

    def read(self, timeout):
        ready, _, _ = select.select([self.sock], [], [], max(0, timeout))
        if not ready:
            return b''
        data = self.sock.recv(65536)
        if not data:
            raise SmokeError('Guest console closed')
        self.record(data)
        self.buffer += data
        if len(self.buffer) > 16 * 1024 * 1024:
            raise SmokeError('Console exceeded bounded capture size')
        return data

    def expect(self, pattern, timeout):
        deadline = time.monotonic() + timeout
        while True:
            match = re.search(pattern, self.buffer)
            if match:
                self.buffer = self.buffer[match.end():]
                return match
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise SmokeError('Timed out waiting for guest console state')
            self.read(min(remaining, 1))

    def command(self, command, timeout=60):
        token = secrets.token_hex(12)
        self.buffer = b''
        self.send(frame_command(command, token) + '\n')
        deadline = time.monotonic() + timeout
        while True:
            parsed = parse_frame(self.buffer, token)
            if parsed is not None:
                status, output = parsed
                self.buffer = b''
                if status:
                    raise SmokeError(f'Guest probe failed with status {status}: {output.decode(errors="replace")[-3000:]}')
                return output.replace(b'\r\n', b'\n')
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise SmokeError('Timed out waiting for framed guest command')
            self.read(min(remaining, 1))


def connect(path, process, timeout=30):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if process.poll() is not None:
            raise SmokeError(f'Test QEMU exited with status {process.returncode}')
        connection = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        try:
            connection.connect(str(path))
            return connection
        except OSError:
            connection.close()
            time.sleep(0.1)
    raise SmokeError('Timed out connecting to private test socket')


class VM:
    def __init__(self, args, scratch, disk, iso, password, phase):
        self.scratch = scratch
        for name in ('serial.sock', 'qmp.sock'):
            (scratch / name).unlink(missing_ok=True)
        self.stderr = open(scratch / f'{phase}-qemu.log', 'ab', buffering=0)
        self.process = subprocess.Popen(qemu_argv(args.qemu, scratch, disk,
            args.ovmf_code, scratch / 'OVMF_VARS.fd', iso, args.memory),
            stdin=subprocess.DEVNULL, stdout=self.stderr, stderr=self.stderr)
        self.console = None
        self.qmp = None
        try:
            self.console = Console(connect(scratch / 'serial.sock', self.process),
                                   scratch / f'{phase}-serial.log', [password])
            ready = connect(scratch / 'qmp.sock', self.process)
            ready.close()
            self.qmp = QMPClient(scratch / 'qmp.sock', timeout=5)
            self.kvm = self.qmp.command('query-kvm')
            if not self.kvm.get('enabled'):
                raise SmokeError('Disposable QEMU is not using KVM')
        except BaseException:
            self.close()
            raise

    def wait_poweroff(self, timeout):
        try:
            rc = self.process.wait(timeout=timeout)
        except subprocess.TimeoutExpired as exc:
            raise SmokeError('Guest failed to power off in time') from exc
        if rc:
            raise SmokeError(f'Test QEMU failed during poweroff ({rc})')

    def close(self):
        # Only the Popen child owned by this harness is signalled.
        if self.qmp:
            if self.process.poll() is None:
                try:
                    self.qmp.command('quit')
                    self.process.wait(timeout=5)
                except (QMPError, subprocess.TimeoutExpired):
                    pass
            self.qmp.close()
        if self.process.poll() is None:
            self.process.terminate()
            try:
                self.process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                self.process.kill()
                self.process.wait(timeout=10)
        if self.console:
            self.console.close()
        self.stderr.close()


LIVE_CHECK = r'''set -eu
[ "$(id -un)" = hydra-live ]
test -e /sys/firmware/efi
test -f /etc/hydra/live-mode/active
findmnt -n -o FSTYPE / | grep -qx overlay
test -f /run/initramfs/live/LiveOS/squashfs.img
sudo -n true
for tool in hydra hydra-install qemu-system-x86_64 qemu-img virt-viewer foot labwc; do command -v "$tool"; done
systemctl is-active NetworkManager hydra-network hydra-dnsmasq
'''

INSTALLED_CHECK = r'''set -eu
[ "$(findmnt -n -o FSTYPE /)" = ext4 ]
[ "$(readlink -f "$(findmnt -n -o SOURCE /)")" = /dev/vda2 ]
! grep -Eq 'root=live:|rd\.live' /proc/cmdline
! findmnt -n -o FSTYPE / | grep -Eq 'overlay|squashfs'
test ! -e /etc/hydra/live-mode
! getent passwd hydra-live
! find /etc/systemd/system /etc/sudoers.d -name 90-hydra-live.conf -o -name 90-hydra-live | grep .
test ! -e /etc/greetd/live.toml
! grep -R -E -- '--autologin|NOPASSWD' /etc/systemd/system/getty@*.service.d /etc/systemd/system/serial-getty@*.service.d /etc/sudoers.d 2>/dev/null
getent shadow root | cut -d: -f2 | grep -Eq '^[!*]'
getent shadow hydratest | cut -d: -f2 | grep -Eq '^\$'
for service in NetworkManager greetd hydra-network hydra-dnsmasq; do systemctl is-active "$service"; done
[ "$(systemctl show hydra-bridge-permissions -p Result --value)" = success ]
[ "$(systemctl show hydra-bridge-permissions -p ExecMainStatus --value)" = 0 ]
ip -4 -o addr show hydrabr0 | grep -F '10.203.0.1/24'
ip -4 route show default | grep '^default '
nft list table ip hydra_nat
nmcli -t -f STATE general | grep -E '^connected'
getent ahostsv4 gentoo.org
for tool in hydra hydra-install qemu-system-x86_64 qemu-img virt-viewer foot labwc; do command -v "$tool"; done
hydra --help
qemu-system-x86_64 --version
qemu-img --version
test -r /etc/greetd/config.toml
find /usr/share -name OVMF_CODE.fd -print -quit | grep .
'''

NETWORK_CHECK = r'''set -eu
[ "$(cat /proc/sys/net/ipv4/ip_forward)" = 1 ]
python3 - <<'NETWORK'
import pathlib, secrets, subprocess, time
tag = secrets.token_hex(3)
lan, guest = 'hl' + tag, 'hv' + tag
uplink, lanpeer = 'hu' + tag, 'hp' + tag
bridgepeer, guestpeer = 'hb' + tag, 'hg' + tag
created = []
links = []
server = None
def run(*args):
    return subprocess.run(args, check=True, capture_output=True, text=True, timeout=20).stdout
try:
    for name in (lan, guest):
        run('ip', 'netns', 'add', name)
        created.append(name)
    run('ip', 'link', 'add', uplink, 'type', 'veth', 'peer', 'name', lanpeer)
    links.append(uplink)
    run('ip', 'link', 'set', lanpeer, 'netns', lan)
    run('ip', 'addr', 'add', '198.18.0.1/30', 'dev', uplink)
    run('ip', 'link', 'set', uplink, 'up')
    run('ip', '-n', lan, 'addr', 'add', '198.18.0.2/30', 'dev', lanpeer)
    run('ip', '-n', lan, 'link', 'set', lanpeer, 'up')
    run('ip', '-n', lan, 'link', 'set', 'lo', 'up')
    run('ip', '-n', lan, 'route', 'add', '10.203.0.0/24', 'via', '198.18.0.1')
    run('ip', 'link', 'add', bridgepeer, 'type', 'veth', 'peer', 'name', guestpeer)
    links.append(bridgepeer)
    run('ip', 'link', 'set', guestpeer, 'netns', guest)
    run('ip', 'link', 'set', bridgepeer, 'master', 'hydrabr0')
    run('ip', 'link', 'set', bridgepeer, 'up')
    run('ip', '-n', guest, 'addr', 'add', '10.203.0.250/24', 'dev', guestpeer)
    run('ip', '-n', guest, 'link', 'set', guestpeer, 'up')
    run('ip', '-n', guest, 'link', 'set', 'lo', 'up')
    run('ip', '-n', guest, 'route', 'add', 'default', 'via', '10.203.0.1')
    port = 43117
    server_code = "import socket; s=socket.socket(); s.settimeout(15); s.bind(('10.203.0.250',43117)); s.listen(1); c,a=s.accept(); c.sendall(b'HYDRA_LAN_TO_VM_OK'); c.close(); s.close()"
    server = subprocess.Popen(['ip', 'netns', 'exec', guest, 'python3', '-c', server_code], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    client_code = "import socket,time; deadline=time.monotonic()+10\nwhile True:\n try:\n  s=socket.create_connection(('10.203.0.250',43117),1); break\n except OSError:\n  if time.monotonic()>deadline: raise\n  time.sleep(.1)\nassert s.recv(100)==b'HYDRA_LAN_TO_VM_OK'; s.close(); print('PASS: routed new LAN-to-VM TCP connection on arbitrary port 43117')"
    print(run('ip', 'netns', 'exec', lan, 'python3', '-c', client_code), end='')
    assert server.wait(timeout=10) == 0
finally:
    if server is not None and server.poll() is None:
        server.terminate()
        try: server.wait(timeout=5)
        except subprocess.TimeoutExpired: server.kill(); server.wait(timeout=5)
    for link in reversed(links):
        subprocess.run(['ip', 'link', 'del', link], check=False, capture_output=True, timeout=10)
    for name in reversed(created):
        subprocess.run(['ip', 'netns', 'del', name], check=False, capture_output=True, timeout=10)
NETWORK
'''

# A nested VM hardware smoke proves KVM initialisation/QMP, not guest OS boot.
NESTED_CHECK = r'''set -eu
if ! test -c /dev/kvm; then printf 'SKIP: guest /dev/kvm unavailable\n'; exit 0; fi
python3 - <<'NESTED'
import json, pathlib, socket, subprocess, tempfile, time
with tempfile.TemporaryDirectory(prefix='hydra-nested-') as tmp:
    root = pathlib.Path(tmp)
    disk = root / 'fresh.qcow2'
    subprocess.run(['qemu-img', 'create', '-f', 'qcow2', str(disk), '64M'], check=True, timeout=30)
    log = open(root / 'qemu.log', 'wb')
    proc = subprocess.Popen(['qemu-system-x86_64', '-machine', 'q35,accel=kvm', '-cpu', 'host', '-m', '128', '-S', '-display', 'none', '-nodefaults', '-nic', 'bridge,br=hydrabr0,model=virtio-net-pci', '-drive', 'if=virtio,format=qcow2,file=' + str(disk), '-qmp', 'unix:' + str(root / 'qmp') + ',server=on,wait=off'], stdin=subprocess.DEVNULL, stdout=log, stderr=log)
    sock = socket.socket(socket.AF_UNIX)
    sock.settimeout(10)
    try:
        deadline = time.monotonic() + 15
        while True:
            try: sock.connect(str(root / 'qmp')); break
            except OSError:
                if proc.poll() is not None or time.monotonic() >= deadline: raise RuntimeError('Nested KVM start failed')
                time.sleep(.1)
        stream = sock.makefile('rwb', buffering=0)
        assert 'QMP' in json.loads(stream.readline())
        for request in ({'execute':'qmp_capabilities'}, {'execute':'query-kvm'}, {'execute':'quit'}):
            stream.write(json.dumps(request).encode() + b'\n')
            while True:
                response = json.loads(stream.readline())
                if 'return' in response or 'error' in response: break
            assert 'error' not in response, response
            if request['execute'] == 'query-kvm':
                assert response['return']['enabled'], response
                print('PASS: normal-user nested KVM VM and bridge helper; fresh qcow2; private QMP')
        assert proc.wait(timeout=10) == 0
    finally:
        sock.close()
        if proc.poll() is None:
            proc.terminate()
            try: proc.wait(timeout=5)
            except subprocess.TimeoutExpired: proc.kill(); proc.wait(timeout=5)
        log.close()
NESTED
'''


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


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--iso', required=True, type=Path)
    parser.add_argument('--ovmf-code', required=True, type=Path)
    parser.add_argument('--ovmf-vars', required=True, type=Path)
    parser.add_argument('--qemu', default='qemu-system-x86_64')
    parser.add_argument('--qemu-img', default='qemu-img')
    parser.add_argument('--scratch-parent', type=Path, default=Path(tempfile.gettempdir()))
    parser.add_argument('--memory', type=int, default=4096)
    parser.add_argument('--boot-timeout', type=int, default=360)
    parser.add_argument('--install-timeout', type=int, default=1800)
    args = parser.parse_args(argv)
    for path in (args.iso, args.ovmf_code, args.ovmf_vars):
        if not path.is_file():
            parser.error(f'Missing input file: {path}')
    if args.memory < 2048 or args.boot_timeout <= 0 or args.install_timeout <= 0:
        parser.error('Memory must be >=2048 MiB and timeouts positive')
    args.iso = args.iso.resolve()
    args.ovmf_code = args.ovmf_code.resolve()
    args.ovmf_vars = args.ovmf_vars.resolve()
    scratch = Path(tempfile.mkdtemp(prefix='hydra-iso-smoke-', dir=args.scratch_parent))
    if len(str(scratch / 'serial.sock').encode()) >= 104:
        parser.error('Scratch path too long for Unix sockets; use /tmp')
    print(f'Acceptance artifacts: {scratch}', flush=True)
    report = {'schema': 1, 'status': 'running', 'scratch': str(scratch),
              'iso': str(args.iso), 'iso_sha256': digest(args.iso),
              'checks': {}, 'physical_hardware': 'pending',
              'nested_scope': 'Normal-user KVM VM launch, bridge helper and QMP; no nested guest OS boot',
              'ovmf_code_sha256': digest(args.ovmf_code),
              'ovmf_vars_input_sha256': digest(args.ovmf_vars)}
    password = secrets.token_urlsafe(24)
    sentinel = secrets.token_hex(32)
    vm = None
    def checkpoint(name, output):
        report['checks'][name] = {'status': 'skipped' if output.startswith(b'SKIP:') else 'passed',
                                  'output': output.decode(errors='replace')}
        (scratch / 'report.json').write_text(json.dumps(report, indent=2) + '\n')
    try:
        disk = scratch / 'fresh-install.qcow2'
        reserve_disk(disk)
        subprocess.run([args.qemu_img, 'create', '-f', 'qcow2', str(disk), '32G'], check=True, timeout=30,
                       stdout=subprocess.DEVNULL)
        shutil.copyfile(args.ovmf_vars, scratch / 'OVMF_VARS.fd')
        os.chmod(scratch / 'OVMF_VARS.fd', 0o600)
        vm = VM(args, scratch, disk, args.iso, password, 'live')
        checkpoint('live_qmp_kvm', json.dumps(vm.kvm).encode())
        console = vm.console
        console.expect(rb'hydra-live[^\r\n]*[#$] ', args.boot_timeout)
        checkpoint('live_boot', console.command(LIVE_CHECK, 90))
        console.buffer = b''
        console.send('sudo -n hydra-install --disk /dev/vda --user hydratest --execute --confirm-erase /dev/vda\n')
        console.expect(rb'Installed account password: ', 120)
        console.send(password + '\n')
        console.expect(rb'Repeat password: ', 30)
        console.send(password + '\n')
        console.expect(rb'Installation completed\. Diagnostic log: [^\r\n]+', args.install_timeout)
        checkpoint('real_installer', b'Installer completed via guest controlling TTY')
        console.command("sudo -n mkdir -p /mnt/hydra-validation && sudo -n mount /dev/vda2 /mnt/hydra-validation && printf %s " +
                        shlex.quote(sentinel) + " | sudo -n tee /mnt/hydra-validation/var/lib/hydra-acceptance-sentinel >/dev/null && sudo -n umount /mnt/hydra-validation", 60)
        console.send('sudo -n poweroff\n')
        vm.wait_poweroff(90)
        vm.close()
        vm = None
        checkpoint('live_poweroff', b'QEMU exited before ISO-free reboot')
        vm = VM(args, scratch, disk, None, password, 'installed')
        checkpoint('installed_qmp_kvm', json.dumps(vm.kvm).encode())
        console = vm.console
        console.expect(rb'login: ', args.boot_timeout)
        console.send('hydratest\n')
        console.expect(rb'Password: ', 30)
        console.send(password + '\n')
        console.expect(rb'hydratest[^\r\n]*[#$] ', 60)
        output = console.command('id -un; if sudo -n true 2>/dev/null; then exit 1; fi', 30)
        if output.strip() != b'hydratest':
            raise SmokeError('Installed account identity mismatch')
        checkpoint('installed_password_login', output)
        console.buffer = b''
        prompt_encoded = base64.b64encode(b'HYDRA_SUDO_PASSWORD: ').decode()
        console.send(f'sudo -p "$(printf %s {prompt_encoded} | base64 -d)" -i\n')
        console.expect(rb'HYDRA_SUDO_PASSWORD: ', 30)
        console.send(password + '\n')
        console.expect(rb'root[^\r\n]*[#$] ', 30)
        observed = console.command('cat /var/lib/hydra-acceptance-sentinel', 30)
        if observed != sentinel.encode():
            raise SmokeError('Persistence sentinel did not survive installed disk boot')
        checkpoint('disk_persistence', b'Sentinel content verified from command output frame')
        checkpoint('installed_contract', console.command(INSTALLED_CHECK, 120))
        checkpoint('lan_to_vm_tcp', console.command(NETWORK_CHECK, 120))
        checkpoint('nested_kvm', console.command('runuser -u hydratest -- sh -c ' + shlex.quote(NESTED_CHECK), 90))
        console.send('poweroff\n')
        vm.wait_poweroff(90)
        checkpoint('installed_poweroff', b'Normal installed guest shutdown')
        report['status'] = 'passed'
        return 0
    except (SmokeError, QMPError, OSError, subprocess.SubprocessError) as exc:
        report['status'] = 'failed'
        report['error'] = str(exc).replace(password, '[REDACTED]')
        print(f'Acceptance failed: {report["error"]}', file=sys.stderr)
        return 1
    finally:
        if vm:
            vm.close()
        report['completed_utc'] = time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())
        (scratch / 'report.json').write_text(json.dumps(report, indent=2) + '\n')
        print(f'Report: {scratch / "report.json"}', flush=True)


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