#!/usr/bin/env python3
"""Hourly VM223 reset, preserving its full disk and pcap first. Never prunes evidence.

--check is read-only. --run stops intake and VM223, archives disk/pcap, validates,
then rolls back decoy-clean. On any failure it leaves VM223/intake stopped and
sets the host STOP latch. RAM is not preserved; collector evidence is untouched.
"""
import argparse
import datetime
import fcntl
import hashlib
import json
import os
from pathlib import Path
import shutil
import stat
import subprocess
import time
import uuid

ROOT = Path('/var/lib/vz/honeypot')
ARCHIVES = ROOT / 'archives'
DISK = Path('/dev/pve/vm-223-disk-0')
LIMIT = 16 * 1024**3
RESERVE = 4 * 1024**3


def command(*args, timeout=60):
    return subprocess.run(args, check=True, capture_output=True, text=True,
                          timeout=timeout).stdout


def qm(*args, timeout=60):
    return command('/usr/sbin/qm', *args, timeout=timeout)


def guest(script):
    result = json.loads(qm('guest', 'exec', '222', '--timeout', '20', '--',
                           '/bin/sh', '-c', script, timeout=25))
    if not result.get('exited') or result.get('exitcode') != 0:
        raise RuntimeError('collector guest operation failed')
    return result.get('out-data', '')


def sync_dir(path):
    fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)


def write_file(path, data):
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
    with os.fdopen(fd, 'wb') as stream:
        stream.write(data)
        stream.flush()
        os.fsync(stream.fileno())
    sync_dir(path.parent)


def bounds(min_free=RESERVE):
    used = sum(p.stat().st_size for p in ARCHIVES.rglob('*') if p.is_file()) if ARCHIVES.exists() else 0
    if used >= LIMIT:
        raise RuntimeError('archive 16GiB budget reached; evidence retained')
    if shutil.disk_usage(ROOT).free < min_free:
        raise RuntimeError('host free-space reserve reached; evidence retained')


def valid_config(config):
    values = dict(line.split(': ', 1) for line in config.splitlines() if ': ' in line)
    if (values.get('name') != 'decoy-execution' or
            not values.get('scsi0', '').startswith('local-lvm:vm-223-disk-0,') or
            'restrict=on' not in values.get('args', '') or
            'file=/var/lib/vz/honeypot/execution.pcap' not in values.get('args', '') or
            any(k.startswith('net') and k[3:].isdigit() for k in values)):
        raise RuntimeError('VM223 identity or isolation mismatch')


def validate():
    if os.geteuid() != 0:
        raise RuntimeError('run on designated Proxmox host as root')
    if ROOT.is_symlink() or not ROOT.is_dir() or (ARCHIVES.exists() and ARCHIVES.is_symlink()):
        raise RuntimeError('invalid evidence directory')
    if (ROOT / 'STOP').exists():
        raise RuntimeError('preexisting host STOP latch requires operator review')
    if not stat.S_ISBLK(DISK.stat().st_mode):
        raise RuntimeError('execution disk is not a block device')
    current = qm('config', '223')
    baseline = qm('config', '223', '--snapshot', 'decoy-clean')
    valid_config(current)
    valid_config(baseline)
    guest('test -f /etc/decoy-collector && test ! -e /var/lib/decoy/STOP')
    return current, baseline


def digest(path):
    hashed = hashlib.sha256()
    with path.open('rb') as stream:
        while chunk := stream.read(1024**2):
            bounds()
            hashed.update(chunk)
        os.fsync(stream.fileno())
    return hashed.hexdigest()


def install_admin_ca():
    public = ROOT / 'admin-ca.authorized'
    if not public.exists():
        return
    line = public.read_text().strip()
    if not line.startswith('cert-authority,principals="decoy-admin",restrict,pty ssh-ed25519 ') or len(line) > 4096 or '\n' in line:
        raise RuntimeError('invalid dedicated admin CA authorization')
    script = ("import pathlib,os; p=pathlib.Path('/root/.ssh/authorized_keys'); "
              "p.parent.mkdir(mode=0o700,exist_ok=True); "
              "lines=p.read_text().splitlines() if p.exists() else []; "
              "line=" + repr(line) + "; "
              "p.write_text('\\n'.join(lines if line in lines else lines+[line])+'\\n'); "
              "p.chmod(0o600)")
    result = json.loads(qm('guest', 'exec', '223', '--timeout', '10', '--', '/usr/bin/python3', '-c', script))
    if not result.get('exited') or result.get('exitcode') != 0:
        raise RuntimeError('admin certificate authorization restore failed')


def run():
    # Permanent lock inode; never unlink it while other cron instances may wait.
    with open('/run/decoy-reset.lock', 'a') as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            print('Reset already in progress')
            return
        validate()
        changed = False
        try:
            command('systemctl', 'stop', 'decoy-ingress.service')
            changed = True
            command('systemctl', 'stop', 'decoy-watchdog.service')
            guest('systemctl stop decoy.service decoy-network.service')
            qm('stop', '223')
            if 'status: stopped' not in qm('status', '223'):
                raise RuntimeError('execution VM did not stop')
            # Recheck after intake stops: never erase a concurrent capture-failure latch.
            current, baseline = validate()
            bounds(8 * 1024**3)
            stamp = datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%dT%H%M%SZ')
            ARCHIVES.mkdir(mode=0o700, exist_ok=True)
            archive = ARCHIVES / (stamp + '-' + uuid.uuid4().hex)
            archive.mkdir(mode=0o700)
            sync_dir(ARCHIVES)
            write_file(archive / 'vm223.conf', current.encode())
            write_file(archive / 'baseline.conf', baseline.encode())
            partial = archive / 'execution.qcow2.partial'
            # Reserve file exclusively before qemu-img writes it; protected parent0700.
            write_file(partial, b'')
            with (archive / 'convert.log').open('xb') as output:
                process = subprocess.Popen(['qemu-img', 'convert', '-c', '-f', 'raw', '-O',
                                            'qcow2', str(DISK), str(partial)], stdout=output, stderr=output)
                deadline = time.monotonic() + 900
                try:
                    while process.poll() is None:
                        bounds()
                        if time.monotonic() > deadline:
                            raise RuntimeError('disk archive timeout')
                        time.sleep(0.5)
                    if process.returncode:
                        raise RuntimeError('disk archive conversion failed; inspect retained convert.log')
                finally:
                    if process.poll() is None:
                        process.kill()
                        process.wait()
                    output.flush()
                    os.fsync(output.fileno())
            bounds()
            check = command('qemu-img', 'check', '--output=json', str(partial), timeout=300)
            command('qemu-img', 'compare', '-f', 'raw', '-F', 'qcow2', str(DISK), str(partial), timeout=300)
            write_file(archive / 'disk-check.json', check.encode())
            disk_hash = digest(partial)
            image = archive / 'execution.qcow2'
            partial.rename(image)
            sync_dir(archive)
            pcap = ROOT / 'execution.pcap'
            if pcap.is_symlink() or not pcap.is_file():
                raise RuntimeError('pcap is missing or not a regular file')
            pcap_hash = digest(pcap)
            pcap_size = pcap.stat().st_size
            if sum(p.stat().st_size for p in ARCHIVES.rglob('*') if p.is_file()) + pcap_size >= LIMIT:
                raise RuntimeError('pcap would exceed archive budget')
            pcap.rename(archive / 'execution.pcap')
            sync_dir(ROOT)
            sync_dir(archive)
            metadata = {'timestamp_utc': stamp, 'vmid': 223, 'baseline': 'decoy-clean',
                        'disk_sha256': disk_hash, 'pcap_sha256': pcap_hash,
                        'disk_bytes': image.stat().st_size, 'pcap_bytes': pcap_size,
                        'memory_preserved': False, 'state': 'validated_before_rollback'}
            write_file(archive / 'manifest.json', (json.dumps(metadata, indent=2) + '\n').encode())
            bounds()
            # A preexisting STOP is never cleared. Any late latch aborts reset.
            if (ROOT / 'STOP').exists():
                raise RuntimeError('host STOP appeared during archive')
            guest('test ! -e /var/lib/decoy/STOP')
            qm('rollback', '223', 'decoy-clean', timeout=300)
            valid_config(qm('config', '223'))
            write_file(ROOT / 'execution.pcap', b'')
            guest('test ! -e /var/lib/decoy/STOP && systemctl start decoy-network.service decoy.service && '
                  'systemctl is-active --quiet decoy-network.service decoy.service')
            qm('start', '223')
            deadline = time.monotonic() + 90
            while True:
                try:
                    ready = json.loads(qm('guest', 'exec', '223', '--timeout', '5', '--',
                        '/bin/sh', '-c', 'test -f /etc/decoy-disposable-guest && systemctl is-active --quiet decoy-relay decoy-guest-network', timeout=10))
                    if ready.get('exited') and ready.get('exitcode') == 0:
                        break
                except (subprocess.SubprocessError, ValueError):
                    pass
                if time.monotonic() >= deadline:
                    raise RuntimeError('restored guest did not become ready')
                time.sleep(2)
            install_admin_ca()
            command('systemctl', 'start', 'decoy-watchdog.service')
            command('systemctl', 'start', 'decoy-ingress.service')
            write_file(archive / 'reset-complete.json',
                       json.dumps({'completed_time_ns': time.time_ns()}).encode())
            print(f'Reset VM223; disk and packet evidence retained at {archive}')
        except BaseException:
            if changed:
                # Best effort fail-closed; never resume intake or roll back on failure.
                for args in [('systemctl', 'stop', 'decoy-ingress.service'),
                             ('/usr/sbin/qm', 'stop', '223')]:
                    try:
                        command(*args)
                    except Exception:
                        pass
                try:
                    guest('systemctl stop decoy.service decoy-network.service')
                except Exception:
                    pass
                try:
                    write_file(ROOT / 'STOP', b'automatic reset failed; operator review required\n')
                except FileExistsError:
                    pass
            raise
        finally:
            if changed:
                command('systemctl', 'start', 'decoy-watchdog.service')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument('--run', action='store_true')
    mode.add_argument('--check', action='store_true')
    args = parser.parse_args()
    os.umask(0o077)
    if args.check:
        validate()
        bounds(8 * 1024**3)
        print('Reset preflight passed; no changes made')
    else:
        run()