#!/usr/bin/python3
#
# SPDX-License-Identifier: GPL-2.0-only
# (c) 2023 Gerd Hoffmann
#
""" print and edit efi sigdb files """
import sys
import argparse

from cryptography.hazmat.primitives import serialization

from virt.firmware.efi import guids
from virt.firmware.efi import efivar
from virt.firmware.efi import siglist
from virt.peutils import pesign

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', dest = 'input', type = str,
                        help = 'read efi sigdb FILE', metavar = 'FILE')
    parser.add_argument('-o', '--output', dest = 'output', type = str,
                        help = 'write efi sigdb FILE.', metavar = 'FILE')
    parser.add_argument('--add-cert', dest = 'certs',  action = 'append', nargs = 2,
                        help = 'add x509 cert to sigdb, loaded in pem format ' +
                        'from FILE and with owner GUID, can be specified multiple times',
                        metavar = ('GUID', 'FILE'))
    parser.add_argument('--add-hash', dest = 'hashes',  action = 'append', nargs = 2,
                        help = 'add sha256 hash to sigdb, with owner GUID, '
                        'can be specified multiple times',
                        metavar = ('GUID', 'HASH'))
    parser.add_argument('-x', '--extract', dest = 'extract',
                        action = 'store_true', default = False,
                        help = 'extract certificates')
    parser.add_argument('-p', '--print', dest = 'print',
                        action = 'store_true', default = False,
                        help = 'print sigdb')
    options = parser.parse_args()

    if options.input:
        with open(options.input, "rb") as f:
            sigdb = siglist.EfiSigDB(f.read())
    else:
        sigdb = siglist.EfiSigDB()

    if options.certs:
        for item in options.certs:
            sigdb.add_cert(guids.parse_str(item[0]), item[1])

    if options.hashes:
        for item in options.hashes:
            sigdb.add_hash(guids.parse_str(item[0]),
                           bytes.fromhex(item[1]))

    if options.print and sigdb:
        for slist in sigdb:
            efivar.EfiVarList.print_siglist(slist)
            if options.extract and slist.x509:
                scn = pesign.cert_common_name(slist.x509.subject)
                fn = "".join(x for x in scn if x.isalnum()) + '.pem'
                print(f'    >>> {fn}')
                with open(fn, 'wb') as f:
                    f.write(slist.x509.public_bytes(serialization.Encoding.PEM))

    if options.output:
        with open(options.output, "wb") as f:
            f.write(bytes(sigdb))

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