189 lines
6 KiB
Python
189 lines
6 KiB
Python
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 overridden 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
|