#!/usr/bin/env python3
"""Exercise the public Linux builder/writer with actual software and regular files."""
import argparse
import hashlib
import json
from pathlib import Path
import subprocess
import sys
import tempfile

project = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(project / 'tools'))
from image_formats import gpt, LINUX_FILESYSTEM
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--cli', type=Path, required=True)
parser.add_argument('--image-tool-runner', type=Path)
parser.add_argument('--void-packages', type=Path, default=project / '.host/void-packages')
parser.add_argument('--xbps-tool-runner', type=Path)
parser.add_argument('--xbps-bin', type=Path)
args = parser.parse_args()
work = Path(tempfile.mkdtemp(prefix='workstation-images.', dir=project / 'out'))
cli = [str(args.cli.resolve())]
if args.image_tool_runner:
    cli += ['--image-tool-runner', str(args.image_tool_runner)]
cli += ['--void-packages', str(args.void_packages.resolve())]
if args.xbps_tool_runner:
    cli += ['--xbps-tool-runner', str(args.xbps_tool_runner.resolve())]
if args.xbps_bin:
    cli += ['--xbps-bin', str(args.xbps_bin.resolve())]
log = (work / 'commands.log').open('w')


def invoke(arguments, ok=True):
    result = subprocess.run([*cli, *map(str, arguments)], capture_output=True, text=True)
    log.write(repr(arguments) + '\n' + result.stdout + result.stderr)
    log.flush()
    assert (result.returncode == 0) == ok, (arguments, result.returncode, result.stdout, result.stderr)
    return result.stdout


def digest(path):
    with path.open('rb') as stream:
        return hashlib.file_digest(stream, 'sha256').hexdigest()


# A user-provided submodule path must fail before copying any package sources.
upstream = project / 'vendor/void-packages'
assert not subprocess.check_output(['git', '-C', str(upstream), 'status', '--porcelain', '--untracked-files=all'])
rejected = subprocess.run([cli[0], '--void-packages', str(upstream), 'software', 'build',
                          str(project / 'examples/software/hello/software.toml'), str(work / 'forbidden-build')],
                         capture_output=True, text=True)
assert rejected.returncode != 0 and 'outside Git submodules' in rejected.stderr, rejected.stderr
assert not (work / 'forbidden-build').exists()

for name in ['hello', 'report']:
    recipe = project / f'examples/software/{name}/software.toml'
    if name == 'report':
        text = recipe.read_text().replace('[commands]', '[commands]\nshell="usr/bin/bash"\nwhere="usr/bin/printenv"')
        text = text.replace('template = "void"', 'template = ' + json.dumps(str(recipe.parent / 'void')))
        recipe = work / 'report-source.toml'; recipe.write_text(text)
    invoke(['software', 'build', recipe, work / f'{name}-installed'])
    inspected = json.loads(invoke(['software', 'inspect', work / f'{name}-installed']))
    assert inspected['installed'] and inspected['packages']
    assert f'fds-demo-{name}-1.0_1' in inspected['packages']
    assert any(p.startswith('glibc-') for p in inspected['packages'])
    assert (work / f'{name}-installed/root/usr/bin/{name}').is_file()
    assert not list((work / f'{name}-installed').glob('*.tar.xz'))
# The actual installed trees, including command links, are integrity checked.
program = work / 'hello-installed/root/usr/bin/hello'
original_program = program.read_bytes()
bad_elf = bytearray(original_program); bad_elf[18:20] = bytes([62, 0])
program.write_bytes(bad_elf)
invoke(['software', 'inspect', work / 'hello-installed'], False)
program.write_bytes(original_program)
escape = work / 'report-installed/root/usr/bin/escape'
escape.symlink_to('/etc/passwd')
invoke(['software', 'inspect', work / 'report-installed'], False)
escape.unlink()
invoke(['software', 'build', project / 'examples/software/hello/software.toml', work / 'hello-installed'], False)
(work / 'cartridge.toml').write_text('format=2\nid="demo.workstation"\nname="Workstation software"\nversion="1.0"\n[[payload]]\nsources=["hello-installed"]\n[[payload]]\nsources=["report-installed"]\n')
image = work / 'software.img'
original = json.loads(invoke(['create', work / 'cartridge.toml', image]))
assert len(original['image']['partitions']) == 3
assert [s['partition'] for s in original['catalogue']['software']] == [2, 3]
observed = json.loads(subprocess.check_output(['sfdisk', '--json', str(image)]))['partitiontable']
assert [p['name'] for p in observed['partitions']] == ['FDS_METADATA', 'FDS_PAYLOAD02', 'FDS_PAYLOAD03']
invoke(['create', work / 'cartridge.toml', work / 'repeat.img'])
assert digest(image) == digest(work / 'repeat.img')
invoke(['create', work / 'cartridge.toml', image], False)
(work / 'collision.toml').write_text((work / 'cartridge.toml').read_text().replace('demo.workstation', 'demo.second'))
invoke(['create', work / 'collision.toml', work / 'collision.img'])
shared = (work / 'cartridge.toml').read_text().replace('sources=["hello-installed"]\n[[payload]]\nsources=["report-installed"]', 'sources=["hello-installed","report-installed"]')
(work / 'shared.toml').write_text(shared)
grouped = json.loads(invoke(['create', work / 'shared.toml', work / 'shared.img']))
assert len(grouped['image']['partitions']) == 2
assert [s['partition'] for s in grouped['catalogue']['software']] == [2, 2]
for name in ['software.img', 'shared.img']:
    subprocess.run(['sfdisk', '--verify', str(work / name)], check=True, stdout=log, stderr=log)
size = image.stat().st_size
for label, extra in [('exact', 0), ('larger', 8 * 1024 * 1024)]:
    target = work / (label + '.target')
    with target.open('xb') as stream:
        stream.truncate(size + extra)
    preview = work / (label + '.preview.json')
    approval = json.loads(invoke(['preview', image, target, preview, '--file-target']))
    before = digest(target)
    invoke(['write', preview, '--confirm', 'WRONG'], False)
    assert digest(target) == before
    invoke(['write', preview, '--confirm', approval['confirmation']])
    subprocess.run(['sfdisk', '--verify', str(target)], check=True, stdout=log, stderr=log)
    assert json.loads(invoke(['inspect', target]))['catalogue'] == original['catalogue']
    invoke(['write', preview, '--confirm', approval['confirmation']], False)

target = work / 'changed-source.target'
with target.open('xb') as stream:
    stream.truncate(size)
source = work / 'changed-source.img'
source.write_bytes(image.read_bytes())
preview = work / 'changed-source.preview.json'
approval = json.loads(invoke(['preview', source, target, preview, '--file-target']))
before = digest(target)
with source.open('r+b') as stream:
    stream.seek(2 * 1024 * 1024 + 8192)
    stream.write(b'X')
invoke(['write', preview, '--confirm', approval['confirmation']], False)
assert digest(target) == before
bad = work / 'corrupt.img'
bad.write_bytes(image.read_bytes())
with bad.open('r+b') as stream:
    stream.seek(1024)
    stream.write(b'corrupt!')
invoke(['inspect', bad], False)
# Rebuild trusted fixture payloads independently, retaining valid GPT and EROFS
# while making the software data invalid. Both host and guest must reject them.
parts = original['image']['partitions']
filesystems = []
for part in parts:
    filesystem = work / f"part{part['number']}.erofs"
    with image.open('rb') as source, filesystem.open('wb') as target:
        source.seek(part['start'])
        target.write(source.read(part['bytes']))
    filesystems.append(filesystem)
tools = [str(args.image_tool_runner.resolve())] if args.image_tool_runner else []
for case in ['program-corrupt', 'catalogue-mapping']:
    number = 2 if case == 'program-corrupt' else 1
    tree = work / (case + '-tree')
    subprocess.run([*tools, 'fsck.erofs', '--extract=' + str(tree), str(filesystems[number-1])], check=True, stdout=log, stderr=log)
    if number == 2:
        program = tree / 'programs/demo.hello/usr/bin/hello'
        content = bytearray(program.read_bytes()); content[len(content)//2] ^= 1
        program.write_bytes(content)
    else:
        metadata = tree / 'FDS/SOFTWARE.TOML'
        metadata.write_text(metadata.read_text().replace('partition = 2', 'partition = 4'))
    changed = work / (case + '.erofs')
    subprocess.run([*tools, 'mkfs.erofs', '--quiet', '-b', '4096', '-T', '0', str(changed), str(tree)], check=True, stdout=log, stderr=log)
    selected = list(filesystems); selected[number-1] = changed
    malformed = work / (case + '.img')
    gpt(malformed, [(part['name'], LINUX_FILESYSTEM, filesystem) for part, filesystem in zip(parts, selected)])
    invoke(['inspect', malformed], False)
assert not subprocess.check_output(['git', '-C', str(upstream), 'status', '--porcelain', '--untracked-files=all'])
record = dict(source_checkout_clean=True, submodule_build_rejected=True, status='passed', work=str(work), cli_sha256=digest(args.cli.resolve()), compiled_aarch64_software=True,
              void_source_packages=True, installed_runtime_dependencies=True, wrong_elf_and_escaping_symlink_rejected=True, shared_and_separate_payload_partitions=True,
              repeat_image_identical=True, direct_installed_erofs_programs=True,
              exact_and_larger_target_readback=True, wrong_confirmation_unchanged=True,
              changed_source_unchanged_target=True, stale_target_preview_rejected=True,
              corrupted_gpt_rejected=True, corrupt_program_and_catalogue_rejected=True, physical_usb_write='not performed')
(work / 'acceptance.json').write_text(json.dumps(record, indent=2) + '\n')
(project / 'out/workstation-images-current.txt').write_text(str(work) + '\n')
print('PASS: workstation software/image/write acceptance:', work)
print('SKIP: physical USB writes and Raspberry Pi execution require hardware')