dissertation-3-evaluation/runners/runners.py

import ipaddress
import os
from datetime import datetime
from typing import Callable, List, Tuple

import proxmoxer

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_dfs(
        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_dfs(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:
    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,

            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 - 1

        self._internet_bridge = internet_bridge

        self._management_bridge = management_bridge
        self._management_initial_ip = management_initial_ip

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

        self._build_bridges()
        for node in nodes:
            self._build_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:
        self._last_bridge += 1

        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',
                )
                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)
        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:
        self._last_node_id += 1

        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__),
                )
                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)
        return new_id

    def _build_node(self, node: structure.Node):
        # Step 1: connect to Internet bridge with DHCP to install packages
        # Step 2: connect to management bridge for correct setup
        pass

    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