#!/usr/bin/python3
#
# SPDX-License-Identifier: GPL-2.0-only
# (c) 2026 Gerd Hoffmann
#
# pylint: disable=too-many-arguments
# pylint: disable=too-many-positional-arguments
# pylint: disable=too-many-public-methods
""" manage qemu firmware config files """
import os
import json
import subprocess

class QemuFirmware:
    """ manage one qemu firmware config file """

    def __init__(self, data = None, filename = None):
        self.data = data
        if self.data is None:
            self.data = {}
        if filename:
            with open(filename, 'r', encoding = 'utf-8') as f:
                self.data = json.loads(f.read())


    def cmdline_flash(self, name, prop, desc, readonly = False, snapshot = False):
        filename = desc.get('filename')
        fformat = desc.get('format')
        drive = f'if=none,id={name},format={fformat},file={filename}'
        if readonly:
            drive += ',readonly=on'
        if snapshot:
            drive += ',snapshot=on'
        cmdline = [ '-drive', drive,
                    '-machine', f'{prop}={name}' ]
        return cmdline

    def cmdline_igvm(self, igvm):
        cmdline = []
        cmdline += [ '-object', f'igvm-cfg,id=igvm0,file={igvm}' ]
        cmdline += [ '-machine', 'igvm-cfg=igvm0' ]
        return cmdline

    def cmdline(self, varsfile = None):
        cmdline = []
        mapping = self.data.get('mapping')
        device = mapping.get('device')
        mode = mapping.get('mode')
        fw_code = mapping.get('executable')
        if device == 'flash' and mode == 'split':
            fw_vars = mapping.get('nvram-template').copy()
            if varsfile:
                fw_vars['filename'] = varsfile
                snapshot = False
            else:
                snapshot = True
            cmdline += self.cmdline_flash('code', 'pflash0', fw_code, readonly = True)
            cmdline += self.cmdline_flash('vars', 'pflash1', fw_vars, snapshot = snapshot)
        elif device == 'flash' and mode == 'stateless':
            cmdline += self.cmdline_flash('code', 'pflash0', fw_code, readonly = True)
        elif device == 'memory':
            cmdline += [ '-bios', mapping.get('filename') ]
        elif device == 'igvm':
            cmdline += self.cmdline_igvm(mapping.get('filename'))
        else:
            raise RuntimeError(f'not supported: device={device} mode={mode}')

        if self.requires_smm():
            cmdline += [ '-machine', 'smm=on' ]
            cmdline += [ '-global', 'driver=cfi.pflash01,property=secure,value=on']
        if self.requires_qemu_vars():
            if self.arch() == 'x86_64':
                device = 'uefi-vars-x64'
            else:
                device = 'uefi-vars-sysbus'
            if varsfile:
                device += f',jsonfile={varsfile}'
            cmdline += [ '-device', device ]

        return cmdline

    def vars_template(self):
        mapping = self.data.get('mapping')
        if mapping.get('device') == 'flash':
            if mapping.get('mode') != 'split':
                return None
            return mapping.get('nvram-template').get('filename')
        if mapping.get('device') == 'memory':
            uefivars = mapping.get('uefi-vars')
            if uefivars:
                return uefivars.get('template')
        return None

    def arch(self):
        targets = self.data.get('targets')[0]
        return targets.get('architecture')

    def is_native(self):
        hostarch = os.uname().machine
        return hostarch == self.arch()

    def machine(self):
        targets = self.data.get('targets')[0]

        if len(targets.get('machines')) == 0:
            if self.arch() == 'x86_64':
                return 'q35'
            return 'virt'

        machine = targets.get('machines')[0]
        if 'pc-i440fx' in machine:
            return 'pc'
        if 'pc-q35' in machine:
            return 'q35'
        if 'virt' in machine:
            return 'virt'
        return machine


    def has_feature(self, name):
        features = self.data.get('features')
        if features is None:
            return False
        return name in features

    def has_sb_enabled(self):
        return self.has_feature('enrolled-keys')

    def has_cc_sev(self):
        return self.has_feature('amd-sev')

    def has_cc_sev_es(self):
        return self.has_feature('amd-sev-es')

    def has_cc_snp(self):
        return self.has_feature('amd-sev-snp')

    def has_cc_tdx(self):
        return self.has_feature('intel-tdx')

    def has_cc(self):
        return self.has_cc_sev() or self.has_cc_sev_es() or self.has_cc_snp() or self.has_cc_tdx()

    def requires_smm(self):
        return self.has_feature('requires-smm')

    def requires_qemu_vars(self):
        return self.has_feature('host-uefi-vars')

    def requires_igvm(self):
        device = self.data.get('mapping').get('device')
        return device == 'igvm'

    def check_qemu(self, binary):
        if self.requires_igvm():
            objects = subprocess.run([ binary, '-object', 'help'],
                                     capture_output = True,
                                     check = True).stdout.decode().strip()
            if not 'igvm-cfg' in objects:
                return 'qemu does not support igvm'

        if self.requires_qemu_vars():
            devices = subprocess.run([ binary, '-device', 'help'],
                                     capture_output = True,
                                     check = True).stdout.decode().strip()
            if not 'uefi-vars' in devices:
                return 'qemu does not support uefi-vars'

        return None

    def quirk_code_filename(self):
        mapping = self.data.get('mapping')
        fw_code = mapping.get('executable')
        if fw_code:
            return fw_code.get('filename')
        return mapping.get('filename')

    # should be available as feature
    def quirk_has_hashes(self):
        return 'amdsev' in self.quirk_code_filename()


    @staticmethod
    def inspect_ovmf(filename):
        inspect = subprocess.run([ 'virt-fw-dump', '--modules', '-i', filename ],
                                 capture_output = True,
                                 check = True).stdout.decode().strip()
        features = []
        if 'name=SmmAccess2Dxe' in inspect:
            features.append('requires-smm')
        if 'name=SecureBootConfigDxe' in inspect:
            features.append('secure-boot')
        if 'name=AmdSevDxe' in inspect:
            features.append('amd-sev')
            features.append('amd-sev-es')
            features.append('amd-sev-snp')
        if 'name=TdxDxe' in inspect:
            features.append('intel-tdx')
        if 'name=VirtMmCommunication' in inspect:
            features.append('host-uefi-vars')
        return features

    @staticmethod
    def inspect_igvm(filename):
        inspect = subprocess.run([ 'igvm-inspect', filename ],
                                 capture_output = True,
                                 check = True).stdout.decode().strip()
        features = []
        if 'platform: SEV_SNP' in inspect:
            features.append('amd-sev-snp')
        return features


    @staticmethod
    def create_flash_one(filename):
        fw_format = 'raw'
        if filename.endswith('.qcow2'):
            fw_format = 'qcow2'
        entry = {
            'filename' : filename,
            'format'   : fw_format
        }
        return entry

    @staticmethod
    def create_targets(arch, machine):
        machines = []
        if machine:
            machines.append(machine)
        target = {
            'architecture': arch,
            'machines': machines,
        }
        return [ target, ]

    @staticmethod
    def create_flash_split(fw_code, fw_vars, arch, machine):
        mapping = {
            'device'         : 'flash',
            'mode'           : 'split',
            'executable'     : QemuFirmware.create_flash_one(fw_code),
            'nvram-template' : QemuFirmware.create_flash_one(fw_vars),
        }
        targets = QemuFirmware.create_targets(arch, machine)
        features = QemuFirmware.inspect_ovmf(fw_code)
        return QemuFirmware(data = { 'mapping' : mapping,
                                     'targets' : targets,
                                     'features' : features })

    @staticmethod
    def create_flash_stateless(fw_code, arch, machine):
        mapping = {
            'device'         : 'flash',
            'mode'           : 'stateless',
            'executable'     : QemuFirmware.create_flash_one(fw_code),
        }
        targets = QemuFirmware.create_targets(arch, machine)
        features = QemuFirmware.inspect_ovmf(fw_code)
        return QemuFirmware(data = { 'mapping' : mapping,
                                     'targets' : targets,
                                     'features' : features })

    @staticmethod
    def create_memory(fw_code, arch, machine):
        mapping = {
            'device'   : 'memory',
            'filename' : fw_code,
        }
        targets = QemuFirmware.create_targets(arch, machine)
        features = QemuFirmware.inspect_ovmf(fw_code)
        return QemuFirmware(data = { 'mapping' : mapping,
                                     'targets' : targets,
                                     'features' : features })

    @staticmethod
    def create_igvm(fw_igvm, arch, machine):
        mapping = {
            'device'   : 'igvm',
            'filename' : fw_igvm,
        }
        targets = QemuFirmware.create_targets(arch, machine)
        features = QemuFirmware.inspect_igvm(fw_igvm)
        return QemuFirmware(data = { 'mapping'  : mapping,
                                     'targets'  : targets,
                                     'features' : features })

    @staticmethod
    def from_env():
        fw_json = os.environ.get('EDK2_JSON')
        if fw_json:
            return QemuFirmware(filename = fw_json)
        fw_code = os.environ.get('EDK2_CODE')
        fw_vars = os.environ.get('EDK2_VARS')
        fw_igvm = os.environ.get('EDK2_IGVM')
        fw_arch = os.environ.get('EDK2_ARCH', os.uname().machine)
        fw_machine = os.environ.get('EDK2_MACHINE')
        if fw_code and fw_vars:
            return QemuFirmware.create_flash_split(fw_code, fw_vars, fw_arch, fw_machine)
        if fw_code:
            return QemuFirmware.create_memory(fw_code, fw_arch, fw_machine)
        if fw_igvm:
            return QemuFirmware.create_igvm(fw_igvm, fw_arch, fw_machine)
        raise RuntimeError('no firmware config found (in EDK2_* env vars)')
