#!/usr/bin/python
""" print and edit ovmf varstore files """
import sys
import json
import logging
import argparse

from virt.firmware.efi import efivar
from virt.firmware.efi import efijson
from virt.firmware.efi import devpath
from virt.firmware.efi import ucs16

from virt.firmware.varstore import edk2
from virt.firmware.varstore import aws


##################################################################################################
# main

# pylint: disable=too-many-branches,too-many-statements
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-l', '--loglevel', dest = 'loglevel', type = str, default = 'info',
                        help = 'set loglevel to LEVEL', metavar = 'LEVEL')
    parser.add_argument('-i', '--input', dest = 'input', type = str,
                        help = 'read edk2 or aws vars from FILE', metavar = 'FILE')
    parser.add_argument('--extract-certs', dest = 'extract',
                        action = 'store_true', default = False,
                        help = 'extract all certificates')

    pgroup = parser.add_argument_group('Variable options')
    pgroup.add_argument('-d', '--delete', dest = 'delete',
                        action = 'append', type = str,
                        help = 'delete variable VAR, can be specified multiple times',
                        metavar = 'VAR')
    pgroup.add_argument('--set-true', dest = 'set_true',
                        action = 'append', type = str,
                        help = 'set variable VAR to true, can be specified multiple times',
                        metavar = 'VAR')
    pgroup.add_argument('--set-false', dest = 'set_false',
                        action = 'append', type = str,
                        help = 'set variable VAR to false, can be specified multiple times',
                        metavar = 'VAR')
    pgroup.add_argument('--set-json', dest = 'set_json', type = str,
                        help = 'set variables from json dump FILE',
                        metavar = 'FILE')

    pgroup = parser.add_argument_group('Boot configuration')
    pgroup.add_argument('--set-boot-uri', dest = 'set_boot_uri',
                        help = 'set network boot uri to LINK (once, using BootNext)',
                        metavar = 'LINK')
    pgroup.add_argument('--append-boot-filepath', dest = 'append_boot_filepath',
                        action = 'append', type = str,
                        help = 'append boot entry for FILE (permanent, using BootOrder)',
                        metavar = 'FILE')

    pgroup = parser.add_argument_group('shim.efi configuration')
    pgroup.add_argument('--set-shim-debug', dest = 'set_shim_debug',
                        action = 'store_true', default = False,
                        help = 'enable shim.efi debugging (pause for debugger attach)')
    pgroup.add_argument('--set-shim-verbose', dest = 'set_shim_verbose',
                        action = 'store_true', default = False,
                        help = 'enable shim.efi verbose messages')
    pgroup.add_argument('--set-fallback-verbose', dest = 'set_fallback_verbose',
                        action = 'store_true', default = False,
                        help = 'enable fallback.efi verbose messages')
    pgroup.add_argument('--set-fallback-no-reboot', dest = 'set_fallback_no_reboot',
                        action = 'store_true', default = False,
                        help = 'disable rebooting for fallback.efi')

    pgroup = parser.add_argument_group('Secure boot setup options')
    pgroup.add_argument('--set-pk', dest = 'pk',  nargs = 2,
                        help = 'set PK to x509 cert, loaded in pem format ' +
                        'from FILE and with owner GUID',
                        metavar = ('GUID', 'FILE'))
    pgroup.add_argument('--add-kek', dest = 'kek',  action = 'append', nargs = 2,
                        help = 'add x509 cert to KEK, loaded in pem format ' +
                        'from FILE and with owner GUID, can be specified multiple times',
                        metavar = ('GUID', 'FILE'))
    pgroup.add_argument('--add-db', dest = 'db',  action = 'append', nargs = 2,
                        help = 'add x509 cert to db, loaded in pem format ' +
                        'from FILE and with owner GUID, can be specified multiple times',
                        metavar = ('GUID', 'FILE'))
    pgroup.add_argument('--set-dbx', dest = 'dbx',
                        help = 'initialize dbx with update from FILE', metavar = 'FILE')
    pgroup.add_argument('--add-mok', dest = 'mok',  action = 'append', nargs = 2,
                        help = 'add x509 cert to MokList, loaded in pem format ' +
                        'from FILE and with owner GUID, can be specified multiple times',
                        metavar = ('GUID', 'FILE'))
    pgroup.add_argument('--add-db-hash', dest = 'db_hash',  action = 'append', nargs = 2,
                        help = 'add sha256 HASH to db, with owner GUID, ' +
                        'can be specified multiple times',
                        metavar = ('GUID', 'HASH'))
    pgroup.add_argument('--add-mok-hash', dest = 'mok_hash',  action = 'append', nargs = 2,
                        help = 'add sha256 HASH to MokList, with owner GUID, ' +
                        'can be specified multiple times',
                        metavar = ('GUID', 'HASH'))

    pgroup = parser.add_argument_group('Secure boot convinience shortcuts')
    pgroup.add_argument('--enroll-redhat', dest = 'redhat',
                        action = 'store_true', default = False,
                        help = 'enroll default certificates for redhat platform')
    pgroup.add_argument('--enroll-cert', dest = 'enroll_cert',
                        help = 'enroll using specified certificate', metavar = "CERT")
    pgroup.add_argument('--enroll-generate', dest = 'enroll_generate',
                        help = 'enroll using generated cert with given common name', metavar = "CN")
    pgroup.add_argument('--no-microsoft', dest = 'microsoft',
                        action = 'store_false', default = True,
                        help = 'do not add microsoft keys')
    pgroup.add_argument('--distro-keys', dest = 'distro', type = str, action = 'append',
                        help = 'add ca keys for DISTRO', metavar = 'DISTRO')
    pgroup.add_argument('--sb', '--secure-boot', dest = 'secureboot',
                        action = 'store_true', default = False,
                        help = 'enable secure boot mode')

    pgroup = parser.add_argument_group('Print options')
    pgroup.add_argument('-p', '--print', dest = 'print',
                        action = 'store_true', default = False,
                        help = 'print varstore')
    pgroup.add_argument('-v', '--verbose', dest = 'verbose',
                        action = 'store_true', default = False,
                        help = 'print varstore verbosely')
    pgroup.add_argument('-x', '--hexdump', dest = 'hexdump',
                        action = 'store_true', default = False,
                        help = 'print variable hexdumps')

    pgroup = parser.add_argument_group('Output options')
    pgroup.add_argument('-o', '--output', dest = 'output', type = str,
                        help = 'write edk2 or aws vars to FILE, using the same format '
                        'the --input FILE has.', metavar = 'FILE')
    pgroup.add_argument('--output-aws', dest = 'output_aws', type = str,
                        help = 'write aws vars to FILE', metavar = 'FILE')
    pgroup.add_argument('--output-json', dest = 'output_json', type = str,
                        help = 'write json dump to FILE', metavar = 'FILE')
    options = parser.parse_args()

    logging.basicConfig(format = '%(levelname)s: %(message)s',
                        level = getattr(logging, options.loglevel.upper()))

    edk2store = None
    awsstore = None
    varlist = efivar.EfiVarList()

    if options.input:
        if edk2.Edk2VarStore.probe(options.input):
            edk2store = edk2.Edk2VarStore(options.input)
            varlist = edk2store.get_varlist()
        elif edk2.Edk2VarStoreQcow2.probe(options.input):
            edk2store = edk2.Edk2VarStoreQcow2(options.input)
            varlist = edk2store.get_varlist()
        elif aws.AwsVarStore.probe(options.input):
            awsstore = aws.AwsVarStore(options.input)
            varlist = awsstore.get_varlist()
        else:
            logging.error("unknown input file format")
            sys.exit(1)

    if options.extract:
        for (key, item) in varlist.items():
            sigdb = item.sigdb
            if sigdb:
                sigdb.extract_certs(key)

    if options.delete:
        for name in options.delete:
            varlist.delete(name)

    if options.set_true:
        for item in options.set_true:
            varlist.set_bool(item, True)

    if options.set_false:
        for item in options.set_false:
            varlist.set_bool(item, False)

    if options.set_boot_uri:
        items = options.set_boot_uri.split('/')
        title = 'netboot ' + items[-1]
        bpath = devpath.DevicePath.uri(options.set_boot_uri)
        varlist.set_boot_entry(0x99, title, bpath)
        varlist.set_boot_next(0x99)

    if options.append_boot_filepath:
        for item in options.append_boot_filepath:
            strings = item.replace('/', '\\').split(' ')
            filepath = strings[0]
            if len(strings) > 1:
                optdata = bytes(ucs16.from_string(strings[1]))
            else:
                optdata = None
            items = filepath.split('\\')
            title = 'file ' + items[-1]
            bpath = devpath.DevicePath.filepath(filepath)
            index = varlist.add_boot_entry(title, bpath, optdata)
            varlist.append_boot_order(index)

    if options.set_shim_debug:
        varlist.set_uint32('SHIM_DEBUG', 1)

    if options.set_shim_verbose:
        varlist.set_uint32('SHIM_VERBOSE', 1)

    if options.set_fallback_verbose:
        varlist.set_uint32('FALLBACK_VERBOSE', 1)

    if options.set_fallback_no_reboot:
        varlist.set_uint32('FB_NO_REBOOT', 1)

    if options.set_json:
        with open(options.set_json, "r", encoding = 'utf-8') as f:
            l = json.loads(f.read(), object_hook = efijson.efi_decode)
        for (key, item) in l.items():
            logging.info('set variable %s from %s', key, options.set_json)
            varlist[key] = item

    if options.redhat:
        varlist.enroll_platform_redhat()
        if options.microsoft:
            varlist.add_microsoft_keys()

    if options.enroll_cert:
        varlist.enroll_platform_with_cert(options.enroll_cert)
        if options.microsoft:
            varlist.add_microsoft_keys()

    if options.enroll_generate:
        varlist.enroll_platform_generate(options.enroll_generate)
        if options.microsoft:
            varlist.add_microsoft_keys()

    if options.distro:
        for item in options.distro:
            varlist.add_distro_keys(item)

    if options.pk:
        varlist.add_cert('PK', options.pk[0], options.pk[1], True)

    if options.kek:
        for item in options.kek:
            varlist.add_cert('KEK', item[0], item[1])

    if options.db:
        for item in options.db:
            varlist.add_cert('db', item[0], item[1])

    if options.dbx:
        logging.info('reading dbx update from %s', options.dbx)
        with open(options.dbx, 'rb') as f:
            dbxupdate = f.read()
        varlist['dbx'] = efivar.EfiVar(ucs16.from_string('dbx'),
                                       authdata = dbxupdate)

    if options.mok:
        for item in options.mok:
            varlist.add_cert('MokList', item[0], item[1])

    if options.db_hash:
        for item in options.db_hash:
            varlist.add_hash('db', item[0], item[1])

    if options.mok_hash:
        for item in options.mok_hash:
            varlist.add_hash('MokList', item[0], item[1])

    if options.secureboot:
        varlist.enable_secureboot()

    if options.print:
        if options.verbose:
            varlist.print_normal(options.hexdump)
        else:
            varlist.print_compact()

    if options.output:
        if edk2store:
            edk2store.write_varstore(options.output, varlist)
        elif awsstore:
            awsstore.write_varstore(options.output, varlist)
        else:
            logging.error("no input file specified (needed as edk2 varstore template)")
            sys.exit(1)

    if options.output_aws:
        if options.output_aws == "-":
            print(aws.AwsVarStore.base64_varstore(varlist).decode())
        else:
            aws.AwsVarStore.write_varstore(options.output_aws, varlist)

    if options.output_json:
        j = json.dumps(varlist, cls=efijson.EfiJSONEncoder, indent = 4)
        if options.output_json == "-":
            print(j)
        else:
            logging.info('writing json varstore to %s', options.output_json)
            with open(options.output_json, "w", encoding = 'utf-8') as f:
                f.write(j)

    return 0

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