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