#!/usr/bin/env python3

#
# Copyright 2017-2020 Mellanox Technologies. All Rights Reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# * Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of Mellanox nor the names of its contributors
#   may be used to endorse or promote products derived from this software
#   without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS
# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
# OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#

"""
Wrapper for UEFI capsule generation.
"""

import hashlib
import os
import struct
import subprocess
import sys
import tempfile
import binascii
from optparse import OptionParser, OptionGroup

# Program version
ProgVersion = "3.3.0"

# GUID, Version and Least Supported Version
FmpGUID = "39342586-4e0e-4833-b4ba-1256b0ffb471"
FmpFwVer = "0x00000001"
FmpLsv = "0x00000001"
FmpImageHeaderVer = "0x00000002"

# Default Keys
MicrosoftCA = "keys/MicCorUEFCA2011_2011-06-27.cer"

# Boot image version and lowest supported version
BootImageDefaultVer = 0x0000001
BootImageDefaultLsv = 0x0000000

# Boot image update policy
BootUpdatePrimaryPartition = 0x01
BootUpdateBackupPartition = 0x02

BootUpdateDefaultPolicy = BootUpdatePrimaryPartition | BootUpdateBackupPartition

# Boot image flags
LockBootPartition = 0x01
SwapBootPartition = 0x02
EnableWatchdog = 0x04
EnableWatchdogSwap = SwapBootPartition | EnableWatchdog

# Key format
KeyFormatPem = 1
KeyFormatCer = 2
KeyFormatBin = 3
KeyFormatRnd = 4
KeyFormatHash = 5

# Key action
KeyActionNone = 0x00
KeyActionAdd = 0x01
KeyActionDelete = 0x02
KeyActionUpdate = 0x03

# Key flags
EnableCryptoDatapath = 0x80
EnableCerberus = 0x40
EnableHmac = 0x20
DisableSecureBoot = 0x10
EnableRma = 0x08
EnableNicSecureBoot = 0x04
EnableArmSecureBoot = 0x02
EnableNvProduction = 0x01
# Extra flags
DisableCryptoDatapath = 0x100
EnableFlashEncryption = 0x200
ProgramManufacturerId = 0x400
ProgramFuseDisable = 0x800
RevokeArmKey = 0x1000
ProgramFeatureDisable = 0x2000
ProgramArmRomPatchFuse = 0x4000
EnableEmmcRpmb = 0x8000
ProgramNicRomPatchFuse = 0x10000

#
# Contents: (command-line, description, image type ID)
# ID=0 is used externally and will not be added into the capsule payload.
#
ImageList = (
    # Arguments passed to GenerateCapsule.py.
    ("engine",            "openssl engine name",                       0),
    ("digest",            "Digest algorithm to use when signing. The default is sha256", 0),
    ("signer-key",        "openssl signer private key (pem file or PKCS11 URI)", 0),
    ("signer-cert",       "openssl signer public cert file (.pem)",    0),
    ("other-public-cert", "openssl other public cert file (.pem)",     0),
    ("pfx-file",          "signtool PFX cert file (.pfx)",             0),
    ("embedded-driver",   "driver to handle this capsule (.efi)",      0),

    # Images included in capsule payload.
    # The images will be sorted by TypeID in the generated capsule.
    # Some holes in the TypeID are reserved for future expansion.
    ("nic-sbkey",         "NIC secure boot key (.bin)",                 4),
    ("arm-sbkey",         "Arm secure boot key (.bin)",                 5),
    ("chip-id",           "The chip identifier",                        6),
    ("db-key",            "UEFI secure boot DB key (.cer)",            10),
    ("db-bin",            "UEFI secure boot DB binary (.bin/.efi)",    10),
    ("db-hash",           "UEFI secure boot DB hash",                  10),
    ("db-delete-key",     "UEFI secure boot DB key (.cer) to delete",  10),
    ("db-delete-hash",    "UEFI secure boot DB hash to delete",        10),
    ("kek-key",           "UEFI secure boot KEK key (.cer)",           11),
    ("kek-delete-key",    "UEFI secure boot KEK key (.cer) to delete", 11),
    ("pk-key",            "UEFI secure boot PK key (.cer)",            12),
    ("uefi-passwd",       "UEFI password",                             13),
    ("arm-oem-sbkey",     "Arm OEM secure boot keys (.bin)",           14),
    ("psc-fuses",         "PSC keystore and external fuses (.bin)",    15),
    ("mek",               "Mellanox Encryption Key (.bin)",            16),
    ("gcm-iv",            "32-bit Input Vector for AES_GCM engine",    17),
    ("manufacturer-id",   "Manufacturer identifier (.bin)",            18),
    ("feature-disable",   "64-bit Feature Disable",                    19),
    ("fuse-disable",      "32-bit Fuse Disable",                       20),
    ("arm-patch-fuses",   "Arm BL1 ROM patch fuses (.bin)",            21),
    ("ddr-config-idx",    "DDR configuration index in decimal",        22),
    ("nic-msv",           "NIC Minimal Security Version in decimal",   23),
    ("cfg-data",          "Configure Eeprom Mfg/Mfg-Extension data",   24),
    ("nic-patch-fuses",   "NIC Boot ROM patch fuses (.bin)",           25),
    ("tls-ca",            "TLS CA cert (.der) for UEFI HTTPS boot",    26),
    ("ddr-screen-rev",    "DDR screening revision in decimal",         27),
    ("dbx-key",           "UEFI secure boot DBX key (.cer)",            3),
    ("dbx-bin",           "UEFI secure boot DBX binary(.bin/.efi)",     3),
    ("dbx-hash",          "UEFI secure boot DBX hash",                  3),
    ("dbx-delete-key",    "UEFI secure boot DBX key (.cer) to delete",  3),
    ("dbx-delete-hash",   "UEFI secure boot DBX hash to delete",        3),
    ("bootimg",           "boot partition image (.bfb)",              100),
    ("firmware",          "NIC firmware image",                       101)
)

#
# The capsule payload contains a list of images. Each image starts with a
# header like below followed by the actual data.
#
class PayloadHeaderClass:
    # ///
    # /// Payload Header (32 bytes)
    # ///
    # typedef struct {
    #   UINT32   Type;             // Image Type ID
    #   UINT32   Length;           // Payload Length
    #   UINT8    PrivateData[24];  // PrivateData
    # };
    #
    _StructFormat = '<LL24s'
    _StructSize = struct.calcsize(_StructFormat)

    def __init__(self):
        self.Type = 0
        self.Length = 0
        self.PrivateData = b''
        self.Payload = b''

    def Encode(self):
        Header = struct.pack(
                            self._StructFormat,
                            self.Type,
                            self.Length,
                            self.PrivateData)
        return Header + self.Payload

#
# Function to initialize the boot image private data. For now, use default
# settings. This might be extended in the future to pass the configuration
# parameters.
#
def SetBootImagePrivateData(Version, Lsv, UpdatePolicy, Flags):
    # ///
    # /// BlueField eMMC Device Image Private Data (24 bytes).
    # ///
    # typedef struct {
    #   UINT32    Version;               // Image Version.
    #   UINT8     UpdatePolicy;          // Device Update Policy.
    #   UINT8     Flags;                 // Flags Used to Update the Device.
    #   BOOLEAN   SystemResetRequired;   // Set to TRUE if a reset is required.
    #   UINT8     Reserved0;             // Reserved byte for future use.
    #   UINT32    LowestSupportedVerion; // Image Last Supported Version.
    #   UINT8     Reserved1[12];         // Reserved bytes for future use.
    # } PMI_EMMC_BOOT_IMAGE_HEADER;
    #
    SystemResetRequired = 0
    Reserved0 = 0
    Reserved1 = b''

    BootImagePrivateData = struct.pack('<L4BL12s',
                                        Version,
                                        UpdatePolicy,
                                        Flags,
                                        SystemResetRequired,
                                        Reserved0,
                                        Lsv,
                                        Reserved1)
    return BootImagePrivateData

#
# The Boot Image Update Configuration
#
class BootImageCfgClass:
    def __init__(self, Version, Lsv, UpdatePolicy, Flags):
        self.Version = Version
        self.Lsv = Lsv
        self.UpdatePolicy = UpdatePolicy
        self.Flags = Flags

#
# The Secure Boot Configuration
#
class SecureBootCfgClass:
    def __init__(self, Flags):
        self.KeyFlags = Flags
        self.ArmKeyUpdate = False
        self.ForceConfig = False

    def HasSecureBootCfg(self):
        if self.KeyFlags:
            return True
        return False

    def HasKeyUpdate(self):
        return self.ArmKeyUpdate

    def SkipUserConfirmation(self):
        if self.ForceConfig:
            return True
        return False

#
# Function to initialize the UEFI secure boot key private data.
# For now, use default settings. This might be extended in the future to pass
# the configuration parameters.
#
def SetSecureBootKeyPrivateData(KeyFormat, KeyActions, KeyFlags, Skip):
    # ///
    # /// Secure Boot Key Private Data (24 bytes).
    # ///
    # typedef struct {
    #   UINT8     Format;               // Key Format.
    #   UINT8     Action;               // Action needed to
    #   UINT8     Flags;                // Flags Used to Manage Keys.
    #   BOOLEAN   SystemResetRequired;  // Set to TRUE if a reset is required.
    #   BOOLEAN   SkipUserConfirmation; // Skip user confirmation, if needed.
    #   UINT32    ExtraFlags;           // Extra flags, if needed.
    #   UINT8     Reserved[15];         // Reserved for future use.
    # } PMI_KEY_PARAMATERS;
    #
    Format = KeyFormat
    Action = KeyActions
    Flags = KeyFlags & 0xff
    SystemResetRequired = 1
    SkipUserConfirmation = Skip
    # Set the first two bytes of extra flags.
    ExtraFlags = (KeyFlags >> 8) & 0xffff
    Reserved = b''

    SecureBootKeyPrivateData = struct.pack('<5BL15s',
                                            Format,
                                            Action,
                                            Flags,
                                            SystemResetRequired,
                                            SkipUserConfirmation,
                                            ExtraFlags,
                                            Reserved)
    return SecureBootKeyPrivateData

# Map command-line option to (description, image ID)
ImageTable = dict([(x, [y, z, None]) for (x, y, z) in ImageList])

# Current path of the script.
CurPath = os.path.dirname(os.path.abspath(__file__))


def SortFunc(e):
    return ImageTable[e[0]][1]


def CreateTmpFile():
    TmpFileHandle, TmpFileName = tempfile.mkstemp()
    os.close(TmpFileHandle)
    return TmpFileName

def CreateCapsulePayload(OutputFileName, Images, SecureBootCfg, BootImageCfg, Options):
    OutputFile = open(OutputFileName, "w+b")
    PayloadHeader = PayloadHeaderClass()

    HasKeyConfig = False
    HasDummyKey = False
    DummyKeyFileName = None
    Skip = SecureBootCfg.SkipUserConfirmation()
    # If secure boot flags are set and there are no keys to be programmed
    # then create a dummy one to handle the fuse configuration.
    if SecureBootCfg.HasSecureBootCfg():
        for (ImageKey, FileName) in Images:
            if ImageKey in [ "arm-sbkey", "nic-sbkey", "chip-id", "mek" ]:
                HasKeyConfig = True
                break
        if not HasKeyConfig:
            DummyKeyFileName = CreateTmpFile()
            Images.append(("arm-sbkey", DummyKeyFileName))
            HasDummyKey = True

    Images.sort(key=SortFunc)

    for (ImageKey, FileName) in Images:
        ImageType = ImageTable[ImageKey][1]
        ImageTable[ImageKey][2] = FileName
        if ImageType == 0:
            continue
        if ImageKey == "uefi-passwd":
            #
            # Password is provided as plain string.
            #
            PassLen = len(FileName)
            if PassLen < 1 or PassLen > 64:
                print("Error: Supplied UEFI password must be >= 1 and <= 64 characters")
                exit(1)
            WidePasswd = ""
            for i in range(0, PassLen):
                WidePasswd = WidePasswd + FileName[i] + "\x00"

            # Import cryptography modules and others here instead; because most of
            # these modules are needed for UEFI password setting only. This script
            # shouldn't complain about missing modules if executed to accomplish
            # other functions.
            from cryptography.hazmat.primitives import hashes
            from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
            from cryptography.hazmat.backends import default_backend
            # Generate the Salt
            SaltBytes = os.urandom(32)
            # Implement PBKDF2, this creates our Key Derivation Function
            # Note: "iterations" below must match iterations defined in
            #     "PasswordLib.h" in Host UEFI
            Kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=SaltBytes,
                iterations=1000,
                backend=default_backend()
                )
            # Derive the key from our WidePasswd
            if sys.version_info[0] < 3:
                Key = Kdf.derive(WidePasswd)
            else:
                Key = Kdf.derive(bytes(WidePasswd, 'utf-8'))
            # We will send the Key and the Salt to the host
            if sys.version_info[0] < 3:
                KeyPlusSalt = '{}{}'.format(binascii.hexlify(Key), binascii.hexlify(SaltBytes))
                PayloadHeader.Payload = KeyPlusSalt
            else:
                KeyPlusSalt = struct.pack('<32s32s', Key, SaltBytes)
                PayloadHeader.Payload = binascii.hexlify(KeyPlusSalt)
            PayloadHeader.Length = len(PayloadHeader.Payload)
        elif ImageKey == "ddr-config-idx":
            #
            # DDR config index is provided in decimal.
            #
            KeyFlags = SecureBootCfg.KeyFlags
            PayloadHeader.Payload = struct.pack("<I", int(FileName))
            PayloadHeader.Length = len(PayloadHeader.Payload)
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                   KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        elif ImageKey == "nic-msv":
            #
            # NIC minimal version is provided in decimal.
            #
            KeyFlags = SecureBootCfg.KeyFlags
            PayloadHeader.Payload = struct.pack("<I", int(FileName))
            PayloadHeader.Length = len(PayloadHeader.Payload)
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                   KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        elif ImageKey == "gcm-iv":
            #
            # GCM is provided in hex.
            #
            KeyFlags = EnableFlashEncryption | SecureBootCfg.KeyFlags
            PayloadHeader.Payload = struct.pack("<I", int(FileName, 16))
            PayloadHeader.Length = len(PayloadHeader.Payload)
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                   KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        elif ImageKey == "chip-id":
            KeyFlags = EnableRma | SecureBootCfg.KeyFlags
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                  KeyFormatBin, KeyActionNone, KeyFlags, Skip)
            ChipId = ""
            for i in range(0, len(FileName)):
                ChipId = ChipId + FileName[i]
            PayloadHeader.Payload = ChipId.decode("hex")
            PayloadHeader.Length = len(PayloadHeader.Payload)
        elif ImageKey == "feature-disable":
            KeyFlags = ProgramFeatureDisable | SecureBootCfg.KeyFlags
            PayloadHeader.Payload = struct.pack("<Q", int(FileName, 16))
            PayloadHeader.Length = len(PayloadHeader.Payload)
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                   KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        elif ImageKey == "fuse-disable":
            KeyFlags = ProgramFuseDisable | SecureBootCfg.KeyFlags
            PayloadHeader.Payload = struct.pack("<I", int(FileName, 16))
            PayloadHeader.Length = len(PayloadHeader.Payload)
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                   KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        elif ImageKey == "ddr-screen-rev":
            #
            # DDR screening revision is provided in decimal.
            #
            KeyFlags = SecureBootCfg.KeyFlags
            PayloadHeader.Payload = struct.pack("<I", int(FileName))
            PayloadHeader.Length = len(PayloadHeader.Payload)
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                   KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        else:
            InFile = open(FileName, "rb")
            PayloadHeader.Length = os.path.getsize(FileName)
            PayloadHeader.Payload = InFile.read(PayloadHeader.Length)
            InFile.close()
        PayloadHeader.Type = ImageType
        if ImageKey in [ "pk-key", "kek-key", "db-key", "db-bin", "db-hash", "dbx-key", "dbx-bin", "dbx-hash" ]:
            #
            # By default, delete existing key (if any) and install new key.
            #
            KeyAction = KeyActionDelete | KeyActionAdd
            if ImageKey == "pk-key" and os.path.getsize(FileName) == 0:
                #
                # If the PK is empty, then delete the existing PK key
                # (if any) and disable UEFI secure boot.
                #
                KeyAction = KeyActionDelete
            elif ImageKey in [ "db-key", "db-bin", "db-hash" ]:
                if Options.db_append:
                    KeyAction = KeyActionAdd
            elif ImageKey in [ "dbx-key", "dbx-bin", "dbx-hash" ]:
                if Options.dbx_append:
                    KeyAction = KeyActionAdd
            KeyFormat = KeyFormatCer
            if ImageKey in [ "dbx-bin", "db-bin" ]:
                KeyFormat = KeyFormatBin
            elif ImageKey in [ "dbx-hash", "db-hash" ]:
                KeyFormat = KeyFormatHash
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                             KeyFormat, KeyAction, 0, False)
        if ImageKey in [ "kek-delete-key", "db-delete-key", "db-delete-hash", "dbx-delete-key", "dbx-delete-hash" ]:
            KeyAction = KeyActionDelete
            KeyFormat = KeyFormatCer
            if ImageKey in [ "db-delete-hash", "dbx-delete-hash" ]:
                KeyFormat = KeyFormatHash
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                             KeyFormat, KeyAction, 0, False)
        if ImageKey == "arm-sbkey":
            KeyFlags = SecureBootCfg.KeyFlags
            KeyAction = KeyActionNone
            KeyFormat = KeyFormatBin
            if HasDummyKey:
                KeyFormat = KeyFormatRnd
            if not HasDummyKey:
                KeyFlags = KeyFlags | EnableArmSecureBoot
                KeyAction = KeyActionAdd
            if SecureBootCfg.HasKeyUpdate():
                KeyAction = KeyActionUpdate
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                         KeyFormat, KeyAction, KeyFlags, Skip)
        if ImageKey == "arm-oem-sbkey":
            KeyFlags = EnableArmSecureBoot
            KeyAction = KeyActionAdd
            KeyFormat = KeyFormatBin
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                      KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        if ImageKey == "nic-sbkey":
            KeyFlags = EnableNicSecureBoot | SecureBootCfg.KeyFlags
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                   KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        if ImageKey == "mek":
            KeyFlags = EnableFlashEncryption | SecureBootCfg.KeyFlags
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                                   KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        if ImageKey == "psc-fuses":
            KeyFlags = EnableNvProduction | SecureBootCfg.KeyFlags
            KeyAction = KeyActionAdd
            KeyFormat = KeyFormatBin
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                    KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        if ImageKey == "manufacturer-id":
            KeyFlags = ProgramManufacturerId | SecureBootCfg.KeyFlags
            KeyAction = KeyActionAdd
            KeyFormat = KeyFormatBin
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                    KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        if ImageKey == "arm-patch-fuses":
            KeyFlags = ProgramArmRomPatchFuse | SecureBootCfg.KeyFlags
            KeyAction = KeyActionAdd
            KeyFormat = KeyFormatBin
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                    KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        if ImageKey == "nic-patch-fuses":
            KeyFlags = ProgramNicRomPatchFuse | SecureBootCfg.KeyFlags
            KeyAction = KeyActionAdd
            KeyFormat = KeyFormatBin
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                    KeyFormatBin, KeyActionAdd, KeyFlags, Skip)
        if ImageKey == "tls-ca":
            if os.path.getsize(FileName) == 0:
                # If the CA is empty, then delete the existing CAs
                KeyAction = KeyActionDelete
            else:
                if Options.tls_ca_append:
                    KeyAction = KeyActionAdd
                else:
                    KeyAction = KeyActionDelete | KeyActionAdd
            KeyFormat = KeyFormatCer
            PayloadHeader.PrivateData = SetSecureBootKeyPrivateData(
                    KeyFormat, KeyAction, 0, False)
        if ImageKey == "bootimg":
            BootImgVer = BootImageCfg.Version
            BootImgLsv = BootImageCfg.Lsv
            UpdatePolicy = BootImageCfg.UpdatePolicy
            UpdateFlags = BootImageCfg.Flags
            PayloadHeader.PrivateData = SetBootImagePrivateData(
                             BootImgVer, BootImgLsv, UpdatePolicy, UpdateFlags)
        Payload = PayloadHeader.Encode()
        OutputFile.write(Payload)

    if DummyKeyFileName is not None:
        os.unlink(DummyKeyFileName)

    OutputFile.close()


def ParseImageOpt(option, opt_str, value, Parser):
    """Handle one of the --<image-type> options."""
    Parser.values.images.append((opt_str.lstrip("-"), value))

def ParseVerStr(ver_str, opt_name, Parser):
    try:
        major, minor, patch, build = ver_str.split(".")
    except ValueError as e:
        Parser.error("{}: bad format, {}".format(opt_name, e.message))

    try:
        ipatch = int(patch)
    except ValueError:
        print("warn: patch {} is not int, defaulting to 0".format(patch))
        ipatch = 0

    try:
        ver = (((int(major)&0xFF) << 24) |
               ((int(minor)&0xFF) << 16) |
               ((ipatch&0xFF) << 8) |
                (int(build)&0xFF))
    except ValueError:
        Parser.error(opt_name + ": major, minor, build must be integers")

    return ver

def main(argv):
    #
    # Set up our option Parser.
    #
    Parser = OptionParser(
            usage="\n  %prog [ options ] <OUTFILE>\n"
                  "  %prog [ -h | --help ]",
            description="%prog creates a UEFI capsule.")

    Parser.set_defaults(images=[])

    iog = OptionGroup(Parser, "Image Options", '')
    for fn in sorted(ImageTable.keys()):
        iog.add_option(
                "--" + fn, type="string",
                action="callback", callback=ParseImageOpt, metavar="FILE",
                help=ImageTable[fn][0])

    Parser.add_option_group(iog)

    Parser.add_option (
                "--version", action = "store_true", dest="version",
                help = "Display " + os.path.basename(__file__) + " version information.")
    Parser.add_option(
                "--img-ver", action="store",
                type="int", dest="img_ver",
                help="boot image version")
    Parser.add_option(
                "--img-ver-str", action="store",
                type="str", dest="img_ver_str",
                help="Major.minor.patch.build version string, alternative to "
                     "--img-ver. Converts the string to a 32-bit integer of the "
                     "form MMmmppbb, where each character is 4 bits. "
                     "Only one of --img-ver-str and --img-ver may be used "
                     "at a time.")
    Parser.add_option(
                "--img-lsv", action="store",
                type="int", dest="img_lsv",
                help="boot image lowest supported version")
    Parser.add_option(
                "--img-lsv-str", action="store", type="str",
                dest="img_lsv_str",
                help="boot image lowest supported version, using the string "
                     "format shown in --img-ver-str. May not be used at the "
                     "same time as --img-lsv.")
    Parser.add_option(
                "--img-policy", action="store",
                type="int", dest="img_policy",
                help="boot image update policy: "
                     "(1) update primary partition -"
                     "(2) update secondary partition -"
                     "(3) update both boot partitions (default)")
    Parser.add_option(
                "--boot-swap", action="store_const",
                dest="boot_swap", const=EnableWatchdogSwap,
                help="flag to enable watchdog-swap")
    Parser.add_option(
                "--boot-lock", action="store_const",
                dest="boot_lock", const=LockBootPartition,
                help="flag to lock bootpartitions")

    Parser.add_option(
                "--crypto", action="store_const",
                dest="crypto", const=EnableCryptoDatapath,
                help="flag to enable crypto")
    Parser.add_option(
                "--crypto-disable", action="store_const",
                dest="crypto_disable", const=DisableCryptoDatapath,
                help="flag to Disable crypto")
    Parser.add_option(
                "--cerberus", action="store_const",
                dest="cerberus", const=EnableCerberus,
                help="flag to enable attestation flow")
    Parser.add_option(
                "--hmac", action="store_const",
                dest="hmac", const=EnableHmac,
                help="flag to enable HMAC authentication")
    Parser.add_option(
                "--hard-nonsecure", action="store_const",
                dest="hard_nonsecure", const=DisableSecureBoot,
                help="flag to put the chip into explicit non secure state")
    Parser.add_option(
                "--update-arm-sbkey", action="store_const",
                dest="update_arm_sbkey", const=KeyActionUpdate,
                help="flag to update Arm secure boot key")
    Parser.add_option(
                "--revoke-arm-sbkey", action="store_const",
                dest="revoke_arm_sbkey", const=RevokeArmKey,
                help="flag to revoke Arm secure boot key")
    Parser.add_option(
                "--enable-rpmb", action="store_const",
                dest="enable_rpmb", const=EnableEmmcRpmb,
                help="flag to enable eMMC RPMB partition")
    Parser.add_option(
                "--force-sb-config", action="store_const",
                dest="force_sb_config", const=True,
                help="flag to skip user confirmation and force operation")
    Parser.add_option(
                "--db-append", action="store_const",
                dest="db_append", const=True,
                help="Append key files provided with --db-{key,bin,hash} to UEFI DB")
    Parser.add_option(
                "--dbx-append", action="store_const",
                dest="dbx_append", const=True,
                help="Append key files provided with --dbx-{key,bin,hash} to UEFI DBX")
    Parser.add_option(
                "--tls-ca-append", action="store_const",
                dest="tls_ca_append", const=True,
                help="Append CA files provided with --tls-ca to UEFI TLS CA")

    # Parse the command line and execute.
    (opt, args) = Parser.parse_args()

    if len(args) != 1:
        if opt.version:
            print(os.path.basename(__file__) + " version " + ProgVersion)
            exit(0)
        else:
            Parser.error("output file not specified")

    if len(args) == 1 and opt.version:
        print("option \'--version\' ignored")

    if opt.img_ver and opt.img_ver_str:
        Parser.error("cannot specify both --img-ver and --img-ver-str")

    if opt.img_lsv and opt.img_lsv_str:
        Parser.error("cannot specify both --img-lsv and --img-lsv-str")

    PayloadTmpFileName = CreateTmpFile()

    # Initialize secure boot flags for extra fuse configuration
    SecureBootKeyFlags = 0
    if opt.crypto:
        SecureBootKeyFlags = SecureBootKeyFlags | int(opt.crypto)
    if opt.crypto_disable:
        SecureBootKeyFlags = SecureBootKeyFlags | int(opt.crypto_disable)
    if opt.cerberus:
        SecureBootKeyFlags = SecureBootKeyFlags | int(opt.cerberus)
    if opt.hmac:
        SecureBootKeyFlags = SecureBootKeyFlags | int(opt.hmac)
    if opt.hard_nonsecure:
        SecureBootKeyFlags = SecureBootKeyFlags | int(opt.hard_nonsecure)
    if opt.revoke_arm_sbkey:
        SecureBootKeyFlags = SecureBootKeyFlags | int(opt.revoke_arm_sbkey)
    if opt.enable_rpmb:
        SecureBootKeyFlags = SecureBootKeyFlags | int(opt.enable_rpmb)

    SbCfg = SecureBootCfgClass(SecureBootKeyFlags)

    if opt.update_arm_sbkey:
        SbCfg.ArmKeyUpdate = True

    if opt.force_sb_config:
        SbCfg.ForceConfig = True

    ImgVer = BootImageDefaultVer
    ImgLsv = BootImageDefaultLsv
    UpdatePolicy = BootUpdateDefaultPolicy
    ImgFlags = 0

    if opt.img_lsv:
        ImgLsv = opt.img_lsv
    elif opt.img_lsv_str:
        ImgLsv = ParseVerStr(opt.img_lsv_str, "--img-lsv-str", Parser)

    if opt.img_ver:
        ImgVer = opt.img_ver
    elif opt.img_ver_str:
        ImgVer = ParseVerStr(opt.img_ver_str, "--img-ver-str", Parser)

    if opt.img_policy:
        UpdatePolicy = opt.img_policy
    if opt.boot_swap:
        ImgFlags = ImgFlags | int(opt.boot_swap)
    if opt.boot_lock:
        ImgFlags = ImgFlags | int(opt.boot_lock)

    if ImgLsv > ImgVer:
        Parser.error("boot image version invalid")

    BootImgCfg = BootImageCfgClass(ImgVer, ImgLsv, UpdatePolicy, ImgFlags)

    # Generate the capsule payload
    CreateCapsulePayload(PayloadTmpFileName, opt.images, SbCfg, BootImgCfg, opt)

    Cmd = [
        "python3",
        CurPath + "/Capsule/GenerateCapsule.py",
        "-o", args[0],
        "-e",
        "--capversion", FmpImageHeaderVer,
        "--guid", FmpGUID,
        "--fw-version", FmpFwVer,
        "--lsv", FmpLsv]

    Pfx_File = ImageTable["pfx-file"][2]
    SignerPrivateKey = ImageTable["signer-key"][2]
    SignerPublicCert = ImageTable["signer-cert"][2]
    OtherPublicCert = ImageTable["other-public-cert"][2]
    if Pfx_File:
        if SignerPrivateKey or SignerPublicCert or OtherPublicCert:
            print("pfx-file can't mix with signer-cert or other-public-cert\n")
        else:
            Cmd = Cmd + ["--pfx-file", Pfx_File]
    elif SignerPrivateKey is not None and SignerPublicCert is not None:
        Cmd = Cmd + [
            "--signer-private-key", SignerPrivateKey,
            "--signer-public-cert", SignerPublicCert]
        if OtherPublicCert is not None:
            Cmd = Cmd + ["--other-public-cert", OtherPublicCert]

    if ImageTable["embedded-driver"][2] is not None:
        Cmd = Cmd + ["--embedded-driver", ImageTable["embedded-driver"][2]]

    if ImageTable["engine"][2] is not None:
        Cmd = Cmd + ["--engine", ImageTable["engine"][2]]

    if ImageTable["digest"][2] is not None:
        Cmd = Cmd + ["--digest", ImageTable["digest"][2]]

    Cmd = Cmd + [PayloadTmpFileName]

    os.environ['PYTHONPATH'] = CurPath
    subprocess.call(Cmd)

    os.unlink(PayloadTmpFileName)


if __name__ == "__main__":
    main(sys.argv)
