#!/usr/bin/python3
""" check kvm capabilities via /dev/kvm ioctls """
import os
import sys
import fcntl
import logging
import argparse

def kvm_io(nr: int) -> int:
    return (0xAE << 8) | nr

# ioctls
KVM_GET_API_VERSION = kvm_io(0x00)
KVM_CHECK_EXTENSION = kvm_io(0x03)

# caps
KVM_CAP_VM_TYPES = 235

# bit positions in the KVM_CAP_VM_TYPES bitmap
VM_TYPE_BITS = {
    "default": 0,
    "pkvm": 1,
    "sev": 2,
    "sev-es": 3,
    "sev-snp": 4,
    "tdx": 5,
}

# pylint: disable=broad-exception-caught
def open_kvm():
    try:
        fd = os.open("/dev/kvm", os.O_RDWR)
    except Exception as e:
        print(
            f'ERROR: open /dev/kvm failed: {e}',
            file=sys.stderr,
        )
        return None

    api_version = fcntl.ioctl(fd, KVM_GET_API_VERSION, 0)
    logging.debug('KVM API version  : %d', api_version)
    if api_version != 12:
        print('kvm version mismatch', file=sys.stderr)
        os.close(fd)
        return None

    return fd

def query_vm_types():
    types = []
    fd = open_kvm()
    if fd is not None:
        vm_types_bitmap = fcntl.ioctl(fd, KVM_CHECK_EXTENSION, KVM_CAP_VM_TYPES)
        logging.debug('KVM_CAP_VM_TYPES : 0x%x', vm_types_bitmap)
        for name, bit in VM_TYPE_BITS.items():
            if bool(vm_types_bitmap & (1 << bit)):
                types.append(name)
        os.close(fd)
    return types

def kvm_supports_sev():
    return 'sev' in query_vm_types()

def kvm_supports_sev_es():
    return 'sev-es' in query_vm_types()

def kvm_supports_sev_snp():
    return 'sev-snp' in query_vm_types()

def kvm_supports_tdx():
    return 'tdx' in query_vm_types()

def main():
    parser = argparse.ArgumentParser(description = "check kvm capabilities")
    parser.add_argument('--quiet', dest = 'quiet',
                        action = 'store_true', default = False,
                        help = 'do not print check results')
    checks = parser.add_mutually_exclusive_group(required=False)
    checks.add_argument('--sev', dest = 'check',
                        action = 'store_const', const = 'sev',
                        help = 'check for AMD SEV support')
    checks.add_argument('--sev-es', dest = 'check',
                        action = 'store_const', const = 'sev-es',
                        help = 'check for AMD SEV-ES support')
    checks.add_argument('--sev-snp', '--snp', dest = 'check',
                        action = 'store_const', const = 'sev-snp',
                        help = 'check for AMD SEV-SNP support')
    checks.add_argument('--tdx', dest = 'check',
                        action = 'store_const', const = 'tdx',
                        help = 'check for Intel TDX support')
    options = parser.parse_args()

    if not options.check:
        # no check requested: print caps
        print(f'vm types: {query_vm_types()}')
        return 0

    if options.check == 'sev':
        result = kvm_supports_sev()
    elif options.check == 'sev-es':
        result = kvm_supports_sev_es()
    elif options.check == 'sev-snp':
        result = kvm_supports_sev_snp()
    elif options.check == 'tdx':
        result = kvm_supports_tdx()
    else:
        # keep pylint happy (can not happen)
        result = False

    if result is True:
        if not options.quiet:
            print(f'PASS ({options.check} supported)')
        rc = 0
    else:
        if not options.quiet:
            print(f'FAIL ({options.check} not supported)')
        rc = 1
    return rc

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