vm

#!/usr/bin/env python3
import sys
import os
import subprocess
import time
import xml.etree.ElementTree as ET

os.environ['LIBVIRT_DEFAULT_URI'] = 'qemu:///system'

def _sh(cmds, wait=None, pipe=False):
    stdout = subprocess.PIPE if pipe else None
    process = subprocess.Popen(cmds, shell=True, universal_newlines=True, stdout=stdout, stderr=stdout)
    if wait == 0: return process
    out, err = process.communicate(timeout=wait)
    if pipe:
        process.stdout = out
        process.stderr = err
    return process
def _sh_out(cmds):
    return _sh(cmds, pipe=True).stdout

class CliBackend:
    def __init__(self):
        self.fast_mode = False

    def get_vm_list(self, state=None):
        texts = _sh_out(f'virsh list --all')
        lines = texts.strip().split('\n')[2:]
        vms = {}
        for line in lines:
            info = line.split(' ')
            _, name, stateStr, *_ = [e for e in info if len(e) > 0]
            running = stateStr == 'running'
            if state is None or state == running: vms[name] = running
        return vms

    def get_vm_config(self, vm):
        texts = _sh_out(f'virsh dumpxml "{vm}"').strip()
        et = ET.fromstring(texts)
        pci = []
        for element in et.findall(".//devices/hostdev/source/address"):
            keys = ['domain', 'bus', 'slot', 'function']
            domain, bus, device, function = [element.attrib[key][2:] for key in keys]
            bdf = f'{domain}:{bus}:{device}.{function}'
            pci.append(bdf)
        mem = int(et.find('.//currentMemory').text) / 1024 / 1024
        cpu = int(et.find('.//vcpu').text)
        return pci, mem, cpu

    def get_ip_of_vm(self, vm):
        cmd = f'virsh domifaddr {vm}'
        line = _sh_out(cmd).strip().split('\n')[-1]
        if not 'ipv4' in line: return
        ip_full = line.split(' ')[-1]
        return ip_full.split('/')[0]
    def get_hdd(self, vm):
        out = _sh_out(f'virsh domblklist {vm} | grep -E "qcow2"').rstrip()
        ret = []
        for line in out.split('\n'):
            parts = line.strip().split(' ')
            ret.append((parts[0], parts[-1]))
        return ret

class ApiBackend:
    """Implementation using libvirt API"""

    def __init__(self):
        self._conn = None
        self.fast_mode = True

    def _get_conn(self):
        """Get cached libvirt connection"""
        if self._conn is None:
            import libvirt
            self._conn = libvirt.open('qemu:///system')
            if self._conn is None:
                raise Exception('Failed to connect to QEMU/KVM')
        return self._conn

    def get_vm_list(self, state=None):
        conn = self._get_conn()
        return {
            domain.name(): domain.isActive()
            for domain in conn.listAllDomains()
            if state is None or state == domain.isActive()
        }

    def get_ip_of_vm(self, vm):
        conn = self._get_conn()
        domain = conn.lookupByName(vm)
        if domain is None: return None

        xml_root = ET.fromstring(domain.XMLDesc(0))
        macs = [interface.get('address') for interface in xml_root.findall(".//interface/mac")]
        if not macs: return None

        for net_name in conn.listNetworks():
            net = conn.networkLookupByName(net_name)
            for lease in net.DHCPLeases():
                if lease['mac'] in macs: return lease['ipaddr']
        return None

    def get_vm_config(self, vm):
        conn = self._get_conn()
        domain = conn.lookupByName(vm)
        if domain is None: return [], 0, 0

        xml_root = ET.fromstring(domain.XMLDesc(0))

        pci = []
        for element in xml_root.findall(".//devices/hostdev/source/address"):
            keys = ['domain', 'bus', 'slot', 'function']
            domain, bus, device, function = [element.attrib[key][2:] for key in keys]
            bdf = f'{domain}:{bus}:{device}.{function}'
            pci.append(bdf)

        mem = int(xml_root.find('.//currentMemory').text) / 1024 / 1024
        cpu = int(xml_root.find('.//vcpu').text)
        return pci, mem, cpu

    def get_hdd(self, vm):
        conn = self._get_conn()
        domain = conn.lookupByName(vm)
        if domain is None: return None

        ret = []
        xml_root = ET.fromstring(domain.XMLDesc(0))
        for disk in xml_root.findall(".//devices/disk"):
            if disk.get("device") != "disk": continue
            dev = disk.find("target").get("dev")
            file = disk.find("source").get("file")
            ret.append((dev, file))
        return ret

backend = CliBackend() if os.environ.get('VMBK') == 'cli' else ApiBackend()

def detach(vm, target):
    _sh(f'virsh detach-disk {vm} {target} --config')

def attach(vm, size_or_file):
    hdds = backend.get_hdd(vm)
    if size_or_file.isdigit():
        dirname = os.path.dirname(hdds[0][1])
        disk = f'{dirname}/{vm}.new_disk.{len(hdds)}.qcow2'
        _sh(f'qemu-img create -f qcow2 {disk} {size_or_file}G')
    else:
        disk = os.path.abspath(size_or_file)
    if not os.path.exists(disk):
        return print(f'Error: disk file {disk} not found')
    target = _disk_get_available_target(hdds)
    _sh(f'virsh attach-disk {vm} {disk} {target} --persistent --subdriver qcow2')
def _disk_get_available_target(hdds):
    used_targets = [hdd[0] for hdd in hdds]
    for i in range(26):
        target = 'vd' + chr(ord('a') + i)
        if target not in used_targets: return target
    return None

def _vm_info(vm):
    vms = backend.get_vm_list()
    hdds = backend.get_hdd(vm)
    pci, mem, cpu = backend.get_vm_config(vm)
    print(f'''
Name: {vm}
Running: {vms[vm]}
CPU cores: {cpu}
System Memory: {mem}
PCI: {pci}
IP: {backend.get_ip_of_vm(vm)}
HDD: {len(hdds)}
'''.strip())
    for idx, (name, hdd) in enumerate(hdds):
        print(f'{idx}. {name}: {hdd}')
        _sh(f'qemu-img info {hdd} | grep -E "backing file:|virtual size|disk size"')

def info(vm=None):
    if vm: return _vm_info(vm)
    print('\n---------------- CPU ----------------'); cpu()
    print('\n---------------- MEM ----------------'); mem()
    print('\n---------------- GPU ----------------'); gpu()
    print('')

def mem(vm=None, size=None):
    if not vm: return _sh('free -h')
    size = int(size) * 1024 * 1024
    _sh(f'''
        virsh setmaxmem {vm} {size} --config > /dev/null
        virsh setmem {vm} {size} --config > /dev/null
        ''')

def cpu(vm=None, count=None):
    if not vm: return _sh(r'lscpu | grep -E "^CPU\(s\):|NUMA node"')
    n = int(count)
    _sh(f'virt-xml {vm}  --edit --vcpus {n},sockets=1,cores={n//2},threads=2')
    if os.environ.get('vcpupin') == '1': _vcpupin(vm)

def _get_cpu_offset():
    node0Str = _sh_out('lscpu | grep node0').strip()
    number = node0Str.split(',')[1].split('-')[0]
    return int(number)
def _vcpupin(vm):
    _, _1, count = backend.get_vm_config(vm)
    offset = _get_cpu_offset()
    cmds = ''
    for i in range(count):
        cpuid = i // 2
        if i % 2: cpuid += offset
        cmds += f" virsh vcpupin {vm} --config {i} {cpuid} > /dev/null &&"
    cmds += 'echo "vcpu pin done"'
    _sh(cmds)

def gpu(vm=None, *devices):
    if not vm: return _sh('lspci | grep -E "acc|Display"')
    cmd = f'virt-xml {vm} --remove-device --host-dev all > /dev/null'
    if devices:
        suffix = ' '.join([f'--host-dev {dev}' for dev in devices])
        cmd += f' && virt-xml {vm} --add-device {suffix} > /dev/null'
    if _sh(cmd).returncode == 0:
        print(f"assigned GPU {devices} to VM '{vm}'.")
    else:
        print(f"Error: failed to assign GPU {devices} to VM '{vm}'.")

def _xf(dev_code, vm=None, *devices):
    output = _sh_out(f'lspci | grep -E "acc|Display" | grep -E ":{dev_code}."')
    lines = output.strip().split('\n')
    if not vm:
        if not output:
            return print(f"Error: no GPU found with code {dev_code}")
        for index, line in enumerate(lines): print(f'{index}: {line}')
        return
    vfs = [ line.split(' ')[0] for line in lines ]
    devs = [ vfs[int(dev)] for dev in devices ]
    gpu(vm, *devs)
def vf(vm=None, *devices): _xf('02', vm, *devices)
def pf(vm=None, *devices):
    _xf('00', vm, *devices)
    if devices and os.environ.get('rom_bar') != '1':
        _disable_rom_bar(vm)
def _disable_rom_bar(vm):
    xml = _sh_out(f'virsh dumpxml {vm} | tee /tmp/backup.vm.xml')
    xml = xml.replace('</hostdev>', '<rom bar="off"/></hostdev>')
    file = '/tmp/tmp.vm.xml'
    with open(file, 'w') as f: f.write(xml)
    _sh(f'virsh undefine {vm} && virsh define {file}')

def ls(**kwargs):
    vms = backend.get_vm_list()
    if not vms: return
    verbose = backend.fast_mode or '-v' in kwargs
    name_len_max = max([len(name) for name in vms])
    print(f'{"NAME":<{name_len_max}}\t STATE\tCPU {"MEM":>6} {"IP":>16}  PCI')
    active = [name for name in vms if vms[name]]
    inactive = [name for name in vms if not vms[name]]
    for name in sorted(active) + sorted(inactive):
        running = vms[name]
        pci, cpu, mem, ip = [], 0, 0, ''
        if verbose:
            pci, mem, cpu = backend.get_vm_config(name)
            ip = backend.get_ip_of_vm(name) or '-'
        print(f'{name:<{name_len_max}}\t {running:5} {cpu:>4} {mem:6.1f} {ip:>16}  {pci}')
cmd_list = ls


class DynamicLog:
    def __init__(self, sleep=1):
        self.start = time.time()
        self.sleep = sleep
        self.newline = False
    def print(self, msg, *args, **kwargs):
        diff = int(time.time() - self.start)
        print(f'\r{msg}, {diff}s...', end='')
        self.newline = True
        time.sleep(self.sleep)
    def done(self):
        if self.newline: print()
def _wait_host(ip):
    log = DynamicLog()
    while _sh(f'nc -zw 1 {ip} 22 >/dev/null 2>&1').returncode:
        log.print(f'sshd not ready on {ip}, try again later')
    log.done()
def _wait_vm_ip(vm):
    log = DynamicLog(2)
    while True:
        ip = backend.get_ip_of_vm(vm)
        if ip: break
        log.print(f'ip not found for {vm}, try again later')
    log.done()
    return ip
def ssh(vm, command=None):
    vms = backend.get_vm_list()
    if not vms[vm]: return print(f'Error: vm {vm} not running')
    ip = _wait_vm_ip(vm)
    _wait_host(ip)
    options = '-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -q'
    cmd = f'sshpass -p amd1234 ssh {options} root@{ip}'
    if command: cmd += f' -T "{command}"'
    return _sh(cmd).returncode
sh = ssh

def run(vm, cmd=None, **kwargs):
    if '--restart' in kwargs: stop(vm)
    vms = backend.get_vm_list()
    if not vms[vm]: start(vm)
    return ssh(vm, cmd)

def _change_name_to_ip(filepath):
    if ':' not in filepath: return filepath
    vm, file = filepath.split(':')
    ip = backend.get_ip_of_vm(vm)
    return f'root@{ip}:{file}'
def scp(src, dst):
    src = _change_name_to_ip(src)
    dst = _change_name_to_ip(dst)
    _sh(f'sshpass -p amd1234 scp -r {src} {dst}')
cp = scp

def xml(vm, edit=''):
    cmd = 'edit' if edit else 'dumpxml'
    _sh(f'virsh {cmd} {vm}')
def console(vm): _sh(f'virsh console --force {vm}')

def _get_gpus():
    texts = _sh_out('lspci | grep -E "acc|Display"')
    out = [line.split(' ')[0] for line in texts.strip().split('\n')]
    return [bdf.replace(":", r"\:") for bdf in out]
def _print_list(words, prefix=''):
    for word in words:
        if word.startswith(prefix): print(word)
def _complete(*_):
    COMP_POINT = int(os.environ['COMP_POINT'])
    COMP_LINE = os.environ['COMP_LINE']
    line_prefix = COMP_LINE[:COMP_POINT]
    words = [word for word in line_prefix.split(' ') if len(word) > 0]
    if line_prefix[-1] == ' ': words.append('')
    word1 = words[-1]
    if len(words) == 2: return _print_list(_get_local_functions(), word1)
    if words[1] == 'gpu' and len(words) > 3:
        return _print_list(_get_gpus(), word1)
    state = None
    if words[1] in ['stop', 'down', 'ssh', 'sh', 'scp', 'cp', 'console', 'restart']: state = True
    if words[1] in ['start', 'up']: state = False
    _print_list(backend.get_vm_list(state), word1)
def install():
    has_apt = _sh('apt --version >/dev/null 2>&1').returncode == 0
    if has_apt: deps = 'apt install -y guestfs-tools sshpass'
    else: deps = 'yum install -y libguestfs-tools sshpass libvirt-devel'
    _sh(f'''set -x;
        sudo usermod -a -G libvirt $USER
        sudo {deps};
        python3 -m pip install libvirt-python
        sudo cp {os.path.abspath(__file__)} /usr/bin/vm;
        sudo chmod a+x /usr/bin/vm;
        sudo ln -sf /usr/bin/vm {sys.path[-1]}/vm.py;
        echo 'complete -C "vm _complete" vm' | tee -a ~/.bashrc''')
    print(f'installed `vm` command, restart shell session to use it.')

def _shadow_clone_disks(disks, vm):
    new_disks = []
    for _, base_hdd in disks:
        backing_file = _sh_out(f'qemu-img info {base_hdd} | grep -E "backing file:"').strip()
        if backing_file:
            print(f'Warn: the base vm disk is derived disk! {base_hdd}, {backing_file}')
        dirname = os.path.dirname(base_hdd)
        basename = os.path.basename(base_hdd)
        new_hdd = f'{dirname}/{vm}.{basename}'
        _sh(f'qemu-img create -F qcow2 -b {base_hdd} -f qcow2 "{new_hdd}"')
        new_disks.append(new_hdd)
    return new_disks

def fork(base, *vms):
    disks = backend.get_hdd(base)
    for vm in vms:
        new_disks = _shadow_clone_disks(disks, vm)
        file_args = ' '.join([f'--file "{disk}" --preserve-data' for disk in new_disks])
        _sh(f'''
            virt-clone --original "{base}" --name "{vm}" {file_args} &&
            virt-sysprep -d {vm} --operation machine-id''')

def import_vm(name, file, centos=True):
    os_variant = 'centos8' if centos else 'ubuntu22.10'
    _sh(f'''
        virt-install \
            --name {name} \
            --memory 32000 \
            --vcpus 32 \
            --disk path={file},format=qcow2 \
            --import \
            --os-variant {os_variant} \
            --check path_in_use=off --wait=0 --noreboot
    ''')
def clone_vm(base, name):
    _sh(f'''
        virt-clone --original {base} --name {name} --file /disk1/vm_images/{name}.qcow2
    ''')

def remove(*vms, **kwargs):
    options = '--remove-all-storage' if '--rs' in kwargs else ''
    for vm in vms: _sh(f'virsh undefine {options} {vm}')
rm = remove

def start(*vms):
    for vm in vms: _sh(f'virsh start {vm} > /dev/null')
up = start

def stop(*vms, **kwargs):
    vms_info = backend.get_vm_list()
    if '-a' in kwargs: vms = vms_info.keys()
    for vm in vms:
        if vms_info[vm]: _sh(f'virsh destroy {vm} > /dev/null')
down = stop

def restart(*vms):
    stop(*vms)
    start(*vms)

def _parse_kwargs(all_args):
    kwargs = {}
    args = []
    for arg in all_args:
        if not arg.startswith('-'): args.append(arg)
        elif not '=' in arg: kwargs[arg] = ''
        else:
            key, value = arg.split('=')
            kwargs[key] = value
    return args, kwargs
def _get_local_functions():
    return [ name for name, obj in globals().items()
            if not name.startswith('_') and callable(obj) and not isinstance(obj, type) ]
def _main():
    if len(sys.argv) < 2: return _print_list(_get_local_functions())
    _, name, *args = sys.argv
    syms = globals()
    sym = syms.get(name) or syms.get('cmd_' + name)
    if not callable(sym): return print(f'Error: invalid function: {name}')
    args, kwargs = _parse_kwargs(args)
    code = sym(*args, **kwargs)
    if isinstance(code, int): sys.exit(code)
if __name__ == "__main__": _main()