#!/usr/bin/env python3

import argparse
import base64
import os
import struct
import subprocess
import sys
import tempfile

# python3 -m pip install pycryptodomex
from Cryptodome.Hash import HMAC, SHA1, SHA256
from Cryptodome.PublicKey import RSA
from Cryptodome.Signature import pkcs1_15

_CONFIG = {'DEBUG': False}
MAX_BIN_SIZE = 110 * 1024

def debug(msg):
    """
    Outputs the msg string to stdout if DEBUG is enabled (via -d)
    """
    if _CONFIG['DEBUG']:
        sys.stderr.write(str(msg) + '\n')

class ImageFile:
    """
    Signed binary image
    """
    def __init__(self, filename, max_size_kb):
        self._filename = filename
        self._bytes_written = 0
        if self._filename is None:
           self._of = sys.stdout
        else:
           self._of = open(self._filename, "wb")
        self._max_size_kb = max_size_kb
        self._bytes = bytearray()

        debug("%8s %20s: [%6s] %s" % ('OFFSET', 'TYPE', 'SIZE', 'DESCRIPTION'))
        debug("")

    def append(self, data):
        """
        Appends a blob of binary data to the image
        """
        self._bytes.extend(data)

    def append_file(self, source_file):
        """
        Appends the binary contents of source_file to the current image. If
        source_file is None then a base64 encoded blob is read from stdin.
        """
        if source_file is None:
           b64 = ""
           for l in sys.stdin.readlines():
              b64 += l
           file_bytes = base64.b64decode(b64)
        else:
           file_bytes = bytearray(open(source_file, 'rb').read())
        size = len(file_bytes)
        debug("%08x %20s: [%6d] %s" % (self.pos(), 'FILE', size, source_file))
        self.append(file_bytes)

    def append_keynum(self, keynum):
        """
        Appends a given key number as a 32-bit LE integer.
        """
        if (keynum < 0 or keynum > 4) and keynum != 16:
            raise Exception("Bad key number %d" % keynum)
        debug("%08x %20s: [%6d] %d" % (self.pos(), "KEYNUM", 4, keynum))
        self.append(struct.pack('<i', keynum))

    def append_version(self, version):
        """
        Appends a version number, 0-32 to avoid firmware rollback, a Raspberry Pi
        with OTP bit n set will not execute a firmware without bit n set.
        """
        if version < 0 or version > 32:
            raise Exception("Bad version number %d must be between 0-32" % version)
        debug("%08x %20s: [%6d] %d" % (self.pos(), "VERSION", 4, version))
        self.append(struct.pack('<i', version))

    def append_length(self):
        """
        Appends the current length to the image as a 32-bit LE integer
        """
        length = len(self._bytes)
        debug("%08x %20s: [%6d] %d" % (self.pos(), "LEN", 4, length))
        self.append(struct.pack('<i', length))

    def append_public_key(self, pem_file):
        """
        Converts an RSA public key into the format expected by the ROM
        and appends it to the image.

        If a private key is passed then only the public key part is extracted.
        """
        arr = bytearray()
        key = RSA.importKey(open(pem_file, 'r').read())

        if key.size_in_bits() != 2048:
            raise Exception("RSA key size must be 2048")

        # Export N and E in little endian format
        arr.extend(key.n.to_bytes(256, byteorder='little'))
        arr.extend(key.e.to_bytes(8, byteorder='little'))
        debug("%08x %20s: [%6d] %s" % (self.pos(), 'RSA', len(arr), pem_file))
        self.append(arr)

        h = SHA256.new()
        h.update(key.n.to_bytes(256, byteorder='little'))
        h.update(key.e.to_bytes(8, byteorder='little'))
        d = h.hexdigest()
        pub_str = ""
        for i in range(int(len(d)/8)):
           pub_str += "0x%s%s%s%s, " % (d[i*8+6:i*8+8], d[i*8+4:i*8+6], d[i*8+2:i*8+4], d[i*8+0:i*8+2])
        debug("Public key SHA256(N,e) = %s" % pub_str)

    def append_rsa_signature_pkcs11(self, hsm_wrapper):
        temp = tempfile.NamedTemporaryFile(delete=False)
        temp.write(self._bytes)
        temp.close() # close and flush before spawning PKCS#11 wrapper

        res = subprocess.run([hsm_wrapper, "-a", "rsa2048-sha256", temp.name], capture_output=True)
        debug(res.stderr)
        if res.returncode != 0:
            os.unlink(temp.name)
            raise Exception(f"HSM wrapper failed with exit code {res.returncode}: {res.stderr.decode()}")
        signature = res.stdout.decode()
        os.unlink(temp.name)
        self.append(bytearray.fromhex(signature))
        debug("PKCS11 %08x %20s: [%6d] signature %s" % (self.pos(), 'RSA2048 - SHA256', len(signature), signature))

    def append_rsa_signature(self, digest_alg, private_pem):
        """
        Append a RSA 2048 signature of the SHA256 of the data so far
        """
        key = RSA.importKey(open(private_pem, 'r').read())
        if key.size_in_bits() != 2048:
            raise Exception("RSA key size must be 2048")

        if digest_alg == 'sha256':
            digest = SHA256.new(self._bytes)
        elif digest_alg == 'sha1':
            digest = SHA1.new(self._bytes)

        signature = pkcs1_15.new(key).sign(digest)
        self.append(signature)
        debug("%08x %20s: [%6d] digest %s signature %s" % (self.pos(), 'RSA2048 - SHA256', len(signature), digest.hexdigest(), signature.hex()))

    def append_digest(self, digest_alg, hmac_keyfile):
        """
        Appends the hash/digest to the image
        """
        if hmac_keyfile is not None:
            hmac_key = str(open(hmac_keyfile, 'r').read()).strip()
            expected_keylen = 40
            if len(hmac_key) != expected_keylen:
                raise Exception("Bad key length %d expected %d" % (len(hmac_key), expected_keylen))

        if digest_alg == 'hmac-sha1':
            h = HMAC.new(base64.b16decode(hmac_key, True), self._bytes, digestmod=SHA1)
        else:
            raise Exception("Digest not supported %s" % (digest_alg))

        debug("%08x %20s: [%6d] %s" % (self.pos(), digest_alg, len(h.digest()), h.hexdigest()))
        self.append(h.digest())

    def pos(self):
        return len(self._bytes)

    def write(self):
        if len(self._bytes) > self._max_size_kb:
            raise Exception("Signed binary size %d is too large. Max size %d" % (len(self._bytes), MAX_BIN_SIZE))
        debug("Image size %d" % len(self._bytes))
        if self._filename is None:
           self._of.buffer.write(base64.b64encode(self._bytes))
        else:
           self._of.write(self._bytes)

    def close(self):
        self._of.close()

def create_2711_image(output, bootcode, private_key=None, private_keynum=0, hmac=None, hsm_wrapper=None):
    """
    Create a 2711 C0 secure-boot compatible seconds stage signed binary.
    """
    image = ImageFile(output, MAX_BIN_SIZE)
    image.append_file(bootcode)
    image.append_length()
    image.append_keynum(private_keynum)
    if hsm_wrapper:
        image.append_rsa_signature_pkcs11(hsm_wrapper)
    else:
        image.append_rsa_signature('sha1', private_key)
    image.append_digest('hmac-sha1', hmac)
    image.write()
    image.close()

def create_2712_image(output, bootcode, private_version=0, public_key=None, private_key=None, private_keynum=0, hsm_wrapper=None):
    """
    Create a prototype 2712 signed bootloader. The HMAC is removed and the
    full public key is appended.
    """
    image = ImageFile(output, MAX_BIN_SIZE)
    image.append_file(bootcode)
    image.append_length()
    image.append_keynum(private_keynum)
    image.append_version(private_version)
    if hsm_wrapper is not None:
        debug(f"Call HSM wrapper {hsm_wrapper}")
        image.append_rsa_signature_pkcs11(hsm_wrapper)
        image.append_public_key(public_key)
    else:
        image.append_rsa_signature('sha256', private_key)
        image.append_public_key(private_key)
    image.write()
    image.close()

def main():
    help_text = """
    Signs a second stage bootloader image.

    Examples:
    
        Customer counter-signed:
          * Exactly the same as Raspberry Pi signing but the input is the Raspberry Pi signed bootcode.bin
          * The key number will probably always be 16 to indicate a customer signing

        rpi-sign-bootcode --debug -c 2712 -i bootcode.bin.sign2 -o bootcode.bin -k customer.pem

        PKCS#1 v1.5 - HSM wrapper:
          * hsm-wrapper takes a single argument which is a temporary filename containing the data to sign
          * hsm-wrapper outputs the PKCS#1 v1.5 signature in hex format
          * hsm-wrapper must return a non-zero exit code if signing fails
          * hsm-wrapper requires the -a rsa2048-sha256 parameter to specify the algorithm
          * There is no facility to pass the private key or custom HSM arguments - the caller should generate a dedicated wrapper script
          * The public key in PEM format MUST be specified with the -p option

        rpi-sign-bootcode --debug -c 2712 -i bootcode.bin.sign2 -o bootcode.bin -p public.pem -H hsm-wrapper
    """
    parser = argparse.ArgumentParser(help_text)
    parser.add_argument('-o', '--output', required=False, help='Output filename. If not specified, the signed image is written to stdout in base64 format')
    parser.add_argument('-c', '--chip', required=True, type=int, help='Chip number')
    parser.add_argument('-i', '--input', required=False, help='Path of the unsigned bootcode.bin file OR RPi signed bootcode file to be signed with the customer key. If NULL, the binary is read from stdin in base64 format')
    parser.add_argument('-m', '--hmac', required=False, help='Path of the HMAC key file')
    parser.add_argument('-k', '--private-key', dest='private_key', required=False, default=None, help='Path of RSA private key (PEM format)')
    parser.add_argument('-p', '--public-key', dest='public_key', required=False, default=None, help='Path of RSA public key (PEM format)')
    parser.add_argument('-n', '--private-keynum', dest='private_keynum', required=False, default=0, type=int, help='ROM key index for RPi signing stage')
    parser.add_argument('-H', '--hsm-wrapper', default=None, required=False, help='Filename of HSM wrapper script which generates a PKCSv1.1 signature as hex')
    parser.add_argument('-d', '--debug', action='store_true')
    parser.add_argument('-v', '--private-version', dest='private_version', required=False, default=0, type=int, help='Version of firmware, stops firmware rollback, only valid 0-31')

    args = parser.parse_args()
    _CONFIG['DEBUG'] = args.debug
    if args.chip == 2711:
        if args.hmac is None:
            raise Exception("HMAC key requried for 2711")
        create_2711_image(args.output, args.input, private_key=args.private_key, private_keynum=args.private_keynum, hmac=args.hmac, hsm_wrapper=args.hsm_wrapper)
    elif args.chip == 2712:
        create_2712_image(args.output, args.input, private_version=args.private_version, public_key=args.public_key, private_key=args.private_key, private_keynum=args.private_keynum, hsm_wrapper=args.hsm_wrapper)

if __name__ == '__main__':
    main()