#!/usr/bin/env python3

"""
SPDX-License-Identifier: Apache-2.0
Copyright 2021 Thore Sommer

A simple script to generate a valid mb_refstate for the example policy
"""

import argparse
import json
import logging
import re
import sys

from packaging.version import Version

from keylime.mba.elparsing.tpm2_tools_elparser import parse_binary_bootlog

logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")


def event_to_sha256(event):
    """
    Extracts the sha256 digest from an event
    """
    for digest in event["Digests"]:
        if digest["AlgorithmId"] == "sha256":
            return {"sha256": f"0x{digest['Digest']}"}


def get_s_crtm(events):
    """
    Find the EV_S_CRTM_VERSION.
    """
    for event in events:
        if event["EventType"] == "EV_S_CRTM_VERSION":
            return {"scrtm": event_to_sha256(event)}


def get_platform_firmware(events):
    """
    Get all platform specific files measured with EV_EFI_PLATFORM_FIRMWARE_BLOB events.
    """
    out = []
    for event in events:
        if event["EventType"] != "EV_EFI_PLATFORM_FIRMWARE_BLOB":
            continue
        out.append(event_to_sha256(event))
    return {"platform_firmware": out}


def variabledata_to_signature(data):
    """
    Converts VariableData entry from EV_EFI_VARIABLE_DRIVER_CONFIG to signature data.
    """
    out = []
    for entry in data:
        for key in entry["Keys"]:
            out.append({"SignatureOwner": key["SignatureOwner"], "SignatureData": f"0x{key['SignatureData']}"})
    return out


def get_keys(events):
    """
    Get valid signatures for UEFI Secure Boot PK, KEK, DB and DBX
    """
    out = {"pk": [], "kek": [], "db": [], "dbx": []}
    for event in events:
        if event["EventType"] != "EV_EFI_VARIABLE_DRIVER_CONFIG":
            continue

        event_name = event["Event"]["UnicodeName"].lower()
        if event_name in out:
            data = event["Event"]["VariableData"]
            if data is not None:
                out[event_name] = variabledata_to_signature(data)
    return out


def get_kernel(events, secure_boot):
    """
    Extract digest for Shim, Grub, Linux Kernel and initrd.
    """
    out = []
    for event in events:
        if event["EventType"] != "EV_EFI_BOOT_SERVICES_APPLICATION":
            continue
        out.append(event_to_sha256(event))

    kernel = {"shim_authcode_sha256": out[0]["sha256"], "grub_authcode_sha256": out[1]["sha256"]}
    if secure_boot:
        assert len(out) in [3, 4], "Expected 3 different UEFI applications to be booted (Shim, Grub, Linux)"
        kernel["kernel_authcode_sha256"] = out[2]["sha256"]
    else:
        assert len(out) == 2, "Expected 2 different UEFI applications to be booted (Shim, Grub)"

    for event in events:
        if event["EventType"] != "EV_IPL" or event["PCRIndex"] != 9:
            continue
        if "initrd" in event["Event"]["String"] or "initramfs" in event["Event"]["String"]:
            kernel["initrd_plain_sha256"] = event_to_sha256(event)["sha256"]
            break

    if not secure_boot:
        logging.info("Adding plain sha256 digest of vmlinuz for GRUB to refstate, because SecureBoot is disabled")
        for event in events:
            if event["EventType"] != "EV_IPL" or event["PCRIndex"] != 9:
                continue
            if "vmlinuz" in event["Event"]["String"]:
                kernel["vmlinuz_plain_sha256"] = event_to_sha256(event)["sha256"]
                break

    for event in events:
        if event["EventType"] != "EV_IPL" or event["PCRIndex"] != 8:
            continue
        if "kernel_cmdline" in event["Event"]["String"]:
            kernel["kernel_cmdline"] = re.escape(event["Event"]["String"][len("kernel_cmdline: ") :])
            break

    return {"kernels": [kernel]}


def get_mok(events):
    """
    Extract digest for MokList and MokListX.
    """
    out = {"mokdig": [], "mokxdig": []}
    for event in events:
        if event["EventType"] != "EV_IPL":
            continue

        if event["Event"]["String"] == "MokList":
            out["mokdig"].append(event_to_sha256(event))
        elif event["Event"]["String"] == "MokListX":
            out["mokxdig"].append(event_to_sha256(event))
    return out


def secureboot_enabled(events):
    for event in events:
        if event["EventType"] == "EV_EFI_VARIABLE_DRIVER_CONFIG" and event["Event"].get("UnicodeName") == "SecureBoot":
            return event["Event"]["VariableData"]["Enabled"] == "Yes"

    raise Exception("SecureBoot state could not be determined")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "eventlog_file",
        type=argparse.FileType("rb"),
        default=sys.stdin,
        help="Binary UEFI eventlog (Normally /sys/kernel/security/tpm0/binary_bios_measurements)",
    )
    parser.add_argument(
        "--without-secureboot",
        "-i",
        action="store_true",
        help="Set if you want to create a mb_refstate without SecureBoot (only MeasuredBoot)",
    )
    parser.add_argument(
        "mb_refstate", type=argparse.FileType("w"), help="Output path for the generated measure boot reference state."
    )
    args = parser.parse_args()
    log_bin = args.eventlog_file.read()

    failure, log_data = parse_binary_bootlog(log_bin)
    if failure or not log_data:
        logging.error(
            f"Parsing of binary boot measurements failed with: {list(map(lambda x: x.context, failure.events))}"
        )
        sys.exit(1)

    events = log_data.get("events")
    if not events:
        logging.error("No events found on binary boot measurements log")
        sys.exit(1)

    has_secureboot = secureboot_enabled(events)
    if not has_secureboot and not args.without_secureboot:
        logging.error("Provided eventlog has SecureBoot disabled, but -i flag was not set")
        sys.exit(1)

    if has_secureboot and args.without_secureboot:
        logging.warning(
            "-i was set to create a refstate without SecureBoot, "
            "but he provided eventlog has SecureBoot enabled. Ignoring this flag!"
        )

    mb_refstate = {
        "has_secureboot": has_secureboot,
        "scrtm_and_bios": [
            {
                **get_s_crtm(events),
                **get_platform_firmware(events),
            }
        ],
        **get_keys(events),
        **get_mok(events),
        **get_kernel(events, has_secureboot),
    }
    json.dump(mb_refstate, args.mb_refstate)


if __name__ == "__main__":
    main()
