#!/usr/bin/env python3
"""Append a test-only stage-2 check while preserving the exported service database."""
import io
import pathlib
import os
import subprocess
import sys
import tarfile

project = pathlib.Path(__file__).resolve().parent.parent
if len(sys.argv) != 2:
    sys.exit('Usage: tools/make-vm-image EMPTY_WORK_DIRECTORY')
work = pathlib.Path(sys.argv[1]).resolve()
if not work.is_dir() or any(work.iterdir()):
    sys.exit('ERROR: VM image creation requires an existing empty work directory')
rootfs = (project / 'out/rootfs-cli.tar').resolve()
stage2 = 'etc/s6-linux-init/current/scripts/rc.init'
wrapper = (f'#!/bin/bash\nset -euo pipefail\n/{stage2}.fds-base "$@"\n'
           'exec /usr/libexec/fds/m2-guest\n').encode()
found_stage2 = False
with tarfile.open(rootfs, 'r') as original, tarfile.open(work / 'test-rootfs.tar', 'w', format=tarfile.PAX_FORMAT) as output:
    for member in original:
        if member.name == stage2:
            # Preserve the actual stage-2 script and run it to completion first.
            member.name += '.fds-base'
            found_stage2 = True
        output.addfile(member, original.extractfile(member) if member.isfile() else None)
    if not found_stage2:
        sys.exit('ERROR: rootfs has no native stage-2 script; rebuild it first')
    def root_owned(info):
        info.uid = info.gid = 0
        info.uname = info.gname = ''
        return info
    hook = tarfile.TarInfo(stage2)
    hook.mode = 0o755
    hook.size = len(wrapper)
    output.addfile(hook, io.BytesIO(wrapper))
    output.add(project / 'tests/integration/m2-guest', arcname='usr/libexec/fds/m2-guest', filter=root_owned)
subprocess.run(['mkfs.ext4', '-q', '-F', '-b', '4096', '-m', '0', '-L', 'FDS_M2_TEST',
                '-E', 'root_owner=0:0,lazy_itable_init=0,lazy_journal_init=0',
                '-d', str(work / 'test-rootfs.tar'), str(work / 'rootfs.ext4'), '768M'],
               env=dict(os.environ, LC_ALL='C.UTF-8'), check=True)
(work / 'command-line').write_text('console=ttyAMA0 root=/dev/sda rootfstype=ext4 ro rootwait init=/sbin/init panic=-1\n')
print('PASS: isolated VM image created from the exported rootfs')