Port sign_uefi.sh to Python

Shell scripts are hard to modify and hard to test, so port sign_uefi.sh
to Python. This is a fairly direct port that attempts to keep all the
behavior the same. In particular, there are no hard errors for missing
EFI/kernel files, or for failing to sign one of those files if it does
exist. It might be good to make the script more strict in the future,
but for now try to match the existing behavior.

Nothing actually calls the new script yet.

Also enable `black_check` in `PRESUBMIT.cfg` to enforce formatting.

BRANCH=none
BUG=b:261631233
TEST=make runtests
TEST=cros lint scripts/image_signing/sign_uefi*.py

Change-Id: I4b9b86607cc403779b0504758dd097b0d7237fef
Signed-off-by: Nicholas Bishop <nicholasbishop@google.com>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/4083506
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>

5 files changed, 260 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
index f5bb5005..1a852c34 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,3 +6,4 @@ ID
 target
 .idea
 *.swp
+__pycache__
diff --git a/Makefile b/Makefile
index cb504ff3..295b680f 100644
--- a/Makefile
+++ b/Makefile
@@ -1247,6 +1247,7 @@ runcgpttests: install_for_test
 .PHONY: runtestscripts
 runtestscripts: install_for_test genfuzztestcases
 	${RUNTEST} ${SRC_RUN}/scripts/image_signing/sign_android_unittests.sh
+	${RUNTEST} ${SRC_RUN}/scripts/image_signing/sign_uefi_unittest.py
 	${RUNTEST} ${SRC_RUN}/tests/load_kernel_tests.sh
 	${RUNTEST} ${SRC_RUN}/tests/run_cgpt_tests.sh ${BUILD_RUN}/cgpt/cgpt
 	${RUNTEST} ${SRC_RUN}/tests/run_cgpt_tests.sh ${BUILD_RUN}/cgpt/cgpt -D 358400
diff --git a/PRESUBMIT.cfg b/PRESUBMIT.cfg
index 9f665453..5ffa5f93 100644
--- a/PRESUBMIT.cfg
+++ b/PRESUBMIT.cfg
@@ -1,4 +1,5 @@
 [Hook Overrides]
+black_check: true
 branch_check: true
 cargo_clippy_check: true
 checkpatch_check: true
diff --git a/scripts/image_signing/sign_uefi.py b/scripts/image_signing/sign_uefi.py
new file mode 100755
index 00000000..53ac87c6
--- /dev/null
+++ b/scripts/image_signing/sign_uefi.py
@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+# Copyright 2022 The ChromiumOS Authors
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Sign the UEFI binaries in the target directory.
+
+The target directory can be either the root of ESP or /boot of root filesystem.
+"""
+
+import argparse
+import logging
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+import tempfile
+from typing import List, Optional
+
+
+def ensure_executable_available(name):
+    """Exit non-zero if the given executable isn't in $PATH.
+
+    Args:
+        name: An executable's file name.
+    """
+    if not shutil.which(name):
+        sys.exit(f"Cannot sign UEFI binaries ({name} not found)")
+
+
+def ensure_file_exists(path, message):
+    """Exit non-zero if the given file doesn't exist.
+
+    Args:
+        path: Path to a file.
+        message: Error message that will be printed if the file doesn't exist.
+    """
+    if not path.is_file():
+        sys.exit(f"{message}: {path}")
+
+
+class Signer:
+    """EFI file signer.
+
+    Attributes:
+        temp_dir: Path of a temporary directory used as a workspace.
+        priv_key: Path of the private key.
+        sign_cert: Path of the signing certificate.
+        verify_cert: Path of the certificate used to verify the signature.
+    """
+
+    def __init__(self, temp_dir, priv_key, sign_cert, verify_cert):
+        self.temp_dir = temp_dir
+        self.priv_key = priv_key
+        self.sign_cert = sign_cert
+        self.verify_cert = verify_cert
+
+    def sign_efi_file(self, target):
+        """Sign an EFI binary file, if possible.
+
+        Args:
+            target: Path of the file to sign.
+        """
+        logging.info("signing efi file %s", target)
+
+        # Allow this to fail, as there maybe no current signature.
+        subprocess.run(["sudo", "sbattach", "--remove", target], check=False)
+
+        signed_file = self.temp_dir / target.name
+        try:
+            subprocess.run(
+                [
+                    "sbsign",
+                    "--key",
+                    self.priv_key,
+                    "--cert",
+                    self.sign_cert,
+                    "--output",
+                    signed_file,
+                    target,
+                ],
+                check=True,
+            )
+        except subprocess.CalledProcessError:
+            logging.warning("cannot sign %s", target)
+            return
+
+        subprocess.run(
+            ["sudo", "cp", "--force", signed_file, target], check=True
+        )
+        try:
+            subprocess.run(
+                ["sbverify", "--cert", self.verify_cert, target], check=True
+            )
+        except subprocess.CalledProcessError:
+            sys.exit("Verification failed")
+
+
+def sign_target_dir(target_dir, key_dir, efi_glob):
+    """Sign various EFI files under |target_dir|.
+
+    Args:
+        target_dir: Path of a boot directory. This can be either the
+            root of the ESP or /boot of the root filesystem.
+        key_dir: Path of a directory containing the key and cert files.
+        efi_glob: Glob pattern of EFI files to sign, e.g. "*.efi".
+    """
+    bootloader_dir = target_dir / "efi/boot"
+    syslinux_dir = target_dir / "syslinux"
+    kernel_dir = target_dir
+
+    verify_cert = key_dir / "db/db.pem"
+    ensure_file_exists(verify_cert, "No verification cert")
+
+    sign_cert = key_dir / "db/db.children/db_child.pem"
+    ensure_file_exists(sign_cert, "No signing cert")
+
+    sign_key = key_dir / "db/db.children/db_child.rsa"
+    ensure_file_exists(sign_key, "No signing key")
+
+    with tempfile.TemporaryDirectory() as working_dir:
+        signer = Signer(Path(working_dir), sign_key, sign_cert, verify_cert)
+
+        for efi_file in sorted(bootloader_dir.glob(efi_glob)):
+            if efi_file.is_file():
+                signer.sign_efi_file(efi_file)
+
+        for syslinux_kernel_file in sorted(syslinux_dir.glob("vmlinuz.?")):
+            if syslinux_kernel_file.is_file():
+                signer.sign_efi_file(syslinux_kernel_file)
+
+        kernel_file = (kernel_dir / "vmlinuz").resolve()
+        if kernel_file.is_file():
+            signer.sign_efi_file(kernel_file)
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """Get CLI parser."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "target_dir",
+        type=Path,
+        help="Path of a boot directory, either the root of the ESP or "
+        "/boot of the root filesystem",
+    )
+    parser.add_argument(
+        "key_dir",
+        type=Path,
+        help="Path of a directory containing the key and cert files",
+    )
+    parser.add_argument(
+        "efi_glob", help="Glob pattern of EFI files to sign, e.g. '*.efi'"
+    )
+    return parser
+
+
+def main(argv: Optional[List[str]] = None) -> Optional[int]:
+    """Sign UEFI binaries.
+
+    Args:
+        argv: Command-line arguments.
+    """
+    logging.basicConfig(level=logging.INFO)
+
+    parser = get_parser()
+    opts = parser.parse_args(argv)
+
+    for tool in ("sbattach", "sbsign", "sbverify"):
+        ensure_executable_available(tool)
+
+    sign_target_dir(opts.target_dir, opts.key_dir, opts.efi_glob)
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
diff --git a/scripts/image_signing/sign_uefi_unittest.py b/scripts/image_signing/sign_uefi_unittest.py
new file mode 100755
index 00000000..2cd4342c
--- /dev/null
+++ b/scripts/image_signing/sign_uefi_unittest.py
@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+# Copyright 2022 The ChromiumOS Authors
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Tests for sign_uefi.py.
+
+This is run as part of `make runtests`, or `make runtestscripts` if you
+want something a little faster.
+"""
+
+from pathlib import Path
+import tempfile
+import unittest
+from unittest import mock
+
+import sign_uefi
+
+
+class Test(unittest.TestCase):
+    """Test sign_uefi.py."""
+
+    @mock.patch.object(sign_uefi.Signer, "sign_efi_file")
+    def test_successful_sign(self, mock_sign):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_dir = Path(tmp_dir)
+
+            # Get key paths.
+            key_dir = tmp_dir / "keys"
+            db_dir = key_dir / "db"
+            db_children_dir = db_dir / "db.children"
+
+            # Get target paths.
+            target_dir = tmp_dir / "boot"
+            syslinux_dir = target_dir / "syslinux"
+            efi_boot_dir = target_dir / "efi/boot"
+
+            # Make test dirs.
+            syslinux_dir.mkdir(parents=True)
+            efi_boot_dir.mkdir(parents=True)
+            db_children_dir.mkdir(parents=True)
+
+            # Make key files.
+            (db_dir / "db.pem").touch()
+            (db_children_dir / "db_child.pem").touch()
+            (db_children_dir / "db_child.rsa").touch()
+
+            # Make EFI files.
+            (efi_boot_dir / "bootia32.efi").touch()
+            (efi_boot_dir / "bootx64.efi").touch()
+            (efi_boot_dir / "testia32.efi").touch()
+            (efi_boot_dir / "testx64.efi").touch()
+            (syslinux_dir / "vmlinuz.A").touch()
+            (syslinux_dir / "vmlinuz.B").touch()
+            (target_dir / "vmlinuz-5.10.156").touch()
+            (target_dir / "vmlinuz").symlink_to(target_dir / "vmlinuz-5.10.156")
+
+            # Set an EFI glob that matches only some of the EFI files.
+            efi_glob = "test*.efi"
+
+            # Sign, but with the actual signing mocked out.
+            sign_uefi.sign_target_dir(target_dir, key_dir, efi_glob)
+
+            # Check that the correct list of files got signed.
+            self.assertEqual(
+                mock_sign.call_args_list,
+                [
+                    # The test*.efi files match the glob,
+                    # the boot*.efi files don't.
+                    mock.call(efi_boot_dir / "testia32.efi"),
+                    mock.call(efi_boot_dir / "testx64.efi"),
+                    # Two syslinux kernels.
+                    mock.call(syslinux_dir / "vmlinuz.A"),
+                    mock.call(syslinux_dir / "vmlinuz.B"),
+                    # One kernel in the target dir.
+                    mock.call(target_dir / "vmlinuz-5.10.156"),
+                ],
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()