dissertation-3-evaluation/runners/runners.py

import ipaddress
import os
import time
from datetime import datetime
from typing import Callable, List, Tuple
from urllib.parse import quote

import proxmoxer
import paramiko

import structure


def check_env(*names: str) -> bool:
    for name in names:
        if name not in os.environ:
            return False
    return True


def bridge_node_search(
        first: structure.Bridge,
        bridge_name_generator: Callable[[structure.Bridge], str],
        node_id_generator: Callable[[structure.Node], int],
) -> Tuple[List[structure.Bridge], List[structure.Node]]:
    bridges: List[structure.Bridge] = []
    nodes: List[structure.Node] = []

    queue: List[structure.Bridge] = [first]
    while len(queue) > 0:
        bridge = queue.pop()
        if bridge.get_name() != '':
            continue

        bridges.append(bridge)
        bridge.set_name(bridge_name_generator(bridge))

        # from this bridge, find all nodes (via all interfaces)
        reachable_nodes: List[structure.Node] = []
        for interface in bridge.get_interfaces():
            node = interface.get_node()
            if node.get_id() is not None:
                continue

            node.set_id(node_id_generator(node))
            reachable_nodes.append(node)
            nodes.append(node)

        # from each node, find all bridges (via all interfaces)
        for node in reachable_nodes:
            for interface in node.get_interfaces():
                bridge = interface.get_bridge()
                if bridge is not None and bridge.get_name() == '':
                    queue.append(bridge)

    return bridges, nodes


class PrintRunner:
    def __init__(self):
        self._last_bridge: int = 0
        self._last_node_id = 0

    def build(self, bridge: structure.Bridge):
        bridges, nodes = bridge_node_search(bridge, lambda _: self.name_bridge(), lambda _: self.id_node())

        print(bridges)
        print(nodes)

    def teardown(self):
        pass

    def name_bridge(self) -> str:
        self._last_bridge += 1
        return 'fake{}'.format(self._last_bridge)

    def id_node(self) -> int:
        self._last_node_id += 1
        return self._last_node_id


class ProxmoxRunner:
    ssh_timeout = 300

    def __init__(
            self,
            host: str,
            node: str,
            user: str,
            token_name: str,
            token_value: str,

            template_id: int,
            initial_vm_id: int,

            internet_bridge: str,

            management_bridge: str,
            management_initial_ip: ipaddress,
            management_gateway: ipaddress,

            verify_ssl: bool = False,
    ):
        self._last_node_id = 0
        self._created_nodes: List[int] = []
        self._last_bridge = 0
        self._created_bridges: List[str] = []

        self._proxmox = proxmoxer.ProxmoxAPI(
            host,
            user=user,
            token_name=token_name,
            token_value=token_value,
            verify_ssl=verify_ssl,
        )

        self._proxmox_node = node

        self._template_id = template_id
        self._initial_vm_id = initial_vm_id

        self._internet_bridge = structure.Bridge()
        self._internet_bridge.set_name(internet_bridge)

        self._management_bridge = structure.Bridge()
        self._management_bridge.set_name(management_bridge)
        self._management_initial_ip = management_initial_ip
        self._management_gateway = management_gateway

        # generate a single use SSH key (we can use any with Proxmox)
        self._private_key = paramiko.RSAKey.generate(3072)
        self._client = paramiko.SSHClient()

    def build(self, bridge: structure.Bridge):
        bridges, nodes = bridge_node_search(bridge, lambda x: self._create_bridge(x), lambda x: self._create_node(x))

        self._build_bridges()

        for node in nodes:
            self._build_node(node)

        # guarantee that setup is not called until all of the nodes are built
        # this means that all will have their final IPs by this point
        for node in nodes:
            self._setup_node(node)

    def _await_task(self, upid, timeout=10):
        t1 = datetime.now()
        while (datetime.now() - t1).seconds < timeout:
            if self._proxmox.nodes(self._proxmox_node).tasks(upid).status.get()['status'] == 'stopped':
                return
        raise TimeoutError

    def _create_bridge(self, bridge: structure.Bridge) -> str:
        while True:
            try:
                self._proxmox.nodes(self._proxmox_node).network.post(
                    iface='vmbr{}'.format(self._last_bridge),
                    type='bridge',
                    autostart=1,
                    comments='Automatically created by Python evaluation',
                )
                self._last_bridge += 1
                break
            except proxmoxer.core.ResourceException as e:
                if 'interface already exists' in str(e):
                    self._last_bridge += 1
                else:
                    raise e

        bridge_name = 'vmbr{}'.format(self._last_bridge - 1)
        self._created_bridges.append(bridge_name)
        return bridge_name

    def _build_bridges(self):
        network_task = self._proxmox.nodes(self._proxmox_node).network.put()
        self._await_task(network_task)

    def _create_node(self, node: structure.Node) -> int:
        while True:
            try:
                clone_task = self._proxmox.nodes(self._proxmox_node).qemu(self._template_id).clone.post(
                    newid=self._initial_vm_id + self._last_node_id,
                    name='Diss-{}-Testing'.format(node.__class__.__name__),
                )
                self._last_node_id += 1
                break
            except proxmoxer.core.ResourceException as e:
                if 'config file already exists' in str(e):
                    self._last_node_id += 1
                else:
                    raise e
        self._await_task(clone_task)

        new_id = self._initial_vm_id + self._last_node_id
        self._created_nodes.append(new_id - 1)
        return new_id - 1

    def _open_ssh(self, node: structure.Node, interface: structure.Interface = None):
        if interface is None:
            for iface in node.get_interfaces():
                if iface.get_method() == structure.IpMethod.Management:
                    interface = iface
                    break
            if interface is None:
                raise RuntimeError('no management interface available')

        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy)

        t1 = datetime.now()
        while (datetime.now() - t1).seconds < ProxmoxRunner.ssh_timeout:
            try:
                client.connect(
                    hostname=str(interface.get_address()),
                    username='python',
                    pkey=self._private_key,
                    banner_timeout=15,
                )
                client.set_missing_host_key_policy(paramiko.RejectPolicy)
                break
            except (paramiko.ssh_exception.AuthenticationException, paramiko.ssh_exception.NoValidConnectionsError):
                time.sleep(10)

        node.client = client

    def _close_ssh(self, node: structure.Node):
        node.client.close()
        del node.client

    def ssh(self, node: structure.Node, command: str, error_stderr=False, error_stdout=False) -> int:
        chan = node.client.get_transport().open_session()

        chan.exec_command(command)
        exit_status = chan.recv_exit_status()

        if exit_status != 0:
            if error_stderr and error_stdout:
                raise Exception(
                    'stdout:\n{}\n\nstderr:\n{}\n'.format(chan.recv(2048).decode(), chan.recv_stderr(2048).decode()))
            if error_stderr:
                raise Exception(chan.recv_stderr(2048).decode())
            if error_stdout:
                raise Exception(chan.recv(2048).decode())

        return exit_status

    def _build_node(self, node: structure.Node):
        # Step 1: Configure access
        self._proxmox.nodes(self._proxmox_node).qemu(node.get_id()).config.put(
            ciuser='python',
            sshkeys=quote('ssh-rsa ' + self._private_key.get_base64(), ''),
            cores=node.get_core_count(),
            sockets=1,
            memory=node.get_memory_mb(),
        )

        # Step 2: connect to Internet bridge with DHCP to install packages
        if node.get_internet_setup() is not None:
            interfaces = node.get_interfaces()
            internet_interface = structure.Interface(structure.IpMethod.Dhcp4)
            internet_interface.set_bridge(self._internet_bridge)
            temp_interfaces = [internet_interface, interfaces[len(interfaces)-1]]

            self._setup_node_interfaces(node, temp_interfaces)

            start_task = self._proxmox.nodes(self._proxmox_node).qemu(node.get_id()).status.start.post()
            self._await_task(start_task)

            self._open_ssh(node)
            self.ssh(node, node.get_internet_setup(), error_stdout=True, error_stderr=True)
            self._close_ssh(node)

            stop_task = self._proxmox.nodes(self._proxmox_node).qemu(node.get_id()).status.shutdown.post()
            self._await_task(stop_task)

        # Step 3: connect to management bridge for final setup
        self._setup_node_interfaces(node)

        start_task = self._proxmox.nodes(self._proxmox_node).qemu(node.get_id()).status.start.post()
        self._await_task(start_task)

        self._open_ssh(node)

    def _setup_node_interfaces(self, node: structure.Node, interfaces: List[structure.Interface] = None):
        if interfaces is None:
            interfaces = node.get_interfaces()

        kwargs = dict()
        for i in range(len(interfaces)):
            interface = interfaces[i]
            method = interface.get_method()

            if method == structure.IpMethod.Manual:
                pass
            elif method == structure.IpMethod.Management:
                interface.set_bridge(self._management_bridge)
                addr = self._management_initial_ip + node.get_id() - self._initial_vm_id

                kwargs['ipconfig{}'.format(i)] = 'ip={}/24,gw={}'.format(addr, self._management_gateway)
                interface.set_address(addr)
            elif method == structure.IpMethod.Auto4:
                bridge = interface.get_bridge()
                addr = bridge.get_ip_address()

                kwargs['ipconfig{}'.format(i)] = 'ip={}/{}'.format(addr, bridge.netmask)
                interface.set_address(addr)
            elif method == structure.IpMethod.Dhcp4:
                kwargs['ipconfig{}'.format(i)] = 'ip=dhcp'
            else:
                raise RuntimeError('not implemented')

            kwargs['net{}'.format(i)] = 'model=virtio,bridge={}'.format(interface.get_bridge().get_name())

        self._proxmox.nodes(self._proxmox_node).qemu(node.get_id()).config.put(**kwargs)

    def _setup_node(self, node: structure.Node):
        if node.get_setup() is not None:
            self.ssh(node, node.get_setup(), error_stdout=True, error_stderr=True)

    def teardown(self):
        for node in self._created_nodes:
            stop_task = self._proxmox.nodes(self._proxmox_node).qemu(node).status.stop.post()
            self._await_task(stop_task)
            delete_task = self._proxmox.nodes(self._proxmox_node).qemu(node).delete()
            self._await_task(delete_task)

        for bridge in self._created_bridges:
            self._proxmox.nodes(self._proxmox_node).network(bridge).delete()

        network_task = self._proxmox.nodes(self._proxmox_node).network.put()
        self._await_task(network_task)

        self._created_nodes = []
        self._last_node_id = 0

        self._created_bridges = []
        self._last_bridge = 0