#!/usr/bin/env python3
# SPDX-License-Identifier: LGPL-2.1-or-later

# The tests can be called via pytest:
#   PATH=build/:$PATH pytest -v src/ukify/test/test_ukify.py
# or directly:
#   PATH=build/:$PATH src/ukify/test/test_ukify.py
# or via the meson test machinery output:
#   meson test -C build test-ukify -v
# or without verbose output:
#   meson test -C build test-ukify

# pylint: disable=unused-import,import-outside-toplevel,useless-else-on-loop
# pylint: disable=consider-using-with,wrong-import-position,unspecified-encoding
# pylint: disable=protected-access,redefined-outer-name

import base64
import glob
import json
import os
import pathlib
import re
import shutil
import subprocess
import sys
import tempfile
import textwrap

try:
    import pytest
except ImportError as e:
    print(str(e), file=sys.stderr)
    sys.exit(77)

try:
    # pyflakes: noqa
    import pefile  # noqa
except ImportError as e:
    print(str(e), file=sys.stderr)
    sys.exit(77)

# We import ukify.py, which is a template file. But only __version__ is
# substituted, which we don't care about here. Having the .py suffix makes it
# easier to import the file.
sys.path.append(os.path.dirname(__file__) + '/..')
import ukify

# Skip if we're running on an architecture that does not use UEFI.
try:
    ukify.guess_efi_arch()
except ValueError as e:
    print(str(e), file=sys.stderr)
    sys.exit(77)

build_root = os.getenv('PROJECT_BUILD_ROOT')
try:
    slow_tests = bool(int(os.getenv('SYSTEMD_SLOW_TESTS', '1')))
except ValueError:
    slow_tests = True

arg_tools = ['--tools', build_root] if build_root else []
if build_root and (
        p := pathlib.Path(f"{build_root}/linux{ukify.guess_efi_arch()}.efi.stub")
).exists():
    arg_tools += ['--stub', p]

def systemd_measure():
    opts = ukify.create_parser().parse_args(arg_tools)
    return ukify.find_tool('systemd-measure', opts=opts)

def test_guess_efi_arch():
    arch = ukify.guess_efi_arch()
    assert arch in ukify.EFI_ARCHES

def test_shell_join():
    assert ukify.shell_join(['a', 'b', ' ']) == "a b ' '"

def test_round_up():
    assert ukify.round_up(0) == 0
    assert ukify.round_up(4095) == 4096
    assert ukify.round_up(4096) == 4096
    assert ukify.round_up(4097) == 8192

def test_namespace_creation():
    ns = ukify.create_parser().parse_args(())
    assert ns.linux is None
    assert ns.initrd is None

def test_config_example():
    ex = ukify.config_example()
    assert '[UKI]' in ex
    assert 'Splash = BMP' in ex

def test_apply_config(tmp_path):
    config = tmp_path / 'config1.conf'
    config.write_text(textwrap.dedent(
        f'''
        [UKI]
        Linux = LINUX
        Initrd = initrd1 initrd2
                 initrd3
        Cmdline = 1 2 3 4 5
                  6 7 8
        OSRelease = @some/path1
        DeviceTree = some/path2
        Splash = some/path3
        Uname = 1.2.3
        EFIArch=arm
        Stub = some/path4
        PCRBanks = sha512,sha1
        SigningEngine = engine1
        SecureBootPrivateKey = some/path5
        SecureBootCertificate = some/path6
        SignKernel = no

        [PCRSignature:NAME]
        PCRPrivateKey = some/path7
        PCRPublicKey = some/path8
        Phases = {':'.join(ukify.KNOWN_PHASES)}
        '''))

    ns = ukify.create_parser().parse_args(['build'])
    ns.linux = None
    ns.initrd = []
    ukify.apply_config(ns, config)

    assert ns.linux == pathlib.Path('LINUX')
    assert ns.initrd == [pathlib.Path('initrd1'),
                         pathlib.Path('initrd2'),
                         pathlib.Path('initrd3')]
    assert ns.cmdline == '1 2 3 4 5\n6 7 8'
    assert ns.os_release == '@some/path1'
    assert ns.devicetree == pathlib.Path('some/path2')
    assert ns.splash == pathlib.Path('some/path3')
    assert ns.efi_arch == 'arm'
    assert ns.stub == pathlib.Path('some/path4')
    assert ns.pcr_banks == ['sha512', 'sha1']
    assert ns.signing_engine == 'engine1'
    assert ns.sb_key == 'some/path5'
    assert ns.sb_cert == 'some/path6'
    assert ns.sign_kernel is False

    assert ns._groups == ['NAME']
    assert ns.pcr_private_keys == ['some/path7']
    assert ns.pcr_public_keys == ['some/path8']
    assert ns.phase_path_groups == [['enter-initrd:leave-initrd:sysinit:ready:shutdown:final']]

    ukify.finalize_options(ns)

    assert ns.linux == pathlib.Path('LINUX')
    assert ns.initrd == [pathlib.Path('initrd1'),
                         pathlib.Path('initrd2'),
                         pathlib.Path('initrd3')]
    assert ns.cmdline == '1 2 3 4 5 6 7 8'
    assert ns.os_release == pathlib.Path('some/path1')
    assert ns.devicetree == pathlib.Path('some/path2')
    assert ns.splash == pathlib.Path('some/path3')
    assert ns.efi_arch == 'arm'
    assert ns.stub == pathlib.Path('some/path4')
    assert ns.pcr_banks == ['sha512', 'sha1']
    assert ns.signing_engine == 'engine1'
    assert ns.sb_key == 'some/path5'
    assert ns.sb_cert == pathlib.Path('some/path6')
    assert ns.sign_kernel is False

    assert ns._groups == ['NAME']
    assert ns.pcr_private_keys == ['some/path7']
    assert ns.pcr_public_keys == ['some/path8']
    assert ns.phase_path_groups == [['enter-initrd:leave-initrd:sysinit:ready:shutdown:final']]

def test_parse_args_minimal():
    with pytest.raises(ValueError):
        ukify.parse_args([])

    opts = ukify.parse_args('arg1 arg2'.split())
    assert opts.linux == pathlib.Path('arg1')
    assert opts.initrd == [pathlib.Path('arg2')]
    assert opts.os_release in (pathlib.Path('/etc/os-release'),
                               pathlib.Path('/usr/lib/os-release'))

def test_parse_args_many_deprecated():
    opts = ukify.parse_args(
        ['/ARG1', '///ARG2', '/ARG3 WITH SPACE',
         '--cmdline=a b c',
         '--os-release=K1=V1\nK2=V2',
         '--devicetree=DDDDTTTT',
         '--splash=splash',
         '--pcrpkey=PATH',
         '--uname=1.2.3',
         '--stub=STUBPATH',
         '--pcr-private-key=PKEY1',
         '--pcr-public-key=PKEY2',
         '--pcr-banks=SHA1,SHA256',
         '--signing-engine=ENGINE',
         '--secureboot-private-key=SBKEY',
         '--secureboot-certificate=SBCERT',
         '--sign-kernel',
         '--no-sign-kernel',
         '--tools=TOOLZ///',
         '--output=OUTPUT',
         '--measure',
         '--no-measure',
         ])
    assert opts.linux == pathlib.Path('/ARG1')
    assert opts.initrd == [pathlib.Path('/ARG2'), pathlib.Path('/ARG3 WITH SPACE')]
    assert opts.cmdline == 'a b c'
    assert opts.os_release == 'K1=V1\nK2=V2'
    assert opts.devicetree == pathlib.Path('DDDDTTTT')
    assert opts.splash == pathlib.Path('splash')
    assert opts.pcrpkey == pathlib.Path('PATH')
    assert opts.uname == '1.2.3'
    assert opts.stub == pathlib.Path('STUBPATH')
    assert opts.pcr_private_keys == ['PKEY1']
    assert opts.pcr_public_keys == ['PKEY2']
    assert opts.pcr_banks == ['SHA1', 'SHA256']
    assert opts.signing_engine == 'ENGINE'
    assert opts.sb_key == 'SBKEY'
    assert opts.sb_cert == pathlib.Path('SBCERT')
    assert opts.sign_kernel is False
    assert opts.tools == [pathlib.Path('TOOLZ/')]
    assert opts.output == pathlib.Path('OUTPUT')
    assert opts.measure is False

def test_parse_args_many():
    opts = ukify.parse_args(
        ['build',
         '--linux=/ARG1',
         '--initrd=///ARG2',
         '--initrd=/ARG3 WITH SPACE',
         '--cmdline=a b c',
         '--os-release=K1=V1\nK2=V2',
         '--devicetree=DDDDTTTT',
         '--splash=splash',
         '--pcrpkey=PATH',
         '--uname=1.2.3',
         '--stub=STUBPATH',
         '--pcr-private-key=PKEY1',
         '--pcr-public-key=PKEY2',
         '--pcr-banks=SHA1,SHA256',
         '--signing-engine=ENGINE',
         '--secureboot-private-key=SBKEY',
         '--secureboot-certificate=SBCERT',
         '--sign-kernel',
         '--no-sign-kernel',
         '--tools=TOOLZ///',
         '--output=OUTPUT',
         '--measure',
         '--no-measure',
         '--policy-digest',
         '--no-policy-digest',
         ])
    assert opts.linux == pathlib.Path('/ARG1')
    assert opts.initrd == [pathlib.Path('/ARG2'), pathlib.Path('/ARG3 WITH SPACE')]
    assert opts.cmdline == 'a b c'
    assert opts.os_release == 'K1=V1\nK2=V2'
    assert opts.devicetree == pathlib.Path('DDDDTTTT')
    assert opts.splash == pathlib.Path('splash')
    assert opts.pcrpkey == pathlib.Path('PATH')
    assert opts.uname == '1.2.3'
    assert opts.stub == pathlib.Path('STUBPATH')
    assert opts.pcr_private_keys == ['PKEY1']
    assert opts.pcr_public_keys == ['PKEY2']
    assert opts.pcr_banks == ['SHA1', 'SHA256']
    assert opts.signing_engine == 'ENGINE'
    assert opts.sb_key == 'SBKEY'
    assert opts.sb_cert == pathlib.Path('SBCERT')
    assert opts.sign_kernel is False
    assert opts.tools == [pathlib.Path('TOOLZ/')]
    assert opts.output == pathlib.Path('OUTPUT')
    assert opts.measure is False
    assert opts.policy_digest is False

def test_parse_sections():
    opts = ukify.parse_args(
        ['build',
         '--linux=/ARG1',
         '--initrd=/ARG2',
         '--section=test:TESTTESTTEST',
         '--section=test2:@FILE',
         ])

    assert opts.linux == pathlib.Path('/ARG1')
    assert opts.initrd == [pathlib.Path('/ARG2')]
    assert len(opts.sections) == 2

    assert opts.sections[0].name == 'test'
    assert isinstance(opts.sections[0].content, pathlib.Path)
    assert opts.sections[0].tmpfile
    assert opts.sections[0].measure is False

    assert opts.sections[1].name == 'test2'
    assert opts.sections[1].content == pathlib.Path('FILE')
    assert opts.sections[1].tmpfile is None
    assert opts.sections[1].measure is False

def test_config_priority(tmp_path):
    config = tmp_path / 'config1.conf'
    # config: use pesign and give certdir + certname
    config.write_text(textwrap.dedent(
        f'''
        [UKI]
        Linux = LINUX
        Initrd = initrd1 initrd2
                 initrd3
        Cmdline = 1 2 3 4 5
                  6 7 8
        OSRelease = @some/path1
        DeviceTree = some/path2
        Splash = some/path3
        Uname = 1.2.3
        EFIArch = arm
        Stub = some/path4
        PCRBanks = sha512,sha1
        SigningEngine = engine1
        SecureBootSigningTool = pesign
        SecureBootCertificateDir = some/path5
        SecureBootCertificateName = some/name1
        SignKernel = no

        [PCRSignature:NAME]
        PCRPrivateKey = some/path7
        PCRPublicKey = some/path8
        Phases = {':'.join(ukify.KNOWN_PHASES)}
        '''))

    # args: use sbsign and give key + cert, should override pesign
    opts = ukify.parse_args(
        ['build',
         '--linux=/ARG1',
         '--initrd=///ARG2',
         '--initrd=/ARG3 WITH SPACE',
         '--cmdline= a  b  c ',
         '--os-release=K1=V1\nK2=V2',
         '--devicetree=DDDDTTTT',
         '--splash=splash',
         '--pcrpkey=PATH',
         '--uname=1.2.3',
         '--stub=STUBPATH',
         '--pcr-private-key=PKEY1',
         '--pcr-public-key=PKEY2',
         '--pcr-banks=SHA1,SHA256',
         '--signing-engine=ENGINE',
         '--signtool=sbsign',
         '--secureboot-private-key=SBKEY',
         '--secureboot-certificate=SBCERT',
         '--sign-kernel',
         '--no-sign-kernel',
         '--tools=TOOLZ///',
         '--output=OUTPUT',
         '--measure',
         ])

    ukify.apply_config(opts, config)
    ukify.finalize_options(opts)

    assert opts.linux == pathlib.Path('/ARG1')
    assert opts.initrd == [pathlib.Path('initrd1'),
                           pathlib.Path('initrd2'),
                           pathlib.Path('initrd3'),
                           pathlib.Path('/ARG2'),
                           pathlib.Path('/ARG3 WITH SPACE')]
    assert opts.cmdline == 'a b c'
    assert opts.os_release == 'K1=V1\nK2=V2'
    assert opts.devicetree == pathlib.Path('DDDDTTTT')
    assert opts.splash == pathlib.Path('splash')
    assert opts.pcrpkey == pathlib.Path('PATH')
    assert opts.uname == '1.2.3'
    assert opts.stub == pathlib.Path('STUBPATH')
    assert opts.pcr_private_keys == ['PKEY1', 'some/path7']
    assert opts.pcr_public_keys == ['PKEY2', 'some/path8']
    assert opts.pcr_banks == ['SHA1', 'SHA256']
    assert opts.signing_engine == 'ENGINE'
    assert opts.signtool == 'sbsign' # from args
    assert opts.sb_key == 'SBKEY' # from args
    assert opts.sb_cert == pathlib.Path('SBCERT') # from args
    assert opts.sb_certdir == 'some/path5' # from config
    assert opts.sb_cert_name == 'some/name1' # from config
    assert opts.sign_kernel is False
    assert opts.tools == [pathlib.Path('TOOLZ/')]
    assert opts.output == pathlib.Path('OUTPUT')
    assert opts.measure is True

def test_help(capsys):
    with pytest.raises(SystemExit):
        ukify.parse_args(['--help'])
    out = capsys.readouterr()
    assert '--section' in out.out
    assert not out.err

def test_help_display(capsys):
    with pytest.raises(SystemExit):
        ukify.parse_args(['inspect', '--help'])
    out = capsys.readouterr()
    assert '--section' in out.out
    assert not out.err

def test_help_error_deprecated(capsys):
    with pytest.raises(SystemExit):
        ukify.parse_args(['a', 'b', '--no-such-option'])
    out = capsys.readouterr()
    assert not out.out
    assert '--no-such-option' in out.err
    assert len(out.err.splitlines()) == 1

def test_help_error(capsys):
    with pytest.raises(SystemExit):
        ukify.parse_args(['build', '--no-such-option'])
    out = capsys.readouterr()
    assert not out.out
    assert '--no-such-option' in out.err
    assert len(out.err.splitlines()) == 1

@pytest.fixture(scope='session')
def kernel_initrd():
    items = sorted(glob.glob('/lib/modules/*/vmlinuz'))
    if not items:
        items = sorted(glob.glob('/boot/vmlinuz*'))
    if not items:
        return None

    # This doesn't necessarily give us the latest version, since we're just
    # using alphanumeric ordering. But this is fine, a predictable result is
    # enough.
    linux = items[-1]

    # We don't look _into_ the initrd. Any file is OK.
    return ['--linux', linux, '--initrd', ukify.__file__]

def test_check_splash():
    try:
        # pyflakes: noqa
        import PIL  # noqa
    except ImportError:
        pytest.skip('PIL not available')

    with pytest.raises(OSError):
        ukify.check_splash(os.devnull)

def test_basic_operation(kernel_initrd, tmp_path):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')

    output = f'{tmp_path}/basic.efi'
    opts = ukify.parse_args([
        'build',
        *kernel_initrd,
        f'--output={output}',
    ])
    try:
        ukify.check_inputs(opts)
    except OSError as e:
        pytest.skip(str(e))

    ukify.make_uki(opts)

    # let's check that objdump likes the resulting file
    subprocess.check_output(['objdump', '-h', output])

    shutil.rmtree(tmp_path)

def test_sections(kernel_initrd, tmp_path):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')

    output = f'{tmp_path}/basic.efi'
    opts = ukify.parse_args([
        'build',
        *kernel_initrd,
        f'--output={output}',
        '--uname=1.2.3',
        '--cmdline=ARG1 ARG2 ARG3',
        '--os-release=K1=V1\nK2=V2\n',
        '--section=.test:CONTENTZ',
    ])

    try:
        ukify.check_inputs(opts)
    except OSError as e:
        pytest.skip(str(e))

    ukify.make_uki(opts)

    # let's check that objdump likes the resulting file
    dump = subprocess.check_output(['objdump', '-h', output], text=True)

    for sect in 'text osrel cmdline linux initrd uname test'.split():
        assert re.search(fr'^\s*\d+\s+\.{sect}\s+[0-9a-f]+', dump, re.MULTILINE)

    shutil.rmtree(tmp_path)

def test_addon(tmp_path):
    output = f'{tmp_path}/addon.efi'
    args = [
        'build',
        f'--output={output}',
        '--cmdline=ARG1 ARG2 ARG3',
        """--sbat=sbat,1,foo
foo,1
bar,2
""",
        '--section=.test:CONTENTZ',
        """--sbat=sbat,1,foo
baz,3
"""
    ]
    if stub := os.getenv('EFI_ADDON'):
        args += [f'--stub={stub}']
        expected_exceptions = ()
    else:
        expected_exceptions = (FileNotFoundError,)

    opts = ukify.parse_args(args)
    try:
        ukify.check_inputs(opts)
    except expected_exceptions as e:
        pytest.skip(str(e))

    ukify.make_uki(opts)

    # let's check that objdump likes the resulting file
    dump = subprocess.check_output(['objdump', '-h', output], text=True)

    for sect in 'text cmdline test sbat'.split():
        assert re.search(fr'^\s*\d+\s+\.{sect}\s+[0-9a-f]+', dump, re.MULTILINE)

    pe = pefile.PE(output, fast_load=True)
    found = False

    for section in pe.sections:
        if section.Name.rstrip(b"\x00").decode() == ".sbat":
            assert found is False
            split = section.get_data().rstrip(b"\x00").decode().splitlines()
            assert split == ["sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md", "foo,1", "bar,2", "baz,3"]
            found = True

    assert found is True

def unbase64(filename):
    tmp = tempfile.NamedTemporaryFile()
    base64.decode(filename.open('rb'), tmp)
    tmp.flush()
    return tmp

def test_uname_scraping(kernel_initrd):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')

    assert kernel_initrd[0] == '--linux'
    uname = ukify.Uname.scrape(kernel_initrd[1])
    assert re.match(r'\d+\.\d+\.\d+', uname)

@pytest.mark.skipif(not slow_tests, reason='slow')
@pytest.mark.parametrize("days", [365*10, None])
def test_efi_signing_sbsign(days, kernel_initrd, tmp_path):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')
    if not shutil.which('sbsign'):
        pytest.skip('sbsign not found')

    ourdir = pathlib.Path(__file__).parent
    cert = unbase64(ourdir / 'example.signing.crt.base64')
    key = unbase64(ourdir / 'example.signing.key.base64')

    output = f'{tmp_path}/signed.efi'
    args = [
        'build',
        *kernel_initrd,
        f'--output={output}',
        '--uname=1.2.3',
        '--cmdline=ARG1 ARG2 ARG3',
        f'--secureboot-certificate={cert.name}',
        f'--secureboot-private-key={key.name}',
    ]
    if days is not None:
        args += [f'--secureboot-certificate-validity={days}']

    opts = ukify.parse_args(args)

    try:
        ukify.check_inputs(opts)
    except OSError as e:
        pytest.skip(str(e))

    ukify.make_uki(opts)

    if shutil.which('sbverify'):
        # let's check that sbverify likes the resulting file
        dump = subprocess.check_output([
            'sbverify',
            '--cert', cert.name,
            output,
        ], text=True)

        assert 'Signature verification OK' in dump

    shutil.rmtree(tmp_path)

@pytest.mark.skipif(not slow_tests, reason='slow')
def test_efi_signing_pesign(kernel_initrd, tmp_path):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')
    if not shutil.which('pesign'):
        pytest.skip('pesign not found')

    nss_db = f'{tmp_path}/nss_db'
    name = 'Test_Secureboot'
    author = 'systemd'

    subprocess.check_call(['mkdir', '-p', nss_db])
    cmd = f'certutil -N --empty-password -d {nss_db}'.split(' ')
    subprocess.check_call(cmd)
    cmd = f'efikeygen -d {nss_db} -S -k -c CN={author} -n {name}'.split(' ')
    subprocess.check_call(cmd)

    output = f'{tmp_path}/signed.efi'
    opts = ukify.parse_args([
        'build',
        *kernel_initrd,
        f'--output={output}',
        '--uname=1.2.3',
        '--signtool=pesign',
        '--cmdline=ARG1 ARG2 ARG3',
        f'--secureboot-certificate-name={name}',
        f'--secureboot-certificate-dir={nss_db}',
    ])

    try:
        ukify.check_inputs(opts)
    except OSError as e:
        pytest.skip(str(e))

    ukify.make_uki(opts)

    # let's check that pesign likes the resulting file
    dump = subprocess.check_output([
        'pesign', '-S',
        '-i', output,
    ], text=True)

    assert f"The signer's common name is {author}" in dump

    shutil.rmtree(tmp_path)

def test_inspect(kernel_initrd, tmp_path, capsys):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')
    if not shutil.which('sbsign'):
        pytest.skip('sbsign not found')

    ourdir = pathlib.Path(__file__).parent
    cert = unbase64(ourdir / 'example.signing.crt.base64')
    key = unbase64(ourdir / 'example.signing.key.base64')

    output = f'{tmp_path}/signed2.efi'
    uname_arg='1.2.3'
    osrel_arg='Linux'
    cmdline_arg='ARG1 ARG2 ARG3'

    args = [
        'build',
        *kernel_initrd,
        f'--cmdline={cmdline_arg}',
        f'--os-release={osrel_arg}',
        f'--uname={uname_arg}',
        f'--output={output}',
    ] + arg_tools
    if slow_tests:
        args += [
            f'--secureboot-certificate={cert.name}',
            f'--secureboot-private-key={key.name}',
        ]

    opts = ukify.parse_args(args)

    ukify.check_inputs(opts)
    ukify.make_uki(opts)

    opts = ukify.parse_args(['inspect', output])
    ukify.inspect_sections(opts)

    text = capsys.readouterr().out

    expected_osrel = f'.osrel:\n  size: {len(osrel_arg)}'
    assert expected_osrel in text
    expected_cmdline = f'.cmdline:\n  size: {len(cmdline_arg)}'
    assert expected_cmdline in text
    expected_uname = f'.uname:\n  size: {len(uname_arg)}'
    assert expected_uname in text

    expected_initrd = '.initrd:\n  size:'
    assert expected_initrd in text
    expected_linux = '.linux:\n  size:'
    assert expected_linux in text

    shutil.rmtree(tmp_path)

@pytest.mark.skipif(not slow_tests, reason='slow')
def test_pcr_signing(kernel_initrd, tmp_path):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')
    try:
        systemd_measure()
    except ValueError:
        pytest.skip('systemd-measure not found')

    ourdir = pathlib.Path(__file__).parent
    pub = unbase64(ourdir / 'example.tpm2-pcr-public.pem.base64')
    priv = unbase64(ourdir / 'example.tpm2-pcr-private.pem.base64')

    output = f'{tmp_path}/signed.efi'
    args = [
        'build',
        *kernel_initrd,
        f'--output={output}',
        '--uname=1.2.3',
        '--cmdline=ARG1 ARG2 ARG3',
        '--os-release=ID=foobar\n',
        '--pcr-banks=sha384',   # sha1 might not be allowed, use something else
        f'--pcr-private-key={priv.name}',
    ] + arg_tools

    # If the public key is not explicitly specified, it is derived
    # automatically. Let's make sure everything works as expected both when the
    # public keys is specified explicitly and when it is derived from the
    # private key.
    for extra in ([f'--pcrpkey={pub.name}', f'--pcr-public-key={pub.name}'], []):
        opts = ukify.parse_args(args + extra)
        try:
            ukify.check_inputs(opts)
        except OSError as e:
            pytest.skip(str(e))

        ukify.make_uki(opts)

        # let's check that objdump likes the resulting file
        dump = subprocess.check_output(['objdump', '-h', output], text=True)

        for sect in 'text osrel cmdline linux initrd uname pcrsig'.split():
            assert re.search(fr'^\s*\d+\s+\.{sect}\s+[0-9a-f]+', dump, re.MULTILINE)

        # objcopy fails when called without an output argument (EPERM).
        # It also fails when called with /dev/null (file truncated).
        # It also fails when called with /dev/zero (because it reads the
        # output file, infinitely in this case.)
        # So let's just call it with a dummy output argument.
        subprocess.check_call([
            'objcopy',
            *(f'--dump-section=.{n}={tmp_path}/out.{n}' for n in (
                'pcrpkey', 'pcrsig', 'osrel', 'uname', 'cmdline')),
            output,
            tmp_path / 'dummy',
        ],
            text=True)

        assert open(tmp_path / 'out.pcrpkey').read() == open(pub.name).read()
        assert open(tmp_path / 'out.osrel').read() == 'ID=foobar\n'
        assert open(tmp_path / 'out.uname').read() == '1.2.3'
        assert open(tmp_path / 'out.cmdline').read() == 'ARG1 ARG2 ARG3'
        sig = open(tmp_path / 'out.pcrsig').read()
        sig = json.loads(sig)
        assert list(sig.keys()) == ['sha384']
        assert len(sig['sha384']) == 4   # four items for four phases

    shutil.rmtree(tmp_path)

@pytest.mark.skipif(not slow_tests, reason='slow')
def test_pcr_signing2(kernel_initrd, tmp_path):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')
    try:
        systemd_measure()
    except ValueError:
        pytest.skip('systemd-measure not found')

    ourdir = pathlib.Path(__file__).parent
    pub = unbase64(ourdir / 'example.tpm2-pcr-public.pem.base64')
    priv = unbase64(ourdir / 'example.tpm2-pcr-private.pem.base64')
    pub2 = unbase64(ourdir / 'example.tpm2-pcr-public2.pem.base64')
    priv2 = unbase64(ourdir / 'example.tpm2-pcr-private2.pem.base64')

    # simulate a microcode file
    with open(f'{tmp_path}/microcode', 'wb') as microcode:
        microcode.write(b'1234567890')

    output = f'{tmp_path}/signed.efi'
    assert kernel_initrd[0] == '--linux'
    opts = ukify.parse_args([
        'build',
        *kernel_initrd[:2],
        f'--initrd={microcode.name}',
        *kernel_initrd[2:],
        f'--output={output}',
        '--uname=1.2.3',
        '--cmdline=ARG1 ARG2 ARG3',
        '--os-release=ID=foobar\n',
        '--pcr-banks=sha384',
        f'--pcrpkey={pub2.name}',
        f'--pcr-public-key={pub.name}',
        f'--pcr-private-key={priv.name}',
        '--phases=enter-initrd enter-initrd:leave-initrd',
        f'--pcr-public-key={pub2.name}',
        f'--pcr-private-key={priv2.name}',
        '--phases=sysinit ready shutdown final',  # yes, those phase paths are not reachable
    ] + arg_tools)

    try:
        ukify.check_inputs(opts)
    except OSError as e:
        pytest.skip(str(e))

    ukify.make_uki(opts)

    # let's check that objdump likes the resulting file
    dump = subprocess.check_output(['objdump', '-h', output], text=True)

    for sect in 'text osrel cmdline linux initrd uname pcrsig'.split():
        assert re.search(fr'^\s*\d+\s+\.{sect}\s+[0-9a-f]+', dump, re.MULTILINE)

    subprocess.check_call([
        'objcopy',
        *(f'--dump-section=.{n}={tmp_path}/out.{n}' for n in (
            'pcrpkey', 'pcrsig', 'osrel', 'uname', 'cmdline', 'initrd')),
        output,
        tmp_path / 'dummy',
    ],
        text=True)

    assert open(tmp_path / 'out.pcrpkey').read() == open(pub2.name).read()
    assert open(tmp_path / 'out.osrel').read() == 'ID=foobar\n'
    assert open(tmp_path / 'out.uname').read() == '1.2.3'
    assert open(tmp_path / 'out.cmdline').read() == 'ARG1 ARG2 ARG3'
    assert open(tmp_path / 'out.initrd', 'rb').read(10) == b'1234567890'

    sig = open(tmp_path / 'out.pcrsig').read()
    sig = json.loads(sig)
    assert list(sig.keys()) == ['sha384']
    assert len(sig['sha384']) == 6   # six items for six phases paths

    shutil.rmtree(tmp_path)

def test_key_cert_generation(tmp_path):
    opts = ukify.parse_args([
        'genkey',
        f"--pcr-public-key={tmp_path / 'pcr1.pub.pem'}",
        f"--pcr-private-key={tmp_path / 'pcr1.priv.pem'}",
        '--phases=enter-initrd enter-initrd:leave-initrd',
        f"--pcr-public-key={tmp_path / 'pcr2.pub.pem'}",
        f"--pcr-private-key={tmp_path / 'pcr2.priv.pem'}",
        '--phases=sysinit ready',
        f"--secureboot-private-key={tmp_path / 'sb.priv.pem'}",
        f"--secureboot-certificate={tmp_path / 'sb.cert.pem'}",
    ])
    assert opts.verb == 'genkey'
    ukify.check_cert_and_keys_nonexistent(opts)

    pytest.importorskip('cryptography')

    ukify.generate_keys(opts)

    if not shutil.which('openssl'):
        return

    for key in (tmp_path / 'pcr1.priv.pem',
                tmp_path / 'pcr2.priv.pem',
                tmp_path / 'sb.priv.pem'):
        out = subprocess.check_output([
            'openssl', 'rsa',
            '-in', key,
            '-text',
            '-noout',
        ], text = True)
        assert 'Private-Key' in out
        assert '2048 bit' in out

    for pub in (tmp_path / 'pcr1.pub.pem',
                tmp_path / 'pcr2.pub.pem'):
        out = subprocess.check_output([
            'openssl', 'rsa',
            '-pubin',
            '-in', pub,
            '-text',
            '-noout',
        ], text = True)
        assert 'Public-Key' in out
        assert '2048 bit' in out

    out = subprocess.check_output([
        'openssl', 'x509',
        '-in', tmp_path / 'sb.cert.pem',
        '-text',
        '-noout',
    ], text = True)
    assert 'Certificate' in out
    assert re.search(r'Issuer: CN\s?=\s?SecureBoot signing key on host', out)

@pytest.mark.skipif(not slow_tests, reason='slow')
def test_join_pcrsig(capsys, kernel_initrd, tmp_path):
    if kernel_initrd is None:
        pytest.skip('linux+initrd not found')
    try:
        systemd_measure()
    except ValueError:
        pytest.skip('systemd-measure not found')

    ourdir = pathlib.Path(__file__).parent
    pub = unbase64(ourdir / 'example.tpm2-pcr-public.pem.base64')

    output = tmp_path / 'basic.efi'
    args = [
        'build',
        *kernel_initrd,
        f'--output={output}',
        f'--pcr-public-key={pub.name}',
        '--json=short',
        '--policy-digest',
    ] + arg_tools
    opts = ukify.parse_args(args)
    try:
        ukify.check_inputs(opts)
    except OSError as e:
        pytest.skip(str(e))

    ukify.make_uki(opts)
    pcrs = json.loads(capsys.readouterr().out)
    for bank, sigs in pcrs.items():
        for sig in sigs:
            sig['sig'] = 'a' * int(bank[3:])

    opts = ukify.parse_args(['inspect', str(output)])
    ukify.inspect_sections(opts)
    text = capsys.readouterr().out
    assert re.search(r'\.pcrpkey', text, re.MULTILINE)
    assert re.search(r'\.pcrsig', text, re.MULTILINE)
    assert not re.search(r'"sig":', text, re.MULTILINE)

    output_sig = tmp_path / 'pcrsig.efi'
    args = [
        'build',
        f'--output={output_sig}',
        f'--join-pcrsig={output}',
        f'--pcrsig={json.dumps(pcrs)}',
        '--json=short',
    ] + arg_tools
    opts = ukify.parse_args(args)
    try:
        ukify.check_inputs(opts)
    except OSError as e:
        pytest.skip(str(e))

    ukify.make_uki(opts)

    opts = ukify.parse_args(['inspect', str(output_sig)])
    ukify.inspect_sections(opts)
    text = capsys.readouterr().out
    assert re.search(r'\.pcrpkey', text, re.MULTILINE)
    assert re.search(r'\.pcrsig', text, re.MULTILINE)
    assert re.search(r'"sig":', text, re.MULTILINE)

    shutil.rmtree(tmp_path)

if __name__ == '__main__':
    sys.exit(pytest.main(sys.argv))
