source: ogClient-Git/src/virtual/ogOperations.py @ b29b2eb

Last change on this file since b29b2eb was b29b2eb, checked in by Roberto Hueso Gómez <rhueso@…>, 5 years ago

Move check_vm_state_loop() into OgVirtualOperations?

Improves code encapsulation by moving check_vm_state_loop method into
OgVirtualOperations? class. This also fixes import error when running ogclient in
'linux' mode.

  • Property mode set to 100644
File size: 21.7 KB
Line 
1#
2# Copyright (C) 2020 Soleta Networks <info@soleta.eu>
3#
4# This program is free software: you can redistribute it and/or modify it under
5# the terms of the GNU Affero General Public License as published by the
6# Free Software Foundation, version 3.
7#
8
9from src.ogRest import ThreadState
10import socket
11import errno
12import select
13import json
14import subprocess
15import shutil
16import os
17import guestfs
18import hivex
19import pathlib
20import re
21import math
22import sys
23import enum
24import time
25
26class OgVM:
27    DEFAULT_CPU = 'host'
28    DEFAULT_VGA = 'virtio-vga'
29    DEFAULT_QMP_IP = 'localhost'
30    DEFAULT_QMP_PORT = 4444
31
32    class State(enum.Enum):
33        STOPPED = 0
34        RUNNING = 1
35
36    def __init__(self,
37                 partition_path,
38                 memory=None,
39                 cpu=DEFAULT_CPU,
40                 vga=DEFAULT_VGA,
41                 qmp_ip=DEFAULT_QMP_IP,
42                 qmp_port=DEFAULT_QMP_PORT,
43                 vnc_params=None):
44        self.partition_path = partition_path
45        self.cpu = cpu
46        self.vga = vga
47        self.qmp_ip = qmp_ip
48        self.qmp_port = qmp_port
49        self.proc = None
50        self.vnc_params = vnc_params
51
52        if memory:
53            self.mem = memory
54        else:
55            available_ram = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
56            available_ram = available_ram / 2 ** 20
57            # Calculate the lower power of 2 amout of RAM memory for the VM.
58            self.mem = 2 ** math.floor(math.log(available_ram) / math.log(2))
59
60
61    def run_vm(self):
62        if self.vnc_params:
63            vnc_str = f'-vnc 0.0.0.0:0,password'
64        else:
65            vnc_str = ''
66
67        cmd = (f'qemu-system-x86_64 -accel kvm -cpu {self.cpu} -smp 4 '
68               f'-drive file={self.partition_path},if=virtio '
69               f'-qmp tcp:localhost:4444,server,nowait '
70               f'-device {self.vga} -display gtk '
71               f'-m {self.mem}M -boot c -full-screen {vnc_str}')
72        self.proc = subprocess.Popen([cmd], shell=True)
73
74        if self.vnc_params:
75            # Wait for QMP to be available.
76            time.sleep(20)
77            cmd = { "execute": "change",
78                    "arguments": { "device": "vnc",
79                                   "target": "password",
80                                   "arg": str(self.vnc_params['pass']) } }
81            with OgQMP(self.qmp_ip, self.qmp_port) as qmp:
82                qmp.talk(str(cmd))
83
84class OgQMP:
85    QMP_TIMEOUT = 5
86    QMP_POWEROFF_TIMEOUT = 300
87
88    def __init__(self, ip, port):
89        self.ip = ip
90        self.port = port
91
92    def connect(self):
93        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
94        self.sock.setblocking(0)
95        try:
96            self.sock.connect((self.ip, self.port))
97        except socket.error as err:
98            if err.errno == errno.ECONNREFUSED:
99                raise Exception('cannot connect to qemu')
100            elif err.errno == errno.EINPROGRESS:
101                pass
102
103        readset = [ self.sock ]
104        readable, writable, exception = select.select(readset,
105                                                      [],
106                                                      [],
107                                                      OgQMP.QMP_TIMEOUT)
108
109        if self.sock in readable:
110            try:
111                out = self.recv()
112            except:
113                pass
114
115        if 'QMP' not in out:
116            raise Exception('cannot handshake qemu')
117
118        out = self.talk(str({"execute": "qmp_capabilities"}))
119        if 'return' not in out:
120            raise Exception('cannot handshake qemu')
121
122    def disconnect(self):
123        try:
124            self.sock.close()
125        except:
126            pass
127
128    def __enter__(self):
129        self.connect()
130        return self
131
132    def __exit__(self, exc_type, exc_val, exc_tb):
133        self.disconnect()
134
135    def talk(self, data, timeout=QMP_TIMEOUT):
136        writeset = [ self.sock ]
137        readable, writable, exception = select.select([],
138                                                      writeset,
139                                                      [],
140                                                      timeout)
141        if self.sock in writable:
142            try:
143                self.sock.send(bytes(data, 'utf-8'))
144            except:
145                raise Exception('cannot talk to qemu')
146        else:
147            raise Exception('timeout when talking to qemu')
148
149        return self.recv(timeout=timeout)
150
151    def recv(self, timeout=QMP_TIMEOUT):
152        readset = [self.sock]
153        readable, _, _ = select.select(readset, [], [], timeout)
154
155        if self.sock in readable:
156            try:
157                out = self.sock.recv(4096).decode('utf-8')
158                out = json.loads(out)
159            except socket.error as err:
160                raise Exception('cannot talk to qemu')
161        else:
162            raise Exception('timeout when talking to qemu')
163        return out
164
165class OgVirtualOperations:
166    def __init__(self):
167        self.IP = '127.0.0.1'
168        self.VIRTUAL_PORT = 4444
169        self.USABLE_DISK = 0.75
170        self.OG_PATH = os.path.dirname(os.path.realpath(sys.argv[0]))
171        self.OG_IMAGES_PATH = f'{self.OG_PATH}/images'
172        self.OG_PARTITIONS_PATH = f'{self.OG_PATH}/partitions'
173        self.OG_PARTITIONS_CFG_PATH = f'{self.OG_PATH}/partitions.json'
174
175        if not os.path.exists(self.OG_IMAGES_PATH):
176            os.mkdir(self.OG_IMAGES_PATH, mode=0o755)
177        if not os.path.exists(self.OG_PARTITIONS_PATH):
178            os.mkdir(self.OG_PARTITIONS_PATH, mode=0o755)
179
180    def poweroff_guest(self):
181        try:
182            with OgQMP(self.IP, self.VIRTUAL_PORT) as qmp:
183                qmp.talk(str({"execute": "system_powerdown"}))
184                out = qmp.recv()
185                assert(out['event'] == 'POWERDOWN')
186                out = qmp.recv(timeout=OgQMP.QMP_POWEROFF_TIMEOUT)
187                assert(out['event'] == 'SHUTDOWN')
188        except:
189            return
190
191    def poweroff_host(self):
192        subprocess.run(['/sbin/poweroff'])
193
194    def poweroff(self):
195        self.poweroff_guest()
196        self.poweroff_host()
197
198    def reboot(self):
199        try:
200            with OgQMP(self.IP, self.VIRTUAL_PORT) as qmp:
201                qmp.talk(str({"execute": "system_reset"}))
202        except:
203            pass
204
205    def check_vm_state(self):
206        try:
207            with OgQMP(self.IP, self.VIRTUAL_PORT) as qmp:
208                pass
209            return OgVM.State.RUNNING
210        except:
211            return OgVM.State.STOPPED
212
213    def get_installed_os(self):
214        installed_os = {}
215        try:
216            with open(self.OG_PARTITIONS_CFG_PATH, 'r') as f:
217                cfg = json.loads(f.read())
218            for part in cfg['partition_setup']:
219                if len(part['os']) > 0:
220                    installed_os[part['os']] = (part['disk'], part['partition'])
221        except:
222            pass
223        return installed_os
224
225    def check_vm_state_loop(self, ogRest):
226        POLLING_WAIT_TIME = 12
227        while True:
228            time.sleep(POLLING_WAIT_TIME)
229            state = self.check_vm_state()
230            installed_os = self.get_installed_os()
231            if state == OgVM.State.STOPPED and \
232               ogRest.state == ThreadState.IDLE and \
233               len(installed_os) > 0:
234                self.poweroff_host()
235
236    def shellrun(self, request, ogRest):
237        return
238
239    def session(self, request, ogRest):
240        disk = request.getDisk()
241        partition = request.getPartition()
242
243        part_path = f'{self.OG_PARTITIONS_PATH}/disk{disk}_part{partition}.qcow2'
244        if ogRest.CONFIG['vnc']['activate']:
245            qemu = OgVM(part_path, vnc_params=ogRest.CONFIG['vnc'])
246        else:
247            qemu = OgVM(part_path)
248        qemu.run_vm()
249
250    def partitions_cfg_to_json(self, data):
251        for part in data['partition_setup']:
252            part.pop('virt-drive')
253            for k, v in part.items():
254                part[k] = str(v)
255        data['disk_setup'] = {k: str(v) for k, v in data['disk_setup'].items()}
256        return data
257
258    def refresh(self, ogRest):
259        try:
260            # Return last partitions setup in case VM is running.
261            with OgQMP(self.IP, self.VIRTUAL_PORT) as qmp:
262                pass
263            with open(self.OG_PARTITIONS_CFG_PATH, 'r') as f:
264                data = json.loads(f.read())
265            data = self.partitions_cfg_to_json(data)
266            return data
267        except:
268            pass
269
270        try:
271            with open(self.OG_PARTITIONS_CFG_PATH, 'r+') as f:
272                data = json.loads(f.read())
273                for part in data['partition_setup']:
274                    if len(part['virt-drive']) > 0:
275                        if not os.path.exists(part['virt-drive']):
276                            part['code'] = '',
277                            part['filesystem'] = 'EMPTY'
278                            part['os'] = ''
279                            part['size'] = 0
280                            part['used_size'] = 0
281                            part['virt-drive'] = ''
282                            continue
283                        g = guestfs.GuestFS(python_return_dict=True)
284                        g.add_drive_opts(part['virt-drive'],
285                                         format="qcow2",
286                                         readonly=1)
287                        g.launch()
288                        devices = g.list_devices()
289                        assert(len(devices) == 1)
290                        partitions = g.list_partitions()
291                        assert(len(partitions) == 1)
292                        filesystems_dict = g.list_filesystems()
293                        assert(len(filesystems_dict) == 1)
294                        g.mount(partitions[0], '/')
295                        used_disk = g.du('/')
296                        g.umount_all()
297                        total_size = g.disk_virtual_size(part['virt-drive']) / 1024
298
299                        part['used_size'] = int(100 * used_disk / total_size)
300                        part['size'] = total_size
301                        root = g.inspect_os()
302                        if len(root) == 1:
303                            part['os'] = f'{g.inspect_get_distro(root[0])} ' \
304                                         f'{g.inspect_get_product_name(root[0])}'
305                        else:
306                            part['os'] = ''
307                        filesystem = [fs for fs in filesystems_dict.values()][0]
308                        part['filesystem'] = filesystem.upper()
309                        if filesystem == 'ext4':
310                            part['code'] = '0083'
311                        elif filesystem == 'ntfs':
312                            part['code'] = '0007'
313                        g.close()
314                f.seek(0)
315                f.write(json.dumps(data, indent=4))
316                f.truncate()
317        except FileNotFoundError:
318            total_disk, used_disk, free_disk = shutil.disk_usage("/")
319            free_disk = int(free_disk * self.USABLE_DISK)
320            data = {'serial_number': '',
321                    'disk_setup': {'disk': 1,
322                                   'partition': 0,
323                                   'code': '0',
324                                   'filesystem': '',
325                                   'os': '',
326                                   'size': int(free_disk / 1024),
327                                   'used_size': int(100 * used_disk / total_disk)},
328                    'partition_setup': []}
329            for i in range(4):
330                part_json = {'disk': 1,
331                             'partition': i + 1,
332                             'code': '',
333                             'filesystem': 'EMPTY',
334                             'os': '',
335                             'size': 0,
336                             'used_size': 0,
337                             'virt-drive': ''}
338                data['partition_setup'].append(part_json)
339            with open(self.OG_PARTITIONS_CFG_PATH, 'w+') as f:
340                f.write(json.dumps(data, indent=4))
341        except:
342            with open(self.OG_PARTITIONS_CFG_PATH, 'r') as f:
343                data = json.load(f)
344
345        data = self.partitions_cfg_to_json(data)
346
347        return data
348
349    def setup(self, request, ogRest):
350        self.poweroff_guest()
351        self.refresh(ogRest)
352
353        part_setup = request.getPartitionSetup()
354        disk = request.getDisk()
355
356        for i, part in enumerate(part_setup):
357            if int(part['format']) == 0:
358                continue
359
360            drive_path = f'{self.OG_PARTITIONS_PATH}/disk{disk}_part{part["partition"]}.qcow2'
361            g = guestfs.GuestFS(python_return_dict=True)
362            g.disk_create(drive_path, "qcow2", int(part['size']) * 1024)
363            g.add_drive_opts(drive_path, format="qcow2", readonly=0)
364            g.launch()
365            devices = g.list_devices()
366            assert(len(devices) == 1)
367            g.part_disk(devices[0], "gpt")
368            partitions = g.list_partitions()
369            assert(len(partitions) == 1)
370            g.mkfs(part["filesystem"].lower(), partitions[0])
371            g.close()
372
373            with open(self.OG_PARTITIONS_CFG_PATH, 'r+') as f:
374                data = json.loads(f.read())
375                if part['code'] == 'LINUX':
376                    data['partition_setup'][i]['code'] = '0083'
377                elif part['code'] == 'NTFS':
378                    data['partition_setup'][i]['code'] = '0007'
379                elif part['code'] == 'DATA':
380                    data['partition_setup'][i]['code'] = '00da'
381                data['partition_setup'][i]['filesystem'] = part['filesystem']
382                data['partition_setup'][i]['size'] = int(part['size'])
383                data['partition_setup'][i]['virt-drive'] = drive_path
384                f.seek(0)
385                f.write(json.dumps(data, indent=4))
386                f.truncate()
387
388        return self.refresh(ogRest)
389
390    def image_create(self, path, request, ogRest):
391        disk = request.getDisk()
392        partition = request.getPartition()
393        name = request.getName()
394        repo = request.getRepo()
395        samba_config = ogRest.samba_config
396
397        self.poweroff_guest()
398
399        self.refresh(ogRest)
400
401        drive_path = f'{self.OG_PARTITIONS_PATH}/disk{disk}_part{partition}.qcow2'
402
403        cmd = f'mount -t cifs //{repo}/ogimages {self.OG_IMAGES_PATH} -o ' \
404              f'rw,nolock,serverino,acl,' \
405              f'username={samba_config["user"]},' \
406              f'password={samba_config["pass"]}'
407        subprocess.run([cmd], shell=True)
408
409        try:
410            shutil.copy(drive_path, f'{self.OG_IMAGES_PATH}/{name}')
411        except:
412            return None
413
414        subprocess.run([f'umount {self.OG_IMAGES_PATH}'], shell=True)
415
416        return True
417
418    def image_restore(self, request, ogRest):
419        disk = request.getDisk()
420        partition = request.getPartition()
421        name = request.getName()
422        repo = request.getRepo()
423        # TODO Multicast? Unicast? Solo copia con samba.
424        ctype = request.getType()
425        profile = request.getProfile()
426        cid = request.getId()
427        samba_config = ogRest.samba_config
428
429        self.poweroff_guest()
430        self.refresh(ogRest)
431
432        drive_path = f'{self.OG_PARTITIONS_PATH}/disk{disk}_part{partition}.qcow2'
433
434        if os.path.exists(drive_path):
435            os.remove(drive_path)
436
437        cmd = f'mount -t cifs //{repo}/ogimages {self.OG_IMAGES_PATH} -o ' \
438              f'ro,nolock,serverino,acl,' \
439              f'username={samba_config["user"]},' \
440              f'password={samba_config["pass"]}'
441        subprocess.run([cmd], shell=True)
442
443        try:
444            shutil.copy(f'{self.OG_IMAGES_PATH}/{name}', drive_path)
445        except:
446            return None
447
448        subprocess.run([f'umount {self.OG_IMAGES_PATH}'], shell=True)
449
450        return True
451
452    def software(self, request, path, ogRest):
453        DPKG_PATH = '/var/lib/dpkg/status'
454
455        disk = request.getDisk()
456        partition = request.getPartition()
457        drive_path = f'{self.OG_PARTITIONS_PATH}/disk{disk}_part{partition}.qcow2'
458        g = guestfs.GuestFS(python_return_dict=True)
459        g.add_drive_opts(drive_path, readonly=1)
460        g.launch()
461        root = g.inspect_os()[0]
462
463        os_type = g.inspect_get_type(root)
464        os_major_version = g.inspect_get_major_version(root)
465        os_minor_version = g.inspect_get_minor_version(root)
466        os_distro = g.inspect_get_distro(root)
467
468        software = []
469
470        if 'linux' in os_type:
471            g.mount_ro(g.list_partitions()[0], '/')
472            try:
473                g.download('/' + DPKG_PATH, 'dpkg_list')
474            except:
475                pass
476            g.umount_all()
477            if os.path.isfile('dpkg_list'):
478                pkg_pattern = re.compile('Package: (.+)')
479                version_pattern = re.compile('Version: (.+)')
480                with open('dpkg_list', 'r') as f:
481                    for line in f:
482                        pkg_match = pkg_pattern.match(line)
483                        version_match = version_pattern.match(line)
484                        if pkg_match:
485                            pkg = pkg_match.group(1)
486                        elif version_match:
487                            version = version_match.group(1)
488                        elif line == '\n':
489                            software.append(pkg + ' ' + version)
490                        else:
491                            continue
492                os.remove('dpkg_list')
493        elif 'windows' in os_type:
494            g.mount_ro(g.list_partitions()[0], '/')
495            hive_file_path = g.inspect_get_windows_software_hive(root)
496            g.download('/' + hive_file_path, 'win_reg')
497            g.umount_all()
498            h = hivex.Hivex('win_reg')
499            key = h.root()
500            key = h.node_get_child (key, 'Microsoft')
501            key = h.node_get_child (key, 'Windows')
502            key = h.node_get_child (key, 'CurrentVersion')
503            key = h.node_get_child (key, 'Uninstall')
504            software += [h.node_name(x) for x in h.node_children(key)]
505            # Just for 64 bit Windows versions, check for 32 bit software.
506            if (os_major_version == 5 and os_minor_version >= 2) or \
507               (os_major_version >= 6):
508                key = h.root()
509                key = h.node_get_child (key, 'Wow6432Node')
510                key = h.node_get_child (key, 'Microsoft')
511                key = h.node_get_child (key, 'Windows')
512                key = h.node_get_child (key, 'CurrentVersion')
513                key = h.node_get_child (key, 'Uninstall')
514                software += [h.node_name(x) for x in h.node_children(key)]
515            os.remove('win_reg')
516
517        return '\n'.join(software)
518
519    def parse_pci(self, path='/usr/share/misc/pci.ids'):
520        data = {}
521        with open(path, 'r') as f:
522            for line in f:
523                if line[0] == '#':
524                    continue
525                elif len(line.strip()) == 0:
526                    continue
527                else:
528                    if line[:2] == '\t\t':
529                        fields = line.strip().split(maxsplit=2)
530                        data[last_vendor][last_device][(fields[0], fields[1])] = fields[2]
531                    elif line[:1] == '\t':
532                        fields = line.strip().split(maxsplit=1)
533                        last_device = fields[0]
534                        data[last_vendor][fields[0]] = {'name': fields[1]}
535                    else:
536                        fields = line.strip().split(maxsplit=1)
537                        if fields[0] == 'ffff':
538                            break
539                        last_vendor = fields[0]
540                        data[fields[0]] = {'name': fields[1]}
541        return data
542
543    def hardware(self, path, ogRest):
544        try:
545            with OgQMP(self.IP, self.VIRTUAL_PORT) as qmp:
546                pci_data = qmp.talk(str({"execute": "query-pci"}))
547                mem_data = qmp.talk(str({"execute": "query-memory-size-summary"}))
548                cpu_data = qmp.talk(str({"execute": "query-cpus-fast"}))
549        except:
550            return
551
552        pci_data = pci_data['return'][0]['devices']
553        pci_list = self.parse_pci()
554        device_names = {}
555        for device in pci_data:
556            vendor_id = hex(device['id']['vendor'])[2:]
557            device_id = hex(device['id']['device'])[2:]
558            subvendor_id = hex(device['id']['subsystem-vendor'])[2:]
559            subdevice_id = hex(device['id']['subsystem'])[2:]
560            description = device['class_info']['desc'].lower()
561            name = ''
562            try:
563                name = pci_list[vendor_id]['name']
564                name += ' ' + pci_list[vendor_id][device_id]['name']
565                name += ' ' + pci_list[vendor_id][device_id][(subvendor_id, subdevice_id)]
566            except KeyError:
567                if vendor_id == '1234':
568                    name = 'VGA Cirrus Logic GD 5446'
569                else:
570                    pass
571
572            if 'usb' in description:
573                device_names['usb'] = name
574            elif 'ide' in description:
575                device_names['ide'] = name
576            elif 'ethernet' in description:
577                device_names['net'] = name
578            elif 'vga' in description:
579                device_names['vga'] = name
580
581            elif 'audio' in description or 'sound' in description:
582                device_names['aud'] = name
583            elif 'dvd' in description:
584                device_names['cdr'] = name
585
586        ram_size = int(mem_data['return']['base-memory']) * 2 ** -20
587        device_names['mem'] = f'QEMU {int(ram_size)}MiB'
588
589        cpu_arch = cpu_data['return'][0]['arch']
590        cpu_target = cpu_data['return'][0]['target']
591        cpu_cores = len(cpu_data['return'])
592        device_names['cpu'] = f'CPU arch:{cpu_arch} target:{cpu_target} ' \
593                              f'cores:{cpu_cores}'
594
595        with open(path, 'w+') as f:
596            f.seek(0)
597            for k, v in device_names.items():
598                f.write(f'{k}={v}\n')
599            f.truncate()
Note: See TracBrowser for help on using the repository browser.