#!/usr/bin/python3

# {{{ CDDL HEADER
#
# This file and its contents are supplied under the terms of the
# Common Development and Distribution License ("CDDL"), version 1.0.
# You may only use this file in accordance with the terms of version
# 1.0 of the CDDL.
#
# A full copy of the text of the CDDL should have accompanied this
# source. A copy of the CDDL is also available via the Internet at
# http://www.illumos.org/license/CDDL.
# }}}

# Portions Copyright 2019 Joyent, Inc. (from uts/i86pc/os/hma.c)
# Copyright 2021 OmniOS Community Edition (OmniOSce) Association.

import re, os, subprocess, struct, sys, ucred
from ctypes import *

from pprint import pprint, pformat

CPUID_DEV                           = "/dev/cpu/self/cpuid"

CPU_INTEL                           = 1
CPU_AMD                             = 2

CPUID_INTC_ECX_VMX                  = 0x00000020

MSR_IA32_FEAT_CTRL                  = 0x3a
IA32_FEAT_CTRL_LOCK                 = 1 << 0
IA32_FEAT_CTRL_VMX_EN               = 1 << 2

MSR_IA32_VMX_BASIC                  = 0x480
IA32_VMX_BASIC_INS_OUTS             = 1 << 54
IA32_VMX_BASIC_TRUE_CTRLS           = 1 << 55

MSR_IA32_VMX_TRUE_PROCBASED_CTLS    = 0x48e
MSR_IA32_VMX_PROCBASED_CTLS         = 0x482
IA32_VMX_PROCBASED_2ND_CTLS         = 1 << 31

MSR_IA32_VMX_PROCBASED2_CTLS        = 0x48b
IA32_VMX_PROCBASED2_EPT             = 1 << 1
IA32_VMX_PROCBASED2_VPID            = 1 << 5

MSR_IA32_VMX_EPT_VPID_CAP           = 0x48c
IA32_VMX_EPT_VPID_INVEPT            = 1 << 20
IA32_VMX_EPT_VPID_INVEPT_SINGLE     = 1 << 25
IA32_VMX_EPT_VPID_INVEPT_ALL        = 1 << 26

CPUID_AMD_ECX_SVM                   = 0x00000004
MSR_AMD_VM_CR                       = 0xc0010114
AMD_VM_CR_SVMDIS                    = 1 << 4
CPUID_AMD_EDX_NESTED_PAGING         = 0x000000001
CPUID_AMD_EDX_NRIPS                 = 0x000000008

def note(msg):
    print("... {}".format(msg))

def rdmsr(msr):
    ret = subprocess.run(['/usr/sbin/rdmsr', hex(msr)], capture_output=True)
    m = rdmsr.p.match(ret.stdout.decode())
    if m:
        return int(m.group(2), 16)
    return 0
rdmsr.p = re.compile(r'(0x[0-9a-f]+): (.*)')

class _cpuid(Structure):
    _fields_ = [
        ('eax', c_uint32),
        ('ebx', c_uint32),
        ('ecx', c_uint32),
        ('edx', c_uint32),
    ]

def cpuid(leaf=0):
    with open(CPUID_DEV, 'rb') as fd:
        fd.seek(leaf)
        x = _cpuid()
        if fd.readinto(x) != sizeof(x):
            return None
    return x

def cpu_vendor():
        x = cpuid()
        s = struct.pack("III", x.ebx, x.edx, x.ecx).decode()
        note(f"CPU vendor string '{s}'")
        if s == 'GenuineIntel':
            return CPU_INTEL
        return CPU_AMD

def vmx_ctl_one_setting(val, flag):
    return val & (flag << 32) != 0

def vmx_check():

    x = cpuid(1)
    if not x.ecx & CPUID_INTC_ECX_VMX:
        note("CPU does not support VMX")
        return False
    note("CPU supports VMX")

    # Has the BIOS set the feature-control lock bit without VMX enabled?
    i = rdmsr(MSR_IA32_FEAT_CTRL)
    if i & IA32_FEAT_CTRL_LOCK and not i & IA32_FEAT_CTRL_VMX_EN:
        note("VMX support not enabled in BIOS (essential)")
        return False
    note("VMX support is enabled in BIOS")

    # The basic INS/OUTS functionality is cited as a necessary prereq
    i = rdmsr(MSR_IA32_VMX_BASIC)
    if not i & IA32_VMX_BASIC_INS_OUTS:
        note("VMX does not support INS/OUTS (essential)")
        return False
    note("VMX supports INS/OUTS")

    # Bit 55 in the VMX_BASIC MSR determines how VMX control information
    # can be queried.
    query_true_ctl = i & IA32_VMX_BASIC_TRUE_CTRLS
    ctlmsr = (MSR_IA32_VMX_TRUE_PROCBASED_CTLS if query_true_ctl
        else MSR_IA32_VMX_PROCBASED_CTLS)

    # Check for EPT and VPID support
    i = rdmsr(ctlmsr)

    if vmx_ctl_one_setting(i, IA32_VMX_PROCBASED_2ND_CTLS):
        i = rdmsr(MSR_IA32_VMX_PROCBASED2_CTLS)
        if vmx_ctl_one_setting(i, IA32_VMX_PROCBASED2_EPT):
            note("VMX supports EPT")
            j = rdmsr(MSR_IA32_VMX_EPT_VPID_CAP)
            if j & IA32_VMX_EPT_VPID_INVEPT:
                note("VMX supports INVEPT")
                if j & IA32_VMX_EPT_VPID_INVEPT_SINGLE:
                    note("VMX supports single INVEPT")
                if j & IA32_VMX_EPT_VPID_INVEPT_ALL:
                    note("VMX supports all INVEPT")
        else:
            note("VMX does not support EPT (optional)")

        if vmx_ctl_one_setting(i, IA32_VMX_PROCBASED2_VPID):
            note("VMX supports VPID")
        else:
            note("VMX does not support VPID (optional)")

    return True

def svm_check():

    x = cpuid(0x80000001)
    if not x.ecx & CPUID_AMD_ECX_SVM:
        note("CPU does not support SVM")
        return False
    note("CPU supports SVM")

    i = rdmsr(MSR_AMD_VM_CR)
    if i & AMD_VM_CR_SVMDIS:
        note("SVM support not enabled in BIOS (essential)")
        return False
    note("SVM support is enabled in BIOS")

    x = cpuid(0x8000000a)
    nasid = x.ebx
    feat = x.edx

    if not nasid:
        note("Not enough ASIDs for guests")
        return False
    if not feat & CPUID_AMD_EDX_NESTED_PAGING:
        note("CPU does not support nested paging")
        return False
    if not feat & CPUID_AMD_EDX_NRIPS:
        note("CPU does not support NRIP save")
        return False

    return True

uc = ucred.get(os.getpid())
if not uc.has_priv("Effective", "sys_config"):
    print("Insufficient privileges, sys_config required")
    sys.exit(0)

vendor = cpu_vendor()
ret = False
if vendor == CPU_INTEL:
    ret = vmx_check()
elif vendor == CPU_AMD:
    ret = svm_check()
else:
    print("Did not detect a supported CPU")
    sys.exit(0)

print("\nbhyve is {}supported on this system.".format("" if ret else "not "))

# Vim hints
# vim:ts=4:sw=4:et:fdm=marker