#!/usr/bin/env python3
#
# Copyright 2026 Canonical Ltd.
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License version 2 or later, as published
# by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranties of MERCHANTABILITY,
# SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Integration tests for the FirmwareSecvarUpdater DXE driver.

Tests that booting firmware with only 2011 Microsoft CAs enrolled
results in the 2023 CAs being appended to KEK and db.
"""

import logging
import os
import subprocess
import sys
import tempfile
import time
import unittest
from collections import UserDict
from pathlib import Path

import pexpect
from virt.firmware.varstore import autodetect

from UEFI.Qemu import QemuEfiMachine, QemuEfiFlashSize
from UEFI import Qemu

TEST_TIMEOUT = 120

# CN strings found in the DER encoding of each certificate.
CERT_CNS = {
    'kek_2011':    b'Microsoft Corporation KEK CA 2011',
    'kek_2023':    b'Microsoft Corporation KEK 2K CA 2023',
    'pca_2011':    b'Microsoft Windows Production PCA 2011',
    'uefi_2011':   b'Microsoft Corporation UEFI CA 2011',
    'uefi_2023':   b'Microsoft UEFI CA 2023',
    'oprom_2023':  b'Microsoft Option ROM UEFI CA 2023',
    'win_2023':    b'Windows UEFI CA 2023',
}

ALL_2023_CNS = [
    CERT_CNS['kek_2023'],
    CERT_CNS['uefi_2023'],
    CERT_CNS['oprom_2023'],
    CERT_CNS['win_2023'],
]

# Expected state after a successful update: every cert that should
# be present in KEK and db, keyed by (var_name, cert_key).
EXPECTED_CERTS = [
    ('KEK', 'kek_2011'),
    ('KEK', 'kek_2023'),
    ('db',  'pca_2011'),
    ('db',  'uefi_2011'),
    ('db',  'win_2023'),
    ('db',  'uefi_2023'),
    ('db',  'oprom_2023'),
]


def has_cert_cn(varstore_path, var_name, cn_bytes):
    """Check if the variable's raw data contains the given CN string."""
    vs = autodetect.open_varstore(varstore_path)
    varlist = vs.get_varlist()
    if var_name not in varlist:
        return False
    return cn_bytes in varlist[var_name].data


def strip_2023_certs(varstore_path, output_path):
    """
    Copy a varstore, removing 2023 Microsoft certificates from KEK
    and db to simulate a legacy VM with only 2011 CAs.
    Returns the number of signature list entries removed.
    """
    vs = autodetect.open_varstore(varstore_path)
    varlist = vs.get_varlist()

    removed = 0
    for var_name in ('KEK', 'db'):
        if var_name not in varlist:
            continue
        kept = []
        for siglist in varlist[var_name].sigdb:
            cert_data = siglist[0]['data'] if len(siglist) > 0 else b''
            if any(cn in cert_data for cn in ALL_2023_CNS):
                removed += 1
            else:
                kept.append(siglist)
        new_sigdb = varlist[var_name].sigdb.__class__()
        new_sigdb.extend(kept)
        varlist[var_name].sigdb_set(new_sigdb)

    vs.write_varstore(output_path, varlist)
    return removed


class DpkgArchDict(UserDict):
    data = {}
    dpkg_arch_out = subprocess.check_output(
        ['dpkg-architecture'],
    ).decode().rstrip().splitlines()
    for var in dpkg_arch_out:
        key, val = var.split('=')
        data[key] = val

    def __init__(self):
        pass


class SecvarUpdateTest(unittest.TestCase):
    DpkgArch = DpkgArchDict()

    def setUp(self):
        self.startTime = time.time()
        self.root = Path(os.environ.get('DEB_EDK2_ROOT', '/'))

    def tearDown(self):
        t = time.time() - self.startTime
        sys.stdout.write("%s runtime: %.3fs\n" % (self.id(), t))

    def boot_and_shutdown(self, cmd):
        """Boot QEMU, wait for shell or no-boot prompt, then shut down."""
        logger.debug(f"Running: {' '.join(cmd)}")
        child = pexpect.spawn(' '.join(cmd), encoding='UTF-8')
        child.logfile = sys.stdout
        try:
            i = child.expect(
                [
                    'Press .* or any other key to continue',
                    'BdsDxe: No bootable option or device was found',
                    pexpect.TIMEOUT,
                ],
                timeout=TEST_TIMEOUT,
            )
            if i == 0:
                child.sendline('\x1b')
                child.expect(['Shell> '], timeout=TEST_TIMEOUT)
                child.sendline('reset -s\r')
            elif i == 1:
                child.close()
        except pexpect.EOF:
            pass

    def make_2011_only_vars(self, template_path):
        """
        Create a temp VARS file with 2023 certs stripped.
        Caller must os.unlink() the returned path.
        """
        f = tempfile.NamedTemporaryFile(suffix='.fd', delete=False)
        f.close()

        self.assertTrue(
            has_cert_cn(str(template_path), 'KEK', CERT_CNS['kek_2011']),
            "2011 KEK CA missing from VARS template")

        removed = strip_2023_certs(str(template_path), f.name)
        self.assertGreater(removed, 0,
            "No 2023 certs to strip - template may be outdated")
        self.assertFalse(
            has_cert_cn(f.name, 'KEK', CERT_CNS['kek_2023']),
            "2023 KEK CA still present after stripping")

        return f.name

    def assert_all_certs_present(self, vars_path):
        """Assert all expected 2011 and 2023 certs are in KEK and db."""
        for var_name, cert_key in EXPECTED_CERTS:
            self.assertTrue(
                has_cert_cn(vars_path, var_name, CERT_CNS[cert_key]),
                f"{cert_key} missing from {var_name}")

    def boot_ovmf(self, vars_path):
        """Boot OVMF Q35 with the given VARS file. Returns the QemuCommand."""
        code = self.root / 'usr/share/OVMF/OVMF_CODE_4M.ms.fd'
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF_Q35,
            code_path=code,
            vars_template_path=Path(vars_path),
            flash_size=QemuEfiFlashSize.SIZE_4MB,
        )
        self.boot_and_shutdown(q.command)
        return q

    @unittest.skipUnless(
        DpkgArch['DEB_BUILD_ARCH_BITS'] == '64', "Requires 64-bit QEMU",
    )
    @unittest.skipUnless(
        DpkgArch['DEB_BUILD_ARCH'] == 'amd64', "amd64-only",
    )
    def test_ovmf_secvar_update_appends_2023_certs(self):
        """Boot OVMF with 2011-only VARS, verify 2023 CAs are appended."""
        ms_vars = self.root / 'usr/share/OVMF/OVMF_VARS_4M.ms.fd'
        vars_path = self.make_2011_only_vars(ms_vars)
        try:
            q = self.boot_ovmf(vars_path)
            self.assert_all_certs_present(q.pflash.varfile_path)
        finally:
            os.unlink(vars_path)

    @unittest.skipUnless(
        DpkgArch['DEB_BUILD_ARCH_BITS'] == '64', "Requires 64-bit QEMU",
    )
    @unittest.skipUnless(
        DpkgArch['DEB_BUILD_ARCH'] == 'amd64', "amd64-only",
    )
    def test_ovmf_secvar_update_version_guard(self):
        """After the driver has run, stripping 2023 certs and rebooting
        should NOT re-add them - the version variable prevents re-running."""
        ms_vars = self.root / 'usr/share/OVMF/OVMF_VARS_4M.ms.fd'
        vars_path = self.make_2011_only_vars(ms_vars)
        try:
            # First boot - driver runs, appends 2023 certs
            q1 = self.boot_ovmf(vars_path)
            self.assert_all_certs_present(q1.pflash.varfile_path)

            # Strip 2023 certs again (simulating user removal), but keep
            # the version variable that the driver wrote
            stripped_path = q1.pflash.varfile_path + '.stripped'
            removed = strip_2023_certs(q1.pflash.varfile_path, stripped_path)
            self.assertGreater(removed, 0)

            # Second boot - driver should skip due to version guard
            q2 = self.boot_ovmf(stripped_path)

            # 2023 certs should NOT be back
            for cn in ALL_2023_CNS:
                self.assertFalse(
                    has_cert_cn(q2.pflash.varfile_path, 'KEK', cn) or
                    has_cert_cn(q2.pflash.varfile_path, 'db', cn),
                    f"Driver re-ran despite version guard: {cn} was re-added")

            os.unlink(stripped_path)
        finally:
            os.unlink(vars_path)

    @unittest.skipUnless(
        DpkgArch['DEB_BUILD_ARCH_BITS'] == '64', "Requires 64-bit QEMU",
    )
    @unittest.skipUnless(
        DpkgArch['DEB_BUILD_ARCH'] == 'arm64', "arm64-only",
    )
    def test_aavmf_secvar_update_appends_2023_certs(self):
        """Boot AAVMF with 2011-only VARS, verify 2023 CAs are appended."""
        code = self.root / 'usr/share/AAVMF/AAVMF_CODE.ms.fd'
        ms_vars = self.root / 'usr/share/AAVMF/AAVMF_VARS.ms.fd'

        vars_path = self.make_2011_only_vars(ms_vars)
        try:
            q = Qemu.QemuCommand(
                QemuEfiMachine.AAVMF,
                code_path=code,
                vars_template_path=Path(vars_path),
            )
            self.boot_and_shutdown(q.command)
            self.assert_all_certs_present(q.pflash.varfile_path)
        finally:
            os.unlink(vars_path)


if __name__ == '__main__':
    logger = logging.getLogger(__name__)
    logging.basicConfig(level=logging.DEBUG)
    unittest.main(verbosity=2)
