Merge pull request #27621 from esposem/ukify_read

ukify: add an option to read the PE sections

2 files changed, 206 insertions, 37 deletions

diff --git a/src/ukify/test/test_ukify.py b/src/ukify/test/test_ukify.py
index e0e85455c5..3cc3492fe0 100755
--- a/src/ukify/test/test_ukify.py
+++ b/src/ukify/test/test_ukify.py
@@ -360,6 +360,13 @@ def test_help(capsys):
     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'])
@@ -603,6 +610,52 @@ def test_efi_signing_pesign(kernel_initrd, tmpdir):
 
     assert f"The signer's common name is {author}" in dump
 
+def test_inspect(kernel_initrd, tmpdir, 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'{tmpdir}/signed2.efi'
+    uname_arg='1.2.3'
+    osrel_arg='Linux'
+    cmdline_arg='ARG1 ARG2 ARG3'
+    opts = ukify.parse_args([
+        'build',
+        *kernel_initrd,
+        f'--cmdline={cmdline_arg}',
+        f'--os-release={osrel_arg}',
+        f'--uname={uname_arg}',
+        f'--output={output}',
+        f'--secureboot-certificate={cert.name}',
+        f'--secureboot-private-key={key.name}',
+    ])
+
+    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
+
+
 def test_pcr_signing(kernel_initrd, tmpdir):
     if kernel_initrd is None:
         pytest.skip('linux+initrd not found')
diff --git a/src/ukify/ukify.py b/src/ukify/ukify.py
index 3b4ee510fc..a6a4a38040 100755
--- a/src/ukify/ukify.py
+++ b/src/ukify/ukify.py
@@ -43,6 +43,8 @@ import socket
 import subprocess
 import sys
 import tempfile
+import textwrap
+from hashlib import sha256
 from typing import (Any,
                     Callable,
                     IO,
@@ -89,7 +91,7 @@ def guess_efi_arch():
             if int(size) == 32:
                 efi_arch = fallback[0]
 
-    print(f'Host arch {arch!r}, EFI arch {efi_arch!r}')
+    # print(f'Host arch {arch!r}, EFI arch {efi_arch!r}')
     return efi_arch
 
 
@@ -243,13 +245,26 @@ class Uname:
                 print(str(e))
         return None
 
+DEFAULT_SECTIONS_TO_SHOW = {
+        '.linux'    : 'binary',
+        '.initrd'   : 'binary',
+        '.splash'   : 'binary',
+        '.dt'       : 'binary',
+        '.cmdline'  : 'text',
+        '.osrel'    : 'text',
+        '.uname'    : 'text',
+        '.pcrpkey'  : 'text',
+        '.pcrsig'   : 'text',
+        '.sbat'     : 'text',
+}
 
 @dataclasses.dataclass
 class Section:
     name: str
-    content: pathlib.Path
+    content: Optional[pathlib.Path]
     tmpfile: Optional[IO] = None
     measure: bool = False
+    output_mode: Optional[str] = None
 
     @classmethod
     def create(cls, name, contents, **kwargs):
@@ -265,7 +280,7 @@ class Section:
         return cls(name, contents, tmpfile=tmp, **kwargs)
 
     @classmethod
-    def parse_arg(cls, s):
+    def parse_input(cls, s):
         try:
             name, contents, *rest = s.split(':')
         except ValueError as e:
@@ -276,7 +291,19 @@ class Section:
         if contents.startswith('@'):
             contents = pathlib.Path(contents[1:])
 
-        return cls.create(name, contents)
+        sec = cls.create(name, contents)
+        sec.check_name()
+        return sec
+
+    @classmethod
+    def parse_output(cls, s):
+        if not (m := re.match(r'([a-zA-Z0-9_.]+):(text|binary)(?:@(.+))?', s)):
+            raise ValueError(f'Cannot parse section spec: {s!r}')
+
+        name, ttype, out = m.groups()
+        out = pathlib.Path(out) if out else None
+
+        return cls.create(name, out, output_mode=ttype)
 
     def size(self):
         return self.content.stat().st_size
@@ -551,6 +578,7 @@ def pe_add_sections(uki: UKI, output: str):
         if offset + new_section.sizeof() > pe.OPTIONAL_HEADER.SizeOfHeaders:
             raise PEError(f'Not enough header space to add section {section.name}.')
 
+        assert section.content
         data = section.content.read_bytes()
 
         new_section.set_file_offset(offset)
@@ -696,31 +724,32 @@ def make_uki(opts):
     # kernel payload signing
 
     sign_tool = None
-    if opts.signtool == 'sbsign':
-        sign_tool = find_sbsign(opts=opts)
-        sign = sbsign_sign
-        verify_tool = SBVERIFY
-    else:
-        sign_tool = find_pesign(opts=opts)
-        sign = pesign_sign
-        verify_tool = PESIGCHECK
-
     sign_args_present = opts.sb_key or opts.sb_cert_name
+    sign_kernel = opts.sign_kernel
+    sign = None
+    linux = opts.linux
+
+    if sign_args_present:
+        if opts.signtool == 'sbsign':
+            sign_tool = find_sbsign(opts=opts)
+            sign = sbsign_sign
+            verify_tool = SBVERIFY
+        else:
+            sign_tool = find_pesign(opts=opts)
+            sign = pesign_sign
+            verify_tool = PESIGCHECK
 
-    if sign_tool is None and sign_args_present:
-        raise ValueError(f'{opts.signtool}, required for signing, is not installed')
+        if sign_tool is None:
+            raise ValueError(f'{opts.signtool}, required for signing, is not installed')
 
-    sign_kernel = opts.sign_kernel
-    if sign_kernel is None and opts.linux is not None and sign_args_present:
-        # figure out if we should sign the kernel
-        sign_kernel = verify(verify_tool, opts)
-
-    if sign_kernel:
-        linux_signed = tempfile.NamedTemporaryFile(prefix='linux-signed')
-        linux = pathlib.Path(linux_signed.name)
-        sign(sign_tool, opts.linux, linux, opts=opts)
-    else:
-        linux = opts.linux
+        if sign_kernel is None and opts.linux is not None:
+            # figure out if we should sign the kernel
+            sign_kernel = verify(verify_tool, opts)
+
+        if sign_kernel:
+            linux_signed = tempfile.NamedTemporaryFile(prefix='linux-signed')
+            linux = pathlib.Path(linux_signed.name)
+            sign(sign_tool, opts.linux, linux, opts=opts)
 
     if opts.uname is None and opts.linux is not None:
         print('Kernel version not specified, starting autodetection 😖.')
@@ -782,16 +811,17 @@ uki,1,UKI,uki,1,https://www.freedesktop.org/software/systemd/man/systemd-stub.ht
 
     if sign_args_present:
         unsigned = tempfile.NamedTemporaryFile(prefix='uki')
-        output = unsigned.name
+        unsigned_output = unsigned.name
     else:
-        output = opts.output
+        unsigned_output = opts.output
 
-    pe_add_sections(uki, output)
+    pe_add_sections(uki, unsigned_output)
 
     # UKI signing
 
     if sign_args_present:
-        sign(sign_tool, unsigned.name, opts.output, opts=opts)
+        assert sign
+        sign(sign_tool, unsigned_output, opts.output, opts=opts)
 
         # We end up with no executable bits, let's reapply them
         os.umask(umask := os.umask(0))
@@ -913,6 +943,59 @@ def generate_keys(opts):
             pub_key.write_bytes(pub_key_pem)
 
 
+def inspect_section(opts, section):
+    name = section.Name.rstrip(b"\x00").decode()
+
+    # find the config for this section in opts and whether to show it
+    config = opts.sections_by_name.get(name, None)
+    show = (config or
+            opts.all or
+            (name in DEFAULT_SECTIONS_TO_SHOW and not opts.sections))
+    if not show:
+        return name, None
+
+    ttype = config.output_mode if config else DEFAULT_SECTIONS_TO_SHOW.get(name, 'binary')
+
+    data = section.get_data(ignore_padding=True)
+    size = section.Misc_VirtualSize
+    digest = sha256(data).hexdigest()
+
+    struct = {
+        'size' : size,
+        'sha256' : digest,
+    }
+
+    if ttype == 'text':
+        try:
+            struct['text'] = data.decode()
+        except UnicodeDecodeError as e:
+            print(f"Section {name!r} is not valid text: {e}")
+            struct['text'] = '(not valid UTF-8)'
+
+    if config and config.content:
+        assert isinstance(config.content, pathlib.Path)
+        config.content.write_bytes(data)
+
+    if opts.json == 'off':
+        print(f"{name}:\n  size: {size} bytes\n  sha256: {digest}")
+        if ttype == 'text':
+            text = textwrap.indent(struct['text'].rstrip(), ' ' * 4)
+            print(f"  text:\n{text}")
+
+    return name, struct
+
+
+def inspect_sections(opts):
+    indent = 4 if opts.json == 'pretty' else None
+
+    for file in opts.files:
+        pe = pefile.PE(file, fast_load=True)
+        gen = (inspect_section(opts, section) for section in pe.sections)
+        descs = {key:val for (key, val) in gen if val}
+        if opts.json != 'off':
+            json.dump(descs, sys.stdout, indent=indent)
+
+
 @dataclasses.dataclass(frozen=True)
 class ConfigItem:
     @staticmethod
@@ -983,6 +1066,7 @@ class ConfigItem:
     default: Any = None
     version: Optional[str] = None
     choices: Optional[tuple[str, ...]] = None
+    const: Optional[Any] = None
     help: Optional[str] = None
 
     # metadata for config file parsing
@@ -1045,7 +1129,7 @@ class ConfigItem:
         return (section_name, key, value)
 
 
-VERBS = ('build', 'genkey')
+VERBS = ('build', 'genkey', 'inspect')
 
 CONFIG_ITEMS = [
     ConfigItem(
@@ -1160,10 +1244,9 @@ CONFIG_ITEMS = [
         '--section',
         dest = 'sections',
         metavar = 'NAME:TEXT|@PATH',
-        type = Section.parse_arg,
         action = 'append',
         default = [],
-        help = 'additional section as name and contents [NAME section]',
+        help = 'section as name and contents [NAME section] or section to print',
     ),
 
     ConfigItem(
@@ -1277,6 +1360,26 @@ CONFIG_ITEMS = [
         action = argparse.BooleanOptionalAction,
         help = 'print systemd-measure output for the UKI',
     ),
+
+    ConfigItem(
+        '--json',
+        choices = ('pretty', 'short', 'off'),
+        default = 'off',
+        help = 'generate JSON output',
+    ),
+    ConfigItem(
+        '-j',
+        dest='json',
+        action='store_const',
+        const='pretty',
+        help='equivalent to --json=pretty',
+    ),
+
+    ConfigItem(
+        '--all',
+        help = 'print all sections',
+        action = 'store_true',
+    ),
 ]
 
 CONFIGFILE_ITEMS = { item.config_key:item
@@ -1381,7 +1484,15 @@ def finalize_options(opts):
     # Figure out which syntax is being used, one of:
     # ukify verb --arg --arg --arg
     # ukify linux initrd…
-    if len(opts.positional) == 1 and opts.positional[0] in VERBS:
+    if opts.positional[0] == 'inspect':
+        opts.verb = opts.positional[0]
+        opts.files = opts.positional[1:]
+        if not opts.files:
+            raise ValueError('file(s) to inspect must be specified')
+        if len(opts.files) > 1 and opts.json != 'off':
+            # We could allow this in the future, but we need to figure out the right structure
+            raise ValueError('JSON output is not allowed with multiple files')
+    elif len(opts.positional) == 1 and opts.positional[0] in VERBS:
         opts.verb = opts.positional[0]
     elif opts.linux or opts.initrd:
         raise ValueError('--linux/--initrd options cannot be used with positional arguments')
@@ -1439,7 +1550,7 @@ def finalize_options(opts):
             raise ValueError('--secureboot-private-key= and --secureboot-certificate= must be specified together when using --signtool=sbsign')
     else:
         if not bool(opts.sb_cert_name):
-            raise ValueError('--certificate-name must be specified when using --signtool=pesign')
+            raise ValueError('--secureboot-certificate-name must be specified when using --signtool=pesign')
 
     if opts.sign_kernel and not opts.sb_key and not opts.sb_cert_name:
         raise ValueError('--sign-kernel requires either --secureboot-private-key= and --secureboot-certificate= (for sbsign) or --secureboot-certificate-name= (for pesign) to be specified')
@@ -1450,8 +1561,11 @@ def finalize_options(opts):
         suffix = '.efi' if opts.sb_key or opts.sb_cert_name else '.unsigned.efi'
         opts.output = opts.linux.name + suffix
 
-    for section in opts.sections:
-        section.check_name()
+    # Now that we know if we're inputting or outputting, really parse section config
+    f = Section.parse_output if opts.verb == 'inspect' else Section.parse_input
+    opts.sections = [f(s) for s in opts.sections]
+    # A convenience dictionary to make it easy to look up sections
+    opts.sections_by_name = {s.name:s for s in opts.sections}
 
     if opts.summary:
         # TODO: replace pprint() with some fancy formatting.
@@ -1474,6 +1588,8 @@ def main():
     elif opts.verb == 'genkey':
         check_cert_and_keys_nonexistent(opts)
         generate_keys(opts)
+    elif opts.verb == 'inspect':
+        inspect_sections(opts)
     else:
         assert False