diff --git a/daemon/core/emulator/distributed.py b/daemon/core/emulator/distributed.py index a9cba815..5eb3cef6 100644 --- a/daemon/core/emulator/distributed.py +++ b/daemon/core/emulator/distributed.py @@ -7,6 +7,7 @@ import os import threading from collections import OrderedDict from tempfile import NamedTemporaryFile +from typing import TYPE_CHECKING, Callable, Dict, Tuple import netaddr from fabric import Connection @@ -17,6 +18,10 @@ from core.errors import CoreCommandError from core.nodes.interface import GreTap from core.nodes.network import CoreNetwork, CtrlNet +if TYPE_CHECKING: + from core.emulator.session import Session + + LOCK = threading.Lock() CMD_HIDE = True @@ -26,7 +31,7 @@ class DistributedServer: Provides distributed server interactions. """ - def __init__(self, name, host): + def __init__(self, name: str, host: str) -> None: """ Create a DistributedServer instance. @@ -38,7 +43,9 @@ class DistributedServer: self.conn = Connection(host, user="root") self.lock = threading.Lock() - def remote_cmd(self, cmd, env=None, cwd=None, wait=True): + def remote_cmd( + self, cmd: str, env: Dict[str, str] = None, cwd: str = None, wait: bool = True + ) -> str: """ Run command remotely using server connection. @@ -73,7 +80,7 @@ class DistributedServer: stdout, stderr = e.streams_for_display() raise CoreCommandError(e.result.exited, cmd, stdout, stderr) - def remote_put(self, source, destination): + def remote_put(self, source: str, destination: str) -> None: """ Push file to remote server. @@ -84,7 +91,7 @@ class DistributedServer: with self.lock: self.conn.put(source, destination) - def remote_put_temp(self, destination, data): + def remote_put_temp(self, destination: str, data: str) -> None: """ Remote push file contents to a remote server, using a temp file as an intermediate step. @@ -106,11 +113,11 @@ class DistributedController: Provides logic for dealing with remote tunnels and distributed servers. """ - def __init__(self, session): + def __init__(self, session: "Session") -> None: """ Create - :param session: + :param session: session """ self.session = session self.servers = OrderedDict() @@ -119,7 +126,7 @@ class DistributedController: "distributed_address", default=None ) - def add_server(self, name, host): + def add_server(self, name: str, host: str) -> None: """ Add distributed server configuration. @@ -132,7 +139,7 @@ class DistributedController: cmd = f"mkdir -p {self.session.session_dir}" server.remote_cmd(cmd) - def execute(self, func): + def execute(self, func: Callable) -> None: """ Convenience for executing logic against all distributed servers. @@ -143,7 +150,7 @@ class DistributedController: server = self.servers[name] func(server) - def shutdown(self): + def shutdown(self) -> None: """ Shutdown logic for dealing with distributed tunnels and server session directories. @@ -165,7 +172,7 @@ class DistributedController: # clear tunnels self.tunnels.clear() - def start(self): + def start(self) -> None: """ Start distributed network tunnels. @@ -184,7 +191,9 @@ class DistributedController: server = self.servers[name] self.create_gre_tunnel(node, server) - def create_gre_tunnel(self, node, server): + def create_gre_tunnel( + self, node: CoreNetwork, server: DistributedServer + ) -> Tuple[GreTap, GreTap]: """ Create gre tunnel using a pair of gre taps between the local and remote server. @@ -222,7 +231,7 @@ class DistributedController: self.tunnels[key] = tunnel return tunnel - def tunnel_key(self, n1_id, n2_id): + def tunnel_key(self, n1_id: int, n2_id: int) -> int: """ Compute a 32-bit key used to uniquely identify a GRE tunnel. The hash(n1num), hash(n2num) values are used, so node numbers may be @@ -239,7 +248,7 @@ class DistributedController: ) return key & 0xFFFFFFFF - def get_tunnel(self, n1_id, n2_id): + def get_tunnel(self, n1_id: int, n2_id: int) -> Tuple[GreTap, GreTap]: """ Return the GreTap between two nodes if it exists. diff --git a/daemon/core/emulator/emudata.py b/daemon/core/emulator/emudata.py index 8929f72a..acf105eb 100644 --- a/daemon/core/emulator/emudata.py +++ b/daemon/core/emulator/emudata.py @@ -1,40 +1,32 @@ +from typing import List, Optional + import netaddr from core import utils +from core.api.grpc.core_pb2 import LinkOptions from core.emane.nodes import EmaneNet from core.emulator.enumerations import LinkTypes +from core.nodes.base import CoreNetworkBase, CoreNode +from core.nodes.interface import CoreInterface from core.nodes.physical import PhysicalNode class IdGen: - def __init__(self, _id=0): + def __init__(self, _id: int = 0) -> None: self.id = _id - def next(self): + def next(self) -> int: self.id += 1 return self.id -def create_interface(node, network, interface_data): - """ - Create an interface for a node on a network using provided interface data. - - :param node: node to create interface for - :param core.nodes.base.CoreNetworkBase network: network to associate interface with - :param core.emulator.emudata.InterfaceData interface_data: interface data - :return: created interface - """ - node.newnetif( - network, - addrlist=interface_data.get_addresses(), - hwaddr=interface_data.mac, - ifindex=interface_data.id, - ifname=interface_data.name, - ) - return node.netif(interface_data.id) - - -def link_config(network, interface, link_options, devname=None, interface_two=None): +def link_config( + network: CoreNetworkBase, + interface: CoreInterface, + link_options: LinkOptions, + devname: str = None, + interface_two: CoreInterface = None, +) -> None: """ Convenience method for configuring a link, @@ -68,7 +60,7 @@ class NodeOptions: Options for creating and updating nodes within core. """ - def __init__(self, name=None, model="PC", image=None): + def __init__(self, name: str = None, model: str = "PC", image: str = None) -> None: """ Create a NodeOptions object. @@ -93,7 +85,7 @@ class NodeOptions: self.image = image self.emane = None - def set_position(self, x, y): + def set_position(self, x: float, y: float) -> None: """ Convenience method for setting position. @@ -104,7 +96,7 @@ class NodeOptions: self.x = x self.y = y - def set_location(self, lat, lon, alt): + def set_location(self, lat: float, lon: float, alt: float) -> None: """ Convenience method for setting location. @@ -123,7 +115,7 @@ class LinkOptions: Options for creating and updating links within core. """ - def __init__(self, _type=LinkTypes.WIRED): + def __init__(self, _type: LinkTypes = LinkTypes.WIRED) -> None: """ Create a LinkOptions object. @@ -148,12 +140,96 @@ class LinkOptions: self.opaque = None +class InterfaceData: + """ + Convenience class for storing interface data. + """ + + def __init__( + self, + _id: int, + name: str, + mac: str, + ip4: str, + ip4_mask: int, + ip6: str, + ip6_mask: int, + ) -> None: + """ + Creates an InterfaceData object. + + :param int _id: interface id + :param str name: name for interface + :param str mac: mac address + :param str ip4: ipv4 address + :param int ip4_mask: ipv4 bit mask + :param str ip6: ipv6 address + :param int ip6_mask: ipv6 bit mask + """ + self.id = _id + self.name = name + self.mac = mac + self.ip4 = ip4 + self.ip4_mask = ip4_mask + self.ip6 = ip6 + self.ip6_mask = ip6_mask + + def has_ip4(self) -> bool: + """ + Determines if interface has an ip4 address. + + :return: True if has ip4, False otherwise + """ + return all([self.ip4, self.ip4_mask]) + + def has_ip6(self) -> bool: + """ + Determines if interface has an ip6 address. + + :return: True if has ip6, False otherwise + """ + return all([self.ip6, self.ip6_mask]) + + def ip4_address(self) -> Optional[str]: + """ + Retrieve a string representation of the ip4 address and netmask. + + :return: ip4 string or None + """ + if self.has_ip4(): + return f"{self.ip4}/{self.ip4_mask}" + else: + return None + + def ip6_address(self) -> Optional[str]: + """ + Retrieve a string representation of the ip6 address and netmask. + + :return: ip4 string or None + """ + if self.has_ip6(): + return f"{self.ip6}/{self.ip6_mask}" + else: + return None + + def get_addresses(self) -> List[str]: + """ + Returns a list of ip4 and ip6 address when present. + + :return: list of addresses + :rtype: list + """ + ip4 = self.ip4_address() + ip6 = self.ip6_address() + return [i for i in [ip4, ip6] if i] + + class IpPrefixes: """ Convenience class to help generate IP4 and IP6 addresses for nodes within CORE. """ - def __init__(self, ip4_prefix=None, ip6_prefix=None): + def __init__(self, ip4_prefix: str = None, ip6_prefix: str = None) -> None: """ Creates an IpPrefixes object. @@ -171,7 +247,7 @@ class IpPrefixes: if ip6_prefix: self.ip6 = netaddr.IPNetwork(ip6_prefix) - def ip4_address(self, node): + def ip4_address(self, node: CoreNode) -> str: """ Convenience method to return the IP4 address for a node. @@ -183,7 +259,7 @@ class IpPrefixes: raise ValueError("ip4 prefixes have not been set") return str(self.ip4[node.id]) - def ip6_address(self, node): + def ip6_address(self, node: CoreNode) -> str: """ Convenience method to return the IP6 address for a node. @@ -195,7 +271,9 @@ class IpPrefixes: raise ValueError("ip6 prefixes have not been set") return str(self.ip6[node.id]) - def create_interface(self, node, name=None, mac=None): + def create_interface( + self, node: CoreNode, name: str = None, mac: str = None + ) -> InterfaceData: """ Creates interface data for linking nodes, using the nodes unique id for generation, along with a random mac address, unless provided. @@ -239,76 +317,22 @@ class IpPrefixes: ) -class InterfaceData: +def create_interface( + node: CoreNode, network: CoreNetworkBase, interface_data: InterfaceData +): """ - Convenience class for storing interface data. + Create an interface for a node on a network using provided interface data. + + :param node: node to create interface for + :param core.nodes.base.CoreNetworkBase network: network to associate interface with + :param core.emulator.emudata.InterfaceData interface_data: interface data + :return: created interface """ - - def __init__(self, _id, name, mac, ip4, ip4_mask, ip6, ip6_mask): - """ - Creates an InterfaceData object. - - :param int _id: interface id - :param str name: name for interface - :param str mac: mac address - :param str ip4: ipv4 address - :param int ip4_mask: ipv4 bit mask - :param str ip6: ipv6 address - :param int ip6_mask: ipv6 bit mask - """ - self.id = _id - self.name = name - self.mac = mac - self.ip4 = ip4 - self.ip4_mask = ip4_mask - self.ip6 = ip6 - self.ip6_mask = ip6_mask - - def has_ip4(self): - """ - Determines if interface has an ip4 address. - - :return: True if has ip4, False otherwise - """ - return all([self.ip4, self.ip4_mask]) - - def has_ip6(self): - """ - Determines if interface has an ip6 address. - - :return: True if has ip6, False otherwise - """ - return all([self.ip6, self.ip6_mask]) - - def ip4_address(self): - """ - Retrieve a string representation of the ip4 address and netmask. - - :return: ip4 string or None - """ - if self.has_ip4(): - return f"{self.ip4}/{self.ip4_mask}" - else: - return None - - def ip6_address(self): - """ - Retrieve a string representation of the ip6 address and netmask. - - :return: ip4 string or None - """ - if self.has_ip6(): - return f"{self.ip6}/{self.ip6_mask}" - else: - return None - - def get_addresses(self): - """ - Returns a list of ip4 and ip6 address when present. - - :return: list of addresses - :rtype: list - """ - ip4 = self.ip4_address() - ip6 = self.ip6_address() - return [i for i in [ip4, ip6] if i] + node.newnetif( + network, + addrlist=interface_data.get_addresses(), + hwaddr=interface_data.mac, + ifindex=interface_data.id, + ifname=interface_data.name, + ) + return node.netif(interface_data.id) diff --git a/daemon/core/emulator/sessionconfig.py b/daemon/core/emulator/sessionconfig.py index eb38474b..b403b8d6 100644 --- a/daemon/core/emulator/sessionconfig.py +++ b/daemon/core/emulator/sessionconfig.py @@ -1,3 +1,5 @@ +from typing import Any + from core.config import ConfigurableManager, ConfigurableOptions, Configuration from core.emulator.enumerations import ConfigDataTypes, RegisterTlvs from core.plugins.sdt import Sdt @@ -60,29 +62,53 @@ class SessionConfig(ConfigurableManager, ConfigurableOptions): ] config_type = RegisterTlvs.UTILITY.value - def __init__(self): + def __init__(self) -> None: super().__init__() self.set_configs(self.default_values()) def get_config( self, - _id, - node_id=ConfigurableManager._default_node, - config_type=ConfigurableManager._default_type, - default=None, - ): + _id: str, + node_id: int = ConfigurableManager._default_node, + config_type: str = ConfigurableManager._default_type, + default: Any = None, + ) -> str: + """ + Retrieves a specific configuration for a node and configuration type. + + :param str _id: specific configuration to retrieve + :param int node_id: node id to store configuration for + :param str config_type: configuration type to store configuration for + :param default: default value to return when value is not found + :return: configuration value + :rtype str + """ value = super().get_config(_id, node_id, config_type, default) if value == "": value = default return value - def get_config_bool(self, name, default=None): + def get_config_bool(self, name: str, default: Any = None) -> bool: + """ + Get configuration value as a boolean. + + :param name: configuration name + :param default: default value if not found + :return: boolean for configuration value + """ value = self.get_config(name) if value is None: return default return value.lower() == "true" - def get_config_int(self, name, default=None): + def get_config_int(self, name: str, default: Any = None) -> int: + """ + Get configuration value as int. + + :param name: configuration name + :param default: default value if not found + :return: int for configuration value + """ value = self.get_config(name, default=default) if value is not None: value = int(value)