"""Regression tests for disposable-VM isolation and trustworthy serial evidence."""
import importlib.util
from pathlib import Path
import socket
import subprocess
import tempfile
import threading
import unittest

SCRIPT = Path(__file__).resolve().parents[1] / 'dev/iso-smoke.py'


class SmokeTests(unittest.TestCase):
    def setUp(self):
        self.assertTrue(SCRIPT.is_file(), 'disposable ISO acceptance harness is missing')
        spec = importlib.util.spec_from_file_location('iso_smoke', SCRIPT)
        self.mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(self.mod)
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.root = Path(self.temp.name)
        self.code = self.root / 'code.fd'
        self.variables = self.root / 'vars.fd'
        self.code.write_bytes(b'firmware')
        self.variables.write_bytes(b'fresh vars')

    def test_existing_disk_or_symlink_cannot_be_a_test_drive(self):
        disk = self.root / 'disk.qcow2'
        disk.write_bytes(b'original')
        with self.assertRaises(self.mod.SmokeError):
            self.mod.reserve_disk(disk)
        link = self.root / 'link.qcow2'
        link.symlink_to(disk)
        with self.assertRaises(self.mod.SmokeError):
            self.mod.reserve_disk(link)
        self.assertEqual(disk.read_bytes(), b'original')

    def test_qemu_only_accepts_owned_regular_disk_and_readonly_iso(self):
        disk = self.root / 'disk.qcow2'
        self.mod.reserve_disk(disk)
        iso = self.root / 'trusted.iso'
        iso.write_bytes(b'ISO')
        argv = self.mod.qemu_argv('qemu', self.root, disk, self.code,
                                  self.variables, iso, 2048)
        drives = [argv[i + 1] for i, value in enumerate(argv[:-1]) if value == '-drive']
        self.assertTrue(any('format=qcow2' in item and str(disk) in item for item in drives))
        self.assertTrue(any('media=cdrom' in item and 'readonly=on' in item for item in drives))
        self.assertIn('accel=kvm', ' '.join(argv))
        disk.unlink()
        disk.symlink_to(iso)
        with self.assertRaises(self.mod.SmokeError):
            self.mod.qemu_argv('qemu', self.root, disk, self.code,
                               self.variables, iso, 2048)

    def test_installed_reboot_has_no_optical_medium(self):
        disk = self.root / 'disk.qcow2'
        self.mod.reserve_disk(disk)
        argv = self.mod.qemu_argv('qemu', self.root, disk, self.code,
                                  self.variables, None, 2048)
        self.assertNotIn('media=cdrom', ' '.join(argv))
        self.assertNotIn('-cdrom', argv)

    def test_hardlinked_disk_is_refused(self):
        disk = self.root / 'disk.qcow2'
        self.mod.reserve_disk(disk)
        (self.root / 'alias').hardlink_to(disk)
        with self.assertRaises(self.mod.SmokeError):
            self.mod.qemu_argv('qemu', self.root, disk, self.code,
                               self.variables, None, 2048)

    def test_result_frame_ignores_command_echo_and_preserves_failure(self):
        command = self.mod.frame_command('printf sentinel; false', 'testtoken')
        result = subprocess.run(['/bin/sh', '-c', command], capture_output=True)
        echo = command.encode() + b'\r\n'
        self.assertIsNone(self.mod.parse_frame(echo, 'testtoken'))
        self.assertEqual(self.mod.parse_frame(echo + result.stdout, 'testtoken'),
                         (1, b'sentinel'))

    def test_console_redacts_password_even_when_split_across_reads(self):
        one, two = socket.socketpair()
        self.addCleanup(one.close)
        self.addCleanup(two.close)
        log = self.root / 'serial.log'
        console = self.mod.Console(one, log, ['split-secret'])
        self.addCleanup(console.close)
        two.sendall(b'hello split-')
        console.read(0.1)
        two.sendall(b'secret goodbye')
        console.read(0.1)
        console.close()
        self.assertNotIn(b'split-secret', log.read_bytes())
        self.assertIn(b'[REDACTED]', log.read_bytes())

    def test_console_timeout_is_bounded_and_partial_frame_not_success(self):
        one, two = socket.socketpair()
        self.addCleanup(one.close)
        self.addCleanup(two.close)
        console = self.mod.Console(one, self.root / 'serial.log', [])
        self.addCleanup(console.close)
        two.sendall(b'\\036token:END:0\\037')
        with self.assertRaises(self.mod.SmokeError):
            console.command('true', timeout=0.05)

    def test_frame_real_shell_roundtrip_does_not_accept_sentinel_in_echo(self):
        one, two = socket.socketpair()
        self.addCleanup(one.close)
        self.addCleanup(two.close)
        console = self.mod.Console(one, self.root / 'serial.log', [])
        self.addCleanup(console.close)
        def guest():
            data = b''
            while not data.endswith(b'\n'):
                data += two.recv(65536)
            two.sendall(data)
            result = subprocess.run(['/bin/sh', '-c', data.decode()], capture_output=True)
            two.sendall(result.stdout)
        thread = threading.Thread(target=guest)
        thread.start()
        self.assertEqual(console.command("printf 'persisted-data'", timeout=2), b'persisted-data')
        thread.join(timeout=2)

    def test_firmware_vars_must_be_private_copy_in_scratch(self):
        disk = self.root / 'disk.qcow2'
        self.mod.reserve_disk(disk)
        self.variables.unlink()
        self.variables.symlink_to(self.code)
        with self.assertRaises(self.mod.SmokeError):
            self.mod.qemu_argv('qemu', self.root, disk, self.code,
                               self.variables, None, 2048)

    def test_drive_outside_scratch_and_ambiguous_qemu_paths_are_refused(self):
        disk = self.root / 'disk.qcow2'
        self.mod.reserve_disk(disk)
        other = self.root / 'other'
        other.mkdir()
        with self.assertRaises(self.mod.SmokeError):
            self.mod.qemu_argv('qemu', other, disk, self.code,
                               self.variables, None, 2048)
        comma = self.root / 'disk,cache=unsafe.qcow2'
        self.mod.reserve_disk(comma)
        with self.assertRaises(self.mod.SmokeError):
            self.mod.qemu_argv('qemu', self.root, comma, self.code,
                               self.variables, None, 2048)

    def probe_fixture(self, case, live=False):
        # Run the production shell against ordinary scratch files. Commands that
        # query guest state are boundary fixtures; grep/find/test are real tools.
        import os
        root = self.root / 'guest'
        for directory in ('etc/hydra', 'etc/greetd', 'etc/sudoers.d',
                          'etc/systemd/system/getty@tty2.service.d',
                          'etc/systemd/system/serial-getty@ttyS0.service.d',
                          'usr/share/firmware', 'proc', 'sys/firmware/efi',
                          'run/initramfs/live/LiveOS'):
            (root / directory).mkdir(parents=True, exist_ok=True)
        (root / 'etc/greetd/config.toml').write_text('normal authentication')
        (root / 'usr/share/firmware/OVMF_CODE.fd').write_text('firmware')
        (root / 'proc/cmdline').write_text('root=UUID=disk ro' if case != 'live-cmdline' else 'root=live:CDLABEL=HYDRA_OS rd.live.image')
        if live or case == 'live-marker':
            (root / 'etc/hydra/live-mode').mkdir(exist_ok=True)
            (root / 'etc/hydra/live-mode/active').touch()
        if live:
            (root / 'run/initramfs/live/LiveOS/squashfs.img').touch()
        files = {
            'live-dropin': ('etc/systemd/system/getty@tty2.service.d/90-hydra-live.conf', 'live'),
            'live-sudoers': ('etc/sudoers.d/90-hydra-live', 'live'),
            'autologin': ('etc/systemd/system/serial-getty@ttyS0.service.d/other.conf', 'ExecStart=agetty --autologin hydratest'),
            'nopasswd': ('etc/sudoers.d/other', 'hydratest ALL=(ALL) NOPASSWD: ALL'),
            'live-greetd': ('etc/greetd/live.toml', 'live'),
        }
        if case in files:
            path, text = files[case]
            (root / path).write_text(text)
        if case == 'config-read-error':
            unreadable = root / 'etc/sudoers.d/unreadable'
            unreadable.write_text('configuration')
            unreadable.chmod(0)
        if case == 'cmdline-read-error':
            (root / 'proc/cmdline').unlink()
        commands = self.root / 'fixture-bin'
        commands.mkdir(exist_ok=True)
        dispatcher = commands / 'fixture'
        dispatcher.write_text('''#!/usr/bin/env python3
import os, pathlib, sys
name = pathlib.Path(sys.argv[0]).name
args = sys.argv[1:]
case = os.environ['SMOKE_CASE']
if name == 'id': print('hydra-live')
elif name == 'findmnt': print('/dev/vda2' if 'SOURCE' in args else ('overlay' if os.environ['SMOKE_LIVE']=='1' else 'ext4'))
elif name == 'systemctl':
    if args[0] == 'is-active':
        states = [unit != case for unit in args[1:]]
        print('\\n'.join('active' if state else 'failed' for state in states))
        sys.exit(0 if any(states) else 3)
    print('success' if 'Result' in args else '0')
elif name == 'getent':
    if args[:2] == ['passwd','hydra-live']:
        if case == 'account-lookup-error': sys.exit(3)
        if case != 'live-account': sys.exit(2)
        print('hydra-live:x:1000:1000::/home/hydra-live:/bin/bash')
    elif args[0] == 'shadow': print(args[1] + ':' + ('!' if args[1]=='root' else '$6$hash') + ':1::::::')
    else: print('127.0.0.1 STREAM example')
elif name == 'ip': print('10.203.0.1/24' if 'addr' in args else 'default via 10.0.2.2')
elif name == 'nmcli': print('connected')
else: print('fixture tool')
''')
        dispatcher.chmod(0o755)
        for command in ('id', 'findmnt', 'systemctl', 'getent', 'ip', 'nft',
                        'nmcli', 'sudo', 'hydra', 'hydra-install',
                        'qemu-system-x86_64', 'qemu-img', 'virt-viewer', 'foot', 'labwc'):
            path = commands / command
            if not path.exists():
                path.symlink_to(dispatcher)
        probe = self.mod.LIVE_CHECK if live else self.mod.INSTALLED_CHECK
        for path in ('/etc/', '/usr/share', '/proc/cmdline', '/sys/firmware/efi', '/run/initramfs/live/LiveOS/squashfs.img'):
            probe = probe.replace(path, str(root) + path)
        env = dict(os.environ, PATH=str(commands) + ':' + os.environ['PATH'],
                   SMOKE_CASE=case, SMOKE_LIVE='1' if live else '0')
        result = subprocess.run(['/bin/sh', '-c', self.mod.frame_command(probe, 'probecheck')],
                                env=env, capture_output=True, timeout=10)
        parsed = self.mod.parse_frame(result.stdout, 'probecheck')
        self.assertIsNotNone(parsed, result.stderr)
        return parsed[0]

    def test_installed_probe_rejects_each_retained_live_artifact(self):
        self.assertEqual(self.probe_fixture('clean'), 0)
        for case in ('live-cmdline', 'live-account', 'live-dropin', 'live-sudoers',
                     'autologin', 'nopasswd', 'live-marker', 'live-greetd'):
            with self.subTest(case=case):
                # Give each negative case an independent clean filesystem.
                import shutil
                shutil.rmtree(self.root / 'guest')
                self.assertNotEqual(self.probe_fixture(case), 0)

    def test_installed_probe_does_not_treat_account_query_error_as_absence(self):
        self.assertNotEqual(self.probe_fixture('account-lookup-error'), 0)

    def test_installed_probe_rejects_unreadable_configuration(self):
        self.assertNotEqual(self.probe_fixture('config-read-error'), 0)

    def test_installed_probe_rejects_unreadable_boot_arguments(self):
        self.assertNotEqual(self.probe_fixture('cmdline-read-error'), 0)

    def test_installed_probe_rejects_each_failed_required_service(self):
        self.assertEqual(self.probe_fixture('clean'), 0)
        for service in ('NetworkManager', 'greetd', 'hydra-network', 'hydra-dnsmasq'):
            with self.subTest(service=service):
                self.assertNotEqual(self.probe_fixture(service), 0)

    def test_live_probe_rejects_each_failed_required_service(self):
        self.assertEqual(self.probe_fixture('clean', live=True), 0)
        for service in ('NetworkManager', 'hydra-network', 'hydra-dnsmasq'):
            with self.subTest(service=service):
                self.assertNotEqual(self.probe_fixture(service, live=True), 0)

    def tcp_probe(self, fragments):
        # Execute the exact Python client embedded in the guest network probe,
        # redirecting only its address to a fresh loopback socket.
        import ast
        import time
        source = self.mod.NETWORK_CHECK.split("<<'NETWORK'\n", 1)[1].rsplit('\nNETWORK\n', 1)[0]
        assignments = (node for node in ast.walk(ast.parse(source)) if isinstance(node, ast.Assign))
        value = next(node.value for node in assignments
                     if any(isinstance(target, ast.Name) and target.id == 'client_code' for target in node.targets))
        client = ast.parse(ast.literal_eval(value))
        listener = socket.socket()
        self.addCleanup(listener.close)
        listener.bind(('127.0.0.1', 0))
        listener.listen(1)
        listener.settimeout(3)
        host, port = listener.getsockname()
        class Redirect(ast.NodeTransformer):
            def visit_Tuple(self, node):
                if isinstance(node.elts[0], ast.Constant) and node.elts[0].value == '10.203.0.250':
                    return ast.Tuple(elts=[ast.Constant(host), ast.Constant(port)], ctx=ast.Load())
                return self.generic_visit(node)
        client = ast.fix_missing_locations(Redirect().visit(client))
        def server():
            connection, _ = listener.accept()
            with connection:
                for fragment in fragments:
                    try:
                        connection.sendall(fragment)
                    except BrokenPipeError:
                        break
                    time.sleep(.05)
        thread = threading.Thread(target=server)
        thread.start()
        result = subprocess.run(['python3', '-c', ast.unparse(client)], capture_output=True, timeout=5)
        thread.join(timeout=3)
        self.assertFalse(thread.is_alive())
        return result

    def test_tcp_probe_accepts_complete_marker_across_fragmented_reads(self):
        result = self.tcp_probe([b'HYDRA_', b'LAN_TO_', b'VM_OK'])
        self.assertEqual(result.returncode, 0, result.stderr)

    def test_tcp_probe_rejects_truncated_marker(self):
        self.assertNotEqual(self.tcp_probe([b'HYDRA_LAN_TO_']).returncode, 0)