#!/usr/bin/python3
#
# pesign_verify_combine_sigs
#
# Helper script for shim-signed
#
# Microsoft currently only returns signed binaries with one signature
# included; if they are signing with more than key/cert, then we will
# get multiple separate signed binaries, one per key/cert.
#
# Check that all our signed shims are signed with an expected key that
# we can remove and re-add; error out otherwise.
#
# Then finally we will add all those signatures to one output binary.
#
# Order of the listed signed shims matters here - list them *in the
# same order* as the signatures we'd like in the final binary. It's
# recommended to do this in the order:
#
# <oldest CA>
# ...
# <newest CA>
#
# as that is most likely to work with older firmware implementations.
# That's most easily achieved by naming the signatures like
# 0001-first-CA.der, 0002-second-CA.der, etc.
#
# Copyright (c) 2026 Steve McIntyre <93sam@debian.org>
# License: BSD-2-Clause
# Modified for Azure Linux with permission

import os
import glob
import re
import sys
import subprocess
import argparse
import shutil
from cryptography import x509
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.serialization import pkcs7

# Allowed certificates - each binary must be signed with a certificate
# from this set (DER-encoded).
SIGN_CERTS = "0*.der"

PESIGN = "pesign"
PESIGCHECK = "pesigcheck"
CERTUTIL = "certutil"


def parse_args():
    parser = argparse.ArgumentParser(description="pesign_verify_combine_sigs")
    parser.add_argument(
        "--efi_arch", "-a", help="EFI architecture for binaries", required=True
    )
    parser.add_argument(
        "--unsigned-path", "-p", help="Path to unsigned shim binaries",
        required=True,
    )
    parser.add_argument("signed", help="signed binaries to verify/merge", nargs="+")
    args = parser.parse_args()
    return args


def setup_nss_db(dbdir: str, certs: list[str]):
    """
    Create an NSS database in dbdir and import all DER certificates.
    """
    os.makedirs(dbdir, exist_ok=True)
    subprocess.check_output(
        [CERTUTIL, "-d", f"sql:{dbdir}", "-N", "--empty-password"],
        stderr=subprocess.STDOUT,
    )
    for cert_file in certs:
        nickname = os.path.splitext(os.path.basename(cert_file))[0]
        subprocess.check_output(
            [
                CERTUTIL, "-d", f"sql:{dbdir}",
                "-A", "-n", nickname,
                "-t", "CT,CT,CT",
                "-i", cert_file,
            ],
            stderr=subprocess.STDOUT,
        )
        print(f"  Imported {cert_file} as '{nickname}'")


def grab_cert_details(check_cert: str) -> (str, str, str):
    """
    Parse a certificate from disk and grab out:
    - a hash of the certificate for comparison later
    - the Subject text
    """
    print(f"Loading details from {check_cert}")
    with open(check_cert, "rb") as inf:
        der_data = inf.read()
    cert = x509.load_der_x509_certificate(der_data)
    subject = "/" + "/".join([x.value for x in cert.subject])
    sha1 = cert.fingerprint(hashes.SHA1()).hex()
    sha256 = cert.fingerprint(hashes.SHA256()).hex()
    print(f" - {subject}")
    print(f" - sha1sum {sha1}")
    print(f" - sha256sum {sha256}")
    return subject, sha1, sha256


def list_signatures(dbdir: str, signed_filename: str):
    """Show signatures on a PE binary using pesign."""
    cmd = [PESIGN, "-n", f"sql:{dbdir}", "-S", "-i", signed_filename]
    output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL, text=True)
    for line in output.splitlines():
        print(line)


def count_signatures(signed_filename: str, dbdir: str) -> int:
    """Count the number of signatures on a PE binary."""
    cmd = [PESIGN, "-n", f"sql:{dbdir}", "-S", "-i", signed_filename]
    output = subprocess.check_output(cmd, text=True, stderr=subprocess.STDOUT)
    # pesign -S prints one block per signature, each delimited by a
    # line of dashes.  Count "certificate address" lines as each
    # signature block contains exactly one.
    count = 0
    for line in output.splitlines():
        if "certificate address" in line:
            count += 1
    return count


def verify_signature(dbdir: str, signed_filename: str):
    """Verify a PE binary signature against certs in the NSS database.

    pesign 116 does not support --verify, so use pesigcheck with DER
    certificate files instead.
    """
    cert_files = sorted(glob.glob("0*.der"))
    if not cert_files:
        raise RuntimeError("No DER certificate files (0*.der) found for verification")
    cmd = [PESIGCHECK, "-i", signed_filename, "-n", "0"]
    for cf in cert_files:
        cmd += ["-c", cf]
    subprocess.check_output(cmd, stderr=subprocess.STDOUT)


def detach_signature(dbdir: str, signed_filename: str, signum: int, outfile: str):
    """
    Detach a numbered signature from a signed PE binary.
    pesign uses 0-based signature numbering.
    """
    cmd = [
        PESIGN,
        "-n", f"sql:{dbdir}",
        "-i", signed_filename,
        "-e", outfile,
        "-u", f"{signum}",
    ]
    subprocess.check_output(cmd, text=True, stderr=subprocess.STDOUT)


def attach_sig(dbdir: str, sigfile: str, infile: str, outfile: str):
    """
    Use pesign to import a detached signature onto a binary.
    pesign requires -i input -o output (not in-place).
    """
    cmd = [
        PESIGN,
        "-n", f"sql:{dbdir}",
        "-i", infile,
        "-o", outfile,
        "-m", sigfile,
    ]
    subprocess.check_output(cmd, text=True, stderr=subprocess.STDOUT)


def certs_in_detached_signature(detached: str) -> list[dict]:
    """
    Extract certificate details from a PKCS7 blob.
    """
    with open(detached, "rb") as inf:
        pkcs7_data = inf.read()

    output = []
    certs = pkcs7.load_der_pkcs7_certificates(pkcs7_data)
    for cert in certs:
        subject = "/" + "/".join([x.value for x in cert.subject])
        sha1 = cert.fingerprint(hashes.SHA1()).hex()
        sha256 = cert.fingerprint(hashes.SHA256()).hex()
        output.append({"sha1": sha1, "sha256": sha256, "subject": subject})

    # We want them in the order CA -> leaf
    output.reverse()
    return output


def pe_hash(filename: str) -> str:
    """
    Get the Authenticode PE hash of a binary using pesign -h.
    This hash covers the binary content but excludes the signature
    table and PE checksum, making it stable across re-signing.
    """
    cmd = [PESIGN, "-h", "-i", filename]
    output = subprocess.check_output(cmd, text=True, stderr=subprocess.DEVNULL)
    # pesign -h outputs something like:
    # "hash 1234abcd..." or just the hex hash
    # Extract the hex hash from the output
    for line in output.splitlines():
        line = line.strip()
        if line:
            # The hash is the first whitespace-delimited token
            return line.split()[0]
    raise RuntimeError(f"pesign -h produced no output for {filename}")


def main():

    args = parse_args()

    # Set up NSS database in CWD
    nssdb = os.path.join(os.getcwd(), "nssdb")
    shutil.rmtree(nssdb, ignore_errors=True)

    print("Loading details of all the expected certificates")
    print("==========")
    known_certs = {}
    cert_files = sorted(glob.glob(SIGN_CERTS))
    for check_cert in cert_files:
        subject, sha1, sha256 = grab_cert_details(check_cert)
        known_certs[subject] = {
            "sha1": sha1,
            "sha256": sha256,
            "filename": check_cert,
        }
    print("")

    print("Setting up NSS database")
    print("==========")
    setup_nss_db(nssdb, cert_files)
    print("")

    print(f"Verifying signatures for arch {args.efi_arch} ...")
    print("==========\n")

    build = "build"
    shutil.rmtree(build, ignore_errors=True)
    os.mkdir(build)

    # Record what we processed, in input (argv) order. The order the signed
    # binaries are passed on the command line is the order their signatures
    # must appear in the final binary (oldest CA first for old-firmware
    # compatibility), so we must not re-derive order from a filename sort.
    processed = []
    matched_cas = set()

    for signed in args.signed:

        print(f"Checking {signed}")
        print("----------\n")

        # Verify that the image is signed and valid
        print("Looking for any valid checksum and signature")
        try:
            list_signatures(nssdb, signed)
            verify_signature(nssdb, signed)
        except subprocess.CalledProcessError as exc:
            print(f"Invalid signature on {signed}: {exc}")
            print(exc.output.decode() if exc.output else "")
            sys.exit(1)

        num = count_signatures(signed, nssdb)
        if num != 1:
            print(f"Only expected 1 signature, but {signed} has {num}!")
            print("Abort")
            sys.exit(1)
        # else
        print(f"{signed} has 1 signature, good!")

        # Now see what signature we have. We'll have to extract the
        # signature table here, then extract the list of certificates
        # included in the 1 signature we have.
        detached_sig = "detached.p7"
        detach_signature(nssdb, signed, 0, detached_sig)
        sig_certs = certs_in_detached_signature(detached_sig)

        matched_filename = None
        matched_sha1 = None

        print("certs attached:")
        for cert in sig_certs:
            print(f'  - {cert["subject"]}')
            print(f'  - sha1 {cert["sha1"]}')
            print(f'  - sha256 {cert["sha256"]}')

        # Compare every embedded certificate against our known CA
        # certificates, matching on the SHA-256 fingerprint. PKCS#7
        # certificate sets are not reliably chain-ordered, so we must not
        # assume sig_certs[0] is the CA; instead we look for any embedded
        # cert whose fingerprint matches a pinned trust anchor.
        for cert in sig_certs:
            for data in known_certs.values():
                if cert["sha256"] == data["sha256"]:
                    print(
                        f'\nsignature chains to a known certificate ({data["filename"]})'
                    )
                    matched_filename = data["filename"]
                    matched_sha1 = data["sha1"]
                    break
            if matched_filename is not None:
                break

        if matched_filename is None:
            print(f"\nERROR: {signed} signature unknown, abort!")
            sys.exit(1)

        # Each input must chain to a *distinct* trust anchor. If two inputs
        # match the same CA (e.g. the same signed binary passed twice), the
        # detached signatures would otherwise overwrite each other and we
        # would silently ship an under-signed binary.
        if matched_filename in matched_cas:
            print(
                f"\nERROR: {signed} matches CA {matched_filename}, which was "
                "already matched by an earlier input; refusing to produce an "
                "under-signed binary. Abort!"
            )
            sys.exit(1)
        matched_cas.add(matched_filename)

        # Move the detached signature to one side, for future use
        new_filename = os.path.join(build, f"detached-{matched_filename}")
        shutil.move(detached_sig, new_filename)

        # And write out the sha1 checksum of the cert for later use
        sha1_filename = os.path.join(build, f"sha1-{matched_filename}")
        with open(sha1_filename, "w") as outf:
            output = ":".join(re.findall("..", matched_sha1))
            outf.write(output)

        # Remember this input's artifacts in argv order for the final combine.
        processed.append(
            {
                "signed": signed,
                "detached": new_filename,
                "sha1_file": sha1_filename,
            }
        )

        # Copy our matching unsigned binary into the ${BUILD} directory.
        unsigned = f"{build}/shim{args.efi_arch}.efi.signed"
        shutil.copy(
            f"{args.unsigned_path}/shim{args.efi_arch}.efi",
            unsigned,
        )

        # Attach the signature to our unsigned binary, so we know that
        # the binary has not been tampered with during the signing
        # process.
        print("Checking the signature applies to our original binary")
        tmp_signed = unsigned + ".tmp"
        attach_sig(nssdb, new_filename, unsigned, tmp_signed)
        # Replace unsigned with the signed version for comparison
        os.replace(tmp_signed, unsigned)
        print("  Signature applies OK")

        # Now compare the PE hash (Authenticode digest) of both binaries.
        # We cannot compare full file checksums because pesign's
        # re-attachment of a detached signature may produce different
        # padding/alignment than the original signing tool.
        print("Comparing PE hashes of signed binaries")
        old_hash = pe_hash(signed)
        print(f"{old_hash}  {signed}")
        new_hash = pe_hash(unsigned)
        print(f"{new_hash}  {unsigned}")

        if old_hash != new_hash:
            print("\nERROR: PE hashes don't match, abort!")
            sys.exit(1)

        print("Binaries match!\n")

    # If we've got this far, then we've checked all the binaries we
    # were given and things look OK. Now we want to build a single
    # output image with all the signatures attached.

    print(f"Building final combined shim for arch {args.efi_arch} ...")
    print("==========")

    # Start fresh from the unsigned binary
    unsigned = f"{build}/shim{args.efi_arch}.efi.signed"
    shutil.copy(
        f"{args.unsigned_path}/shim{args.efi_arch}.efi",
        unsigned,
    )

    for entry in processed:
        sig = entry["detached"]
        print(f"Adding signature {sig}")
        tmp_out = unsigned + ".tmp"
        attach_sig(nssdb, sig, unsigned, tmp_out)
        os.replace(tmp_out, unsigned)

    # Sanity-check that the final binary carries exactly one signature per
    # input. This catches any silent loss of a signature during the combine.
    final_count = count_signatures(unsigned, nssdb)
    if final_count != len(args.signed):
        print(
            f"\nERROR: expected {len(args.signed)} signatures on the combined "
            f"binary, but it has {final_count}. Abort!"
        )
        sys.exit(1)

    # Stick the signature fingerprints together in a file for us to
    # use later in packaging.
    with open(f"{unsigned}-signatures", "w") as outf:
        for entry in processed:
            with open(entry["sha1_file"]) as inf:
                fingerprint = inf.read()
            outf.write(fingerprint + "\n")

    # And finally show the list of signatures
    print(f"Signatures on {unsigned} :")
    list_signatures(nssdb, unsigned)

    # Clean up NSS database
    shutil.rmtree(nssdb, ignore_errors=True)


if __name__ == "__main__":
    main()
