path: root/src/virtual/ogOperations.py

#
# Copyright (C) 2020 Soleta Networks <info@soleta.eu>
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the
# Free Software Foundation, version 3.
#

import socket
import errno
import enum
import json
import subprocess
import shutil
import os
import guestfs
import hivex
import pathlib
import re
import math

class OgQMP:
    class State(enum.Enum):
        CONNECTING = 0
        RECEIVING = 1
        FORCE_DISCONNECTED = 2

    def __init__(self, ip, port):
        self.ip = ip
        self.port = port

    def connect(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.state = self.State.CONNECTING
        self.data = ""

        try:
            self.sock.connect((self.ip, self.port))
        except socket.error as err:
            print('Error connection' + str(err))
            return None

    def recv(self):
        self.data = self.sock.recv(50024).decode('utf-8')
        return self.data


    def send(self, data=None):
        if not data:
            return None

        self.sock.send(bytes(data, 'utf-8'))
        return len(data)

    def disconnect(self):
        self.state = self.State.FORCE_DISCONNECTED
        self.sock.close()

class OgVirtualOperations:
    def __init__(self):
        self.IP = '127.0.0.1'
        self.VIRTUAL_PORT = 4444

    def poweroff(self):
        qmp = OgQMP(self.IP, self.VIRTUAL_PORT)
        qmp.connect()
        qmp.recv()
        qmp.send(str({"execute": "qmp_capabilities"}))
        qmp.recv()
        qmp.send(str({"execute": "system_powerdown"}))
        qmp.disconnect()

    def reboot(self):
        qmp = OgQMP(self.IP, self.VIRTUAL_PORT)
        qmp.connect()
        qmp.recv()
        qmp.send(str({"execute": "qmp_capabilities"}))
        qmp.recv()
        qmp.send(str({"execute": "system_reset"}))
        qmp.disconnect()

    def session(self, request, ogRest):
        disk = request.getDisk()
        partition = request.getPartition()

        available_ram = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
        available_ram_mib = available_ram / 1024 ** 2
        # Calculate the lower power of 2 amout of RAM memory for the VM.
        vm_ram_mib = 2 ** math.floor(math.log(available_ram_mib) / math.log(2))

        cmd = (f'qemu-system-x86_64 -hda disk{disk}_part{partition}.qcow2 '
               f'-qmp tcp:localhost:4444,server,nowait --enable-kvm '
               f'-display gtk -cpu host -m {vm_ram_mib}M -boot c')
        subprocess.Popen([cmd], shell=True)

    def refresh(self, ogRest):
        path = 'partitions.json'
        temp_mount_dir = 'mnt'
        try:
            with open(path, 'r+') as f:
                data = json.loads(f.read())
                for part in data['partition_setup']:
                    if len(part['virt-drive']) > 0:
                        subprocess.run([f'qemu-nbd -c /dev/nbd0 {part["virt-drive"]}'],
                                       shell=True)
                        subprocess.run([f'partprobe /dev/nbd0'],
                                       shell=True)
                        subprocess.run([f'mount /dev/nbd0p1 {temp_mount_dir}'],
                                       shell=True)
                        total_disk, used_disk, free_disk = shutil.disk_usage(temp_mount_dir)
                        subprocess.run([f'umount {temp_mount_dir}'],
                                       shell=True)
                        subprocess.run(['qemu-nbd -d /dev/nbd0'],
                                       shell=True)
                        part['size'] = int(free_disk / 1024)
                        part['used_size'] = int(100 * used_disk / free_disk)
                f.seek(0)
                f.write(json.dumps(data, indent=4))
                f.truncate()
        except FileNotFoundError:
            total_disk, used_disk, free_disk = shutil.disk_usage("/")
            data = {'serial_number': '',
                    'disk_setup': {'disk': 1,
                                   'partition': 0,
                                   'code': '0',
                                   'filesystem': '',
                                   'os': '',
                                   'size': int(free_disk / 1024),
                                   'used_size': int(100 * used_disk / free_disk)},
                    'partition_setup': []}
            for i in range(4):
                part_json = {'disk': 1,
                             'partition': i + 1,
                             'code': '',
                             'filesystem': 'EMPTY',
                             'os': '',
                             'size': 0,
                             'used_size': 0,
                             'virt-drive': ''}
                data['partition_setup'].append(part_json)
            with open(path, 'w+') as f:
                f.write(json.dumps(data, indent=4))

        # TODO no debería ser necesario eliminar virt-drive ni transformar a strings
        for part in data['partition_setup']:
            part.pop('virt-drive')
            for k, v in part.items():
                part[k] = str(v)
        data['disk_setup'] = {k: str(v) for k, v in data['disk_setup'].items()}

        return data

    def setup(self, request, ogRest):
        path = 'partitions.json'
        refresh(ogRest)

        part_setup = request.getPartitionSetup()
        disk = request.getDisk()

        for i, part in enumerate(part_setup):
            if int(part['format']) == 0:
                continue

            drive_path = f'disk{disk}_part{part["partition"]}.qcow2'
            sequence = [f'qemu-img create -f qcow2 {drive_path} {part["size"]}K',
                        f'qemu-nbd -c /dev/nbd0 {drive_path}',
                        f'parted /dev/nbd0 mklabel gpt',
                        f'parted -a opt /dev/nbd0 mkpart primary {part["filesystem"]} 2048s 100%',
                        f'mkfs.{part["filesystem"].lower()} /dev/nbd0p1',
                        f'qemu-nbd -d /dev/nbd0']
            for cmd in sequence:
                process = subprocess.run([cmd], shell=True)
                if process.returncode != 0:
                    raise RuntimeError

            with open(path, 'r+') as f:
                data = json.loads(f.read())
                if part['code'] == 'LINUX':
                    data['partition_setup'][i]['code'] = '0083'
                elif part['code'] == 'NTFS':
                    data['partition_setup'][i]['code'] = '0007'
                elif part['code'] == 'DATA':
                    data['partition_setup'][i]['code'] = '00da'
                data['partition_setup'][i]['filesystem'] = part['filesystem']
                data['partition_setup'][i]['size'] = int(part['size'])
                data['partition_setup'][i]['virt-drive'] = drive_path
                f.seek(0)
                f.write(json.dumps(data, indent=4))
                f.truncate()

        return refresh(ogRest)

    def image_create(self, path, request, ogRest):
        disk = request.getDisk()
        partition = request.getPartition()
        name = request.getName()
        repo = request.getRepo()

        # Check if VM is running.
        qmp = OgQMP(self.IP, self.VIRTUAL_PORT)
        if qmp.connect() != None:
            qmp.disconnect()
            return None

        refresh(ogRest)

        drive_path = f'disk{disk}_part{partition}.qcow2'
        repo_path = f'images'

        try:
            shutil.copy(drive_path, f'{repo_path}/{name}')
        except:
            return None

        return True

    def image_restore(self, request, ogRest):
        disk = request.getDisk()
        partition = request.getPartition()
        name = request.getName()
        repo = request.getRepo()
        # TODO Multicast? Unicast? Solo copia con samba.
        ctype = request.getType()
        profile = request.getProfile()
        cid = request.getId()

        # Check if VM is running.
        qmp = OgQMP(self.IP, self.VIRTUAL_PORT)
        if qmp.connect() != None:
            qmp.disconnect()
            return None

        refresh(ogRest)

        drive_path = f'disk{disk}_part{partition}.qcow2'
        repo_path = f'images'

        if os.path.exists(drive_path):
            os.remove(drive_path)

        try:
            shutil.copy(f'{repo_path}/{name}', drive_path)
        except:
            return None

        return True

    def software(self, request, path, ogRest):
        DPKG_PATH = '/var/lib/dpkg/status'

        disk = request.getDisk()
        partition = request.getPartition()
        drive_path = f'disk{disk}_part{partition}.qcow2'
        g = guestfs.GuestFS(python_return_dict=True)
        g.add_drive_opts(drive_path, readonly=1)
        g.launch()
        root = g.inspect_os()[0]

        os_type = g.inspect_get_type(root)
        os_major_version = g.inspect_get_major_version(root)
        os_minor_version = g.inspect_get_minor_version(root)
        os_distro = g.inspect_get_distro(root)

        software = []

        if 'linux' in os_type:
            g.mount_ro(g.list_partitions()[0], '/')
            try:
                g.download('/' + DPKG_PATH, 'dpkg_list')
            except:
                pass
            g.umount_all()
            if os.path.isfile('dpkg_list'):
                pkg_pattern = re.compile('Package: (.+)')
                version_pattern = re.compile('Version: (.+)')
                with open('dpkg_list', 'r') as f:
                    for line in f:
                        pkg_match = pkg_pattern.match(line)
                        version_match = version_pattern.match(line)
                        if pkg_match:
                            pkg = pkg_match.group(1)
                        elif version_match:
                            version = version_match.group(1)
                        elif line == '\n':
                            software.append(pkg + ' ' + version)
                        else:
                            continue
                os.remove('dpkg_list')
        elif 'windows' in os_type:
            g.mount_ro(g.list_partitions()[0], '/')
            hive_file_path = g.inspect_get_windows_software_hive(root)
            g.download('/' + hive_file_path, 'win_reg')
            g.umount_all()
            h = hivex.Hivex('win_reg')
            key = h.root()
            key = h.node_get_child (key, 'Microsoft')
            key = h.node_get_child (key, 'Windows')
            key = h.node_get_child (key, 'CurrentVersion')
            key = h.node_get_child (key, 'Uninstall')
            software += [h.node_name(x) for x in h.node_children(key)]
            # Just for 64 bit Windows versions, check for 32 bit software.
            if (os_major_version == 5 and os_minor_version >= 2) or \
               (os_major_version >= 6):
                key = h.root()
                key = h.node_get_child (key, 'Wow6432Node')
                key = h.node_get_child (key, 'Microsoft')
                key = h.node_get_child (key, 'Windows')
                key = h.node_get_child (key, 'CurrentVersion')
                key = h.node_get_child (key, 'Uninstall')
                software += [h.node_name(x) for x in h.node_children(key)]
            os.remove('win_reg')

        with open(path, 'w+') as f:
            f.seek(0)
            for program in software:
                f.write(f'{program}\n')
            f.truncate()

    def parse_pci(self, path='/usr/share/misc/pci.ids'):
        data = {}
        with open(path, 'r') as f:
            for line in f:
                if line[0] == '#':
                    continue
                elif len(line.strip()) == 0:
                    continue
                else:
                    if line[:2] == '\t\t':
                        fields = line.strip().split(maxsplit=2)
                        data[last_vendor][last_device][(fields[0], fields[1])] = fields[2]
                    elif line[:1] == '\t':
                        fields = line.strip().split(maxsplit=1)
                        last_device = fields[0]
                        data[last_vendor][fields[0]] = {'name': fields[1]}
                    else:
                        fields = line.strip().split(maxsplit=1)
                        if fields[0] == 'ffff':
                            break
                        last_vendor = fields[0]
                        data[fields[0]] = {'name': fields[1]}
        return data

    def hardware(self, path, ogRest):
        qmp = OgQMP(self.IP, self.VIRTUAL_PORT)
        qmp.connect()
        qmp.recv()
        qmp.send(str({"execute": "qmp_capabilities"}))
        qmp.recv()
        qmp.send(str({"execute": "query-pci"}))
        data = json.loads(qmp.recv())
        data = data['return'][0]['devices']
        pci_list = parse_pci()
        device_names = {}
        for device in data:
            vendor_id = hex(device['id']['vendor'])[2:]
            device_id = hex(device['id']['device'])[2:]
            subvendor_id = hex(device['id']['subsystem-vendor'])[2:]
            subdevice_id = hex(device['id']['subsystem'])[2:]
            description = device['class_info']['desc'].lower()
            name = ''
            try:
                name = pci_list[vendor_id]['name']
                name += ' ' + pci_list[vendor_id][device_id]['name']
                name += ' ' + pci_list[vendor_id][device_id][(subvendor_id, subdevice_id)]
            except KeyError:
                if vendor_id == '1234':
                    name = 'VGA Cirrus Logic GD 5446'
                else:
                    pass

            if 'usb' in description:
                device_names['usb'] = name
            elif 'ide' in description:
                device_names['ide'] = name
            elif 'ethernet' in description:
                device_names['net'] = name
            elif 'vga' in description:
                device_names['vga'] = name

            elif 'audio' in description or 'sound' in description:
                device_names['aud'] = name
            elif 'dvd' in description:
                device_names['cdr'] = name

        qmp.send(str({"execute": "query-memory-size-summary"}))
        data = json.loads(qmp.recv())
        ram_size = int(data['return']['base-memory']) * 2 ** -20
        device_names['mem'] = f'QEMU {int(ram_size)}MiB'

        qmp.send(str({"execute": "query-cpus-fast"}))
        data = json.loads(qmp.recv())
        qmp.disconnect()
        cpu_arch = data['return'][0]['arch']
        cpu_target = data['return'][0]['target']
        cpu_cores = len(data['return'])
        device_names['cpu'] = f'CPU arch:{cpu_arch} target:{cpu_target} ' \
                              f'cores:{cpu_cores}'

        with open(path, 'w+') as f:
            f.seek(0)
            for k, v in device_names.items():
                f.write(f'{k}={v}\n')
            f.truncate()