#!/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('--cc', nargs='+', default=['aarch64-linux-gnu-gcc'])
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)]
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()


(work / 'hello-root/bin').mkdir(parents=True)
(work / 'report-root/bin').mkdir(parents=True)
(work / 'hello.c').write_text('#include <stdio.h>\nint main(void) { puts("FDS cartridge AArch64 hello"); return 0; }\n')
script = work / 'report-root/bin/report'
script.write_text('#!/bin/sh\nprintf "FDS cartridge script report\\n"\nid\n[ "${1-}" != hold ] || exec tail -f /dev/null\n')
script.chmod(0o755)
recipe = f'''format=1
id="demo.hello"
name="AArch64 hello"
version="1.0"
architecture="aarch64"
root="hello-root"
[commands]
hello="bin/hello"
[build]
directory={json.dumps(str(project))}
command={json.dumps([*args.cc, '-O2', str(work / 'hello.c'), '-o', str(work / 'hello-root/bin/hello')])}
'''
(work / 'hello.toml').write_text(recipe)
(work / 'report.toml').write_text('format=1\nid="demo.report"\nname="Script report"\nversion="1.0"\narchitecture="any"\nroot="report-root"\n[commands]\nreport="bin/report"\n')
invoke(['software', 'build', work / 'hello.toml', work / 'hello-bundle'])
invoke(['software', 'pack', work / 'report.toml', work / 'report-bundle'])
# Architecture-independent bundles cannot hide actual ELF executables.
(work / 'wrong-architecture.toml').write_text(recipe.replace('architecture="aarch64"', 'architecture="any"'))
invoke(['software', 'pack', work / 'wrong-architecture.toml', work / 'rejected-architecture'], False)
assert not (work / 'rejected-architecture').exists()
# Source links outside the software root must never become bundle contents.
(work / 'report-root/bin/escape').symlink_to('/etc/passwd')
invoke(['software', 'pack', work / 'report.toml', work / 'rejected-symlink'], False)
assert not (work / 'rejected-symlink').exists()
(work / 'report-root/bin/escape').unlink()
invoke(['software', 'pack', work / 'report.toml', work / 'report-bundle'], False)
(work / 'cartridge.toml').write_text('format=1\nid="demo.workstation"\nname="Workstation software"\nversion="1.0"\n[[payload]]\nbundles=["hello-bundle"]\n[[payload]]\nbundles=["report-bundle"]\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)
shared = (work / 'cartridge.toml').read_text().replace('bundles=["hello-bundle"]\n[[payload]]\nbundles=["report-bundle"]', 'bundles=["hello-bundle","report-bundle"]')
(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)
for directory in ['hello-bundle', 'report-bundle']:
    archive = next((work / directory).glob('*.tar.xz'))
    subprocess.run(['xz', '--test', str(archive)], check=True)
    subprocess.run(['tar', '-tJf', str(archive)], check=True, stdout=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 ['archive-corrupt', 'catalogue-mapping']:
    number = 2 if case == 'archive-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:
        archive = next((tree / 'bundles').glob('*.tar.xz'))
        content = bytearray(archive.read_bytes()); content[len(content)//2] ^= 1
        archive.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)
record = dict(status='passed', work=str(work), cli_sha256=digest(args.cli.resolve()), compiled_aarch64_software=True,
              script_bundle=True, elf_in_any_bundle_and_escaping_source_symlink_rejected=True, shared_and_separate_payload_partitions=True,
              repeat_image_identical=True, independent_gpt_xz_tar_checks=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_archive_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')