vm

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

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

def _sh(cmds, wait=None, pipe=False):
    stdout = subprocess.PIPE if pipe else None
    process = subprocess.Popen(cmds, shell=True, text=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

def _get_vm_list(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 _vm_info(vm):
    vms = _get_vm_list()
    hdd = _get_hdd(vm)
    pci, mem, cpu = _get_vm_config(vm)
    print(f'''
Name: {vm}
Running: {vms[vm]}
CPU cores: {cpu}
System Memory: {mem}
PCI: {pci}
IP: {_get_ip_of_vm(vm)}
HDD: {hdd}
'''.strip())
    _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
    print(f"<memory unit='KiB'>{size}</memory><currentMemory unit='KiB'>{size}</currentMemory>")
    input("press Enter to start edit xml file:")
    _sh(f'virsh edit {vm}')

def cpu(vm=None, count=None):
    if not vm: return _sh(r'lscpu | grep -E "^CPU\(s\):|NUMA node"')
    _sh(f'virt-xml {vm}  --edit --vcpus {count}')

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\n'
    if devices:
        suffix = [f'--host-dev {dev}' for dev in devices]
        cmd += f'virt-xml {vm} --add-device ' + ' '.join(suffix)
    _sh(cmd)

def ls(**kwargs):
    vms = _get_vm_list()
    if not vms: return
    verbose = '-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')
    for name in vms:
        running = vms[name]
        pci, cpu, mem = [], 0, 0
        if verbose: pci, mem, cpu = _get_vm_config(name)
        ip = _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

def _get_vm_config(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(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]

class DynamicLog:
    def __init__(self, sleep=1):
        self.start = time.time()
        self.sleep = sleep
    def print(self, msg, *args, **kwargs):
        diff = int(time.time() - self.start)
        print(f'\r{msg}, {diff}s...', end='')
        time.sleep(self.sleep)
    def done(self): 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 = _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 = _get_vm_list()
    if not vms[vm]: return print(f'Error: vm {vm} not running')
    ip = _wait_vm_ip(vm)
    _wait_host(ip)
    cmd = f'sshpass -p amd1234 ssh -o StrictHostKeyChecking=no root@{ip}'
    if command: cmd += f' -t "{command}"'
    _sh(cmd)
sh = ssh

def run(vm, cmd=None):
    vms = _get_vm_list()
    if not vms[vm]: start(vm)
    ssh(vm, cmd)

def _change_name_to_ip(filepath):
    if ':' not in filepath: return filepath
    vm, file = filepath.split(':')
    ip = _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 _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']: state = True
    if words[1] in ['start', 'up']: state = False
    _print_list(_get_vm_list(state), word1)
def install():
    _sh(f'''set -x;
        sudo usermod -a -G libvirt $USER
        sudo apt install -y guestfs-tools sshpass;
        sudo cp {os.path.abspath(__file__)} /usr/bin/vm;
        sudo chmod a+x /usr/bin/vm;
        echo 'complete -C "vm _complete" vm' | tee -a ~/.bashrc''')
    print(f'installed `vm` command, restart shell session to use it.')

def _get_hdd(vm):
    out = _sh_out(f'virsh domblklist {vm} | grep -E "vda|hda"')
    return out.strip().split(' ')[-1]

def _fork_one_vm(base, base_hdd, vm):
    new_hdd = f'{os.path.dirname(base_hdd)}/{vm}.qcow2'
    _sh(f'''qemu-img create -f qcow2 -F qcow2 -b {base_hdd} "{new_hdd}" &&
        virt-clone --original "{base}" --name "{vm}" --file "{new_hdd}" --preserve-data &&
        virt-sysprep -d {vm} --operation machine-id''')
def fork(base, *vms):
    base_hdd = _get_hdd(base)
    backing_file = _sh_out(f'qemu-img info {base_hdd} | grep -E "backing file:"').strip()
    if backing_file:
        print(f'Warn: the disk of the vm is derived disk!\n{base}: {base_hdd}\n{backing_file}')
    for vm in vms: _fork_one_vm(base, base_hdd, vm)

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}')
up = start

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

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) ]
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)
    sym(*args, **kwargs)
if __name__ == "__main__": _main()