import logging from typing import TYPE_CHECKING, List, Optional, Set from PIL.ImageTk import PhotoImage from core.api.grpc.wrappers import Node, NodeType from core.gui import images from core.gui.appconfig import CustomNode, GuiConfig from core.gui.images import ImageEnum logger = logging.getLogger(__name__) if TYPE_CHECKING: from core.gui.app import Application NODES: List["NodeDraw"] = [] NETWORK_NODES: List["NodeDraw"] = [] NODE_ICONS = {} CONTAINER_NODES: Set[NodeType] = {NodeType.DEFAULT, NodeType.DOCKER, NodeType.LXC} IMAGE_NODES: Set[NodeType] = {NodeType.DOCKER, NodeType.LXC} WIRELESS_NODES: Set[NodeType] = { NodeType.WIRELESS_LAN, NodeType.EMANE, NodeType.WIRELESS, } RJ45_NODES: Set[NodeType] = {NodeType.RJ45} BRIDGE_NODES: Set[NodeType] = {NodeType.HUB, NodeType.SWITCH} IGNORE_NODES: Set[NodeType] = {NodeType.CONTROL_NET} MOBILITY_NODES: Set[NodeType] = {NodeType.WIRELESS_LAN, NodeType.EMANE} NODE_MODELS: Set[str] = {"router", "PC", "mdr", "prouter"} ROUTER_NODES: Set[str] = {"router", "mdr"} ANTENNA_ICON: Optional[PhotoImage] = None def setup() -> None: global ANTENNA_ICON nodes = [ (ImageEnum.PC, NodeType.DEFAULT, "PC", "PC"), (ImageEnum.MDR, NodeType.DEFAULT, "MDR", "mdr"), (ImageEnum.ROUTER, NodeType.DEFAULT, "Router", "router"), (ImageEnum.PROUTER, NodeType.DEFAULT, "PRouter", "prouter"), (ImageEnum.DOCKER, NodeType.DOCKER, "Docker", None), (ImageEnum.LXC, NodeType.LXC, "LXC", None), ] for image_enum, node_type, label, model in nodes: node_draw = NodeDraw.from_setup(image_enum, node_type, label, model) NODES.append(node_draw) NODE_ICONS[(node_type, model)] = node_draw.image network_nodes = [ (ImageEnum.HUB, NodeType.HUB, "Hub"), (ImageEnum.SWITCH, NodeType.SWITCH, "Switch"), (ImageEnum.WLAN, NodeType.WIRELESS_LAN, "WLAN"), (ImageEnum.WIRELESS, NodeType.WIRELESS, "Wireless"), (ImageEnum.EMANE, NodeType.EMANE, "EMANE"), (ImageEnum.RJ45, NodeType.RJ45, "RJ45"), (ImageEnum.TUNNEL, NodeType.TUNNEL, "Tunnel"), ] for image_enum, node_type, label in network_nodes: node_draw = NodeDraw.from_setup(image_enum, node_type, label) NETWORK_NODES.append(node_draw) NODE_ICONS[(node_type, None)] = node_draw.image ANTENNA_ICON = images.from_enum(ImageEnum.ANTENNA, width=images.ANTENNA_SIZE) def is_bridge(node: Node) -> bool: return node.type in BRIDGE_NODES def is_mobility(node: Node) -> bool: return node.type in MOBILITY_NODES def is_router(node: Node) -> bool: return is_model(node) and node.model in ROUTER_NODES def should_ignore(node: Node) -> bool: return node.type in IGNORE_NODES def is_container(node: Node) -> bool: return node.type in CONTAINER_NODES def is_model(node: Node) -> bool: return node.type == NodeType.DEFAULT def has_image(node_type: NodeType) -> bool: return node_type in IMAGE_NODES def is_wireless(node: Node) -> bool: return node.type in WIRELESS_NODES def is_rj45(node: Node) -> bool: return node.type in RJ45_NODES def is_custom(node: Node) -> bool: return is_model(node) and node.model not in NODE_MODELS def is_iface_node(node: Node) -> bool: return is_container(node) or is_bridge(node) def get_custom_services(gui_config: GuiConfig, name: str) -> List[str]: for custom_node in gui_config.nodes: if custom_node.name == name: return custom_node.services return [] def _get_custom_file(config: GuiConfig, name: str) -> Optional[str]: for custom_node in config.nodes: if custom_node.name == name: return custom_node.image return None def get_icon(node: Node, app: "Application") -> PhotoImage: scale = app.app_scale image = None # node icon was overriden with a specific value if node.icon: try: image = images.from_file(node.icon, width=images.NODE_SIZE, scale=scale) except OSError: logger.error("invalid icon: %s", node.icon) # custom node elif is_custom(node): image_file = _get_custom_file(app.guiconfig, node.model) logger.info("custom node file: %s", image_file) if image_file: image = images.from_file(image_file, width=images.NODE_SIZE, scale=scale) # built in node else: image = images.from_node(node, scale=scale) # default image, if everything above fails if not image: image = images.from_enum( ImageEnum.EDITNODE, width=images.NODE_SIZE, scale=scale ) return image class NodeDraw: def __init__(self) -> None: self.custom: bool = False self.image: Optional[PhotoImage] = None self.image_enum: Optional[ImageEnum] = None self.image_file: Optional[str] = None self.node_type: Optional[NodeType] = None self.model: Optional[str] = None self.services: Set[str] = set() self.label: Optional[str] = None @classmethod def from_setup( cls, image_enum: ImageEnum, node_type: NodeType, label: str, model: str = None, tooltip: str = None, ) -> "NodeDraw": node_draw = NodeDraw() node_draw.image_enum = image_enum node_draw.image = images.from_enum(image_enum, width=images.NODE_SIZE) node_draw.node_type = node_type node_draw.label = label node_draw.model = model node_draw.tooltip = tooltip return node_draw @classmethod def from_custom(cls, custom_node: CustomNode) -> "NodeDraw": node_draw = NodeDraw() node_draw.custom = True node_draw.image_file = custom_node.image node_draw.image = images.from_file(custom_node.image, width=images.NODE_SIZE) node_draw.node_type = NodeType.DEFAULT node_draw.services = set(custom_node.services) node_draw.label = custom_node.name node_draw.model = custom_node.name node_draw.tooltip = custom_node.name return node_draw