From 1ae7a5e78a6689c6bc7e7fe54ea34f55c78352bb Mon Sep 17 00:00:00 2001 From: Dmitrii Gudov Date: Tue, 24 Feb 2026 21:48:30 +0300 Subject: [PATCH] refactor code --- .gitignore | 2 +- pyproject.toml | 10 +++ src/netdiag/args.py | 22 +++++ src/netdiag/base.py | 28 +++--- src/netdiag/domain/models.py | 17 ++-- src/netdiag/output/d2.py | 137 +++++++++++++++++++++++++++++ src/netdiag/output/file_convert.py | 90 ++++++++++++++----- src/netdiag/output/graphviz.py | 20 ++++- src/netdiag/parse/__init__.py | 5 +- src/netdiag/parse/convert_raw.py | 10 +-- 10 files changed, 286 insertions(+), 55 deletions(-) create mode 100644 src/netdiag/args.py create mode 100644 src/netdiag/output/d2.py diff --git a/.gitignore b/.gitignore index 648167a..f0ef48f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,4 @@ -tmp +data/ # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/pyproject.toml b/pyproject.toml index 7e77bf5..6758f0a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,8 +2,18 @@ name = "netdiag" version = "0.2.0" +[project.scripts] +netdiag = "netdiag.base:run" + [tool.setuptools] package-dir = {"" = "src"} [tool.setuptools.packages.find] where = ["src"] + +[tool.black] +line-length = 88 + +[tool.isort] +profile = "black" +line_length = 88 diff --git a/src/netdiag/args.py b/src/netdiag/args.py new file mode 100644 index 0000000..14e407c --- /dev/null +++ b/src/netdiag/args.py @@ -0,0 +1,22 @@ +import argparse + + +def parse_args(argv=None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Network Diagrams Tool - Generate network diagrams from CSV input" + ) + parser.add_argument( + "-i", + "--input", + type=str, + default="data/input/table.csv", + help="Path to the input CSV file (default: data/input/table.csv)", + ) + parser.add_argument( + "-o", + "--output", + type=str, + default="data/output", + help="Directory for output files (default: data/output)", + ) + return parser.parse_args(args=argv) diff --git a/src/netdiag/base.py b/src/netdiag/base.py index 999c201..8aca63c 100644 --- a/src/netdiag/base.py +++ b/src/netdiag/base.py @@ -1,25 +1,31 @@ +import logging +from pathlib import Path + +from .args import parse_args +from .output.d2 import generate_d2_diagram +from .output.file_convert import make_yaml +from .output.graphviz import generate_diagram from .parse import parse_csv -from .domain.models import Topology from .parse.convert_raw import ( convert_raw_topology, ) -from .output.graphviz import generate_diagram -from .output.file_convert import make_yaml -import logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) -def run(): - raw_devices = parse_csv("data/input/table.csv") - # print(raw_devices) - topology = convert_raw_topology(raw_devices) - # print(topology) +def run(argv: list[str] | None = None) -> None: + args = parse_args(argv) - generate_diagram(topology, "data/output/diagram.png") - make_yaml(topology, "data/output/topology.yaml") + raw_devices = parse_csv(Path(args.input)) + topology = convert_raw_topology(raw_devices) + + generate_diagram(topology, Path(args.output) / "diagram.png") + make_yaml(topology, Path(args.output) / "topology.yaml") + # generate_d2_diagram(topology, Path(args.output) / "diagram.d2") + + logging.info("All tasks completed successfully.") if __name__ == "__main__": diff --git a/src/netdiag/domain/models.py b/src/netdiag/domain/models.py index ed4788a..35e3c0c 100644 --- a/src/netdiag/domain/models.py +++ b/src/netdiag/domain/models.py @@ -1,5 +1,4 @@ -from typing import List, Dict, Any, Optional - +from typing import Any, Dict, List, Optional # =========================== @@ -10,9 +9,7 @@ class Interface: network: Optional[str] default_gateway: Optional[str] # --- - device: Optional[ - "Device" - ] # set by Device.add_interface() when the interface is added to a device + device: "Device" # set by Device.add_interface() when the interface is added to a device def __init__( self, @@ -48,6 +45,7 @@ class VirtualInterface(Interface): class Device: name: str + role: str = "device" interfaces: Dict[str, Interface] def __init__(self, name: str): @@ -66,7 +64,8 @@ class Device: def rm_interface(self, interface: Interface): if interface.name in self.interfaces: del self.interfaces[interface.name] - interface.device = None # clear the device attribute of the interface + # interface.device = None # clear the device attribute of the interface + del interface else: raise ValueError( f"Interface '{interface.name}' not found in device '{self.name}'" @@ -77,15 +76,15 @@ class Device: class Host(Device): - pass + role: str = "host" class Router(Device): - pass + role: str = "router" class Switch(Device): - pass + role: str = "switch" # =========================== diff --git a/src/netdiag/output/d2.py b/src/netdiag/output/d2.py new file mode 100644 index 0000000..3225c71 --- /dev/null +++ b/src/netdiag/output/d2.py @@ -0,0 +1,137 @@ +import shutil +import subprocess +from pathlib import Path + +from ..domain.models import Topology + +#! WIP + +""" +D2 example: + +com_left; com_right + +VLAN 2: { + PC2: |md + # PC2 + 10.0.0.2/24 + | + PC2.shape: rectangle +} + +VLAN 3: { + PC3: |md + # PC3 + 10.0.0.3/24 + | + PC3.shape: rectangle +} + +VLAN 4: { + PC1: |md + # PC1 + 10.0.0.1/24 + | + PC1.shape: rectangle + PC4: |md + # PC4 + 10.0.0.4/24 + | + PC4.shape: rectangle +} + +com_left -- com_right : { + source-arrowhead.label: eth1 + target-arrowhead.label: eth1 +} + +VLAN 2.PC2 -- com_left : { + source-arrowhead.label: eth1 + target-arrowhead.label: eth3 +} +VLAN 4.PC1 -- com_left : { + source-arrowhead.label: eth1 + target-arrowhead.label: eth2 +} +VLAN 4.PC4 -- com_right : { + source-arrowhead.label: eth1 + target-arrowhead.label: eth2 +} +VLAN 3.PC3 -- com_right : { + source-arrowhead.label: eth1 + target-arrowhead.label: eth3 +} + +""" + +# https://d2lang.com/tour/themes/ +THEME_NUMBER = 200 + +# https://icons.terrastruct.com/ +icons = { + "router": "https://icons.terrastruct.com/tech%2Frouter.svg", + "host": "https://icons.terrastruct.com/tech%2F065-monitor-4.svg", + "switch": "https://icons.terrastruct.com/tech%2Fswitch.svg", + "device": "https://icons.terrastruct.com/azure%2FCompute%20Service%20Color%2FVM%2FVM-non-azure.svg", + "vlan": "https://icons.terrastruct.com/azure%2FNetworking%20Service%20Color%2FVirtual%20Networks.svg", +} + + +def _check_d2_installed() -> bool: + return shutil.which("d2") is not None + + +def generate_d2_diagram(topology: Topology, output_path: Path) -> None: + if not _check_d2_installed(): + raise EnvironmentError( + "D2 is not installed or 'd2' command is not found in PATH. Please install D2 to use this feature." + ) + + diagram_path = Path() + picture_path = Path() + + if output_path.is_dir(): + diagram_path = output_path / "diagram.d2" + picture_path = output_path / "diagram.png" + else: + diagram_path = output_path.with_suffix(".d2") + picture_path = output_path.with_suffix(".png") + + file: list[str] = [] + + # some logic + + # --- + + with open(diagram_path, "w", encoding="utf-8") as f: + f.write("\n".join(file)) + + _generate_picture(diagram_path, picture_path) + + +def _generate_picture(diagram: Path, output_path: Path) -> None: + res = subprocess.run( + ["d2", "validate", str(diagram)], + check=True, + capture_output=True, + ) + + if res.returncode != 0: + raise RuntimeError( + f"D2 validation failed (code {res.returncode})\n" + f"Stdout: {res.stdout.decode()}\n" + f"Stderr: {res.stderr.decode()}" + ) + + res = subprocess.run( + ["d2", f"--theme={THEME_NUMBER}", str(diagram), str(output_path)], + check=True, + capture_output=True, + ) + + if res.returncode != 0: + raise RuntimeError( + f"D2 diagram generation failed (code {res.returncode})\n" + f"Stdout: {res.stdout.decode()}\n" + f"Stderr: {res.stderr.decode()}" + ) diff --git a/src/netdiag/output/file_convert.py b/src/netdiag/output/file_convert.py index 9e444d3..e142111 100644 --- a/src/netdiag/output/file_convert.py +++ b/src/netdiag/output/file_convert.py @@ -1,30 +1,72 @@ +from pathlib import Path + import yaml from ..domain.models import Topology -""" -Example: -links: -networks: - - A: [PC1.eth0, PC2.eth0] -meta: - id: host-host.csv - name: host-host.csv -nodes: -- role: host - name: PC1 - interfaces: - - eth0: - - ip: 10.0.12.1/24 - network: A -- role: host - name: PC2 - interfaces: - - eth0: - - ip: 10.0.12.2/24 - network: A -""" +def make_yaml(topology: Topology, output_path: Path) -> None: + data = dict() -def make_yaml(topology: Topology, output_path: str) -> None: - pass + data["meta"] = { + "id": output_path.name, + "name": output_path.name, + } + + data["networks"] = [] + for _, network in topology.networks.items(): + interfaces_with_device = [ + iface for iface in network.interfaces if iface.device is not None + ] + + if len(interfaces_with_device) >= 2: + data["networks"].append( + { + network.name: [ + f"{iface.device.name}.{iface.name}" + for iface in interfaces_with_device + ] + } + ) + + data["nodes"] = [] + + for _, device in topology.devices.items(): + data["nodes"].append( + { + "role": device.role, + "name": device.name, + "interfaces": [ + { + interface_name: [ + { + "ip": ( + interface.ip_address + if interface.ip_address + else None + ), + "network": ( + interface.network if interface.network else None + ), + "gateway": ( + interface.default_gateway + if interface.default_gateway + else None + ), + } + ] + } + for interface_name, interface in device.interfaces.items() + ], + } + ) + + with open(str(output_path), "w", encoding="utf-8") as f: + yaml.safe_dump( + data, + f, + allow_unicode=True, + sort_keys=False, + default_flow_style=False, + indent=2, + ) diff --git a/src/netdiag/output/graphviz.py b/src/netdiag/output/graphviz.py index 67c159b..fea6f11 100644 --- a/src/netdiag/output/graphviz.py +++ b/src/netdiag/output/graphviz.py @@ -1,8 +1,19 @@ from ..domain.models import Topology +from pathlib import Path import graphviz +import shutil -def generate_diagram(topology: Topology, output_path: str) -> None: +def _check_graphviz_installed() -> bool: + return shutil.which("dot") is not None + + +def generate_diagram(topology: Topology, output_path: Path) -> None: + if not _check_graphviz_installed(): + raise EnvironmentError( + "Graphviz is not installed or 'dot' command is not found in PATH. Please install Graphviz to use this feature." + ) + dot = graphviz.Graph(name="Network Topology", format="png", engine="neato") dot.attr(overlap="false", splines="true") @@ -20,7 +31,10 @@ def generate_diagram(topology: Topology, output_path: str) -> None: label = network.name or "" dot.edge(iface_a.device.name, iface_b.device.name, label=label) - output_path = output_path[:-4] if output_path.endswith(".png") else output_path - dot.render(output_path, cleanup=True) + output_path = ( + output_path.with_suffix("") if output_path.suffix == ".png" else output_path + ) + + dot.render(str(output_path), cleanup=True) return diff --git a/src/netdiag/parse/__init__.py b/src/netdiag/parse/__init__.py index 6d95e00..f316b0e 100644 --- a/src/netdiag/parse/__init__.py +++ b/src/netdiag/parse/__init__.py @@ -1,6 +1,7 @@ -from typing import Dict, Any, List import csv import logging +from pathlib import Path +from typing import Any, Dict, List class RawDevices: @@ -22,7 +23,7 @@ class RawDevices: return f"RawDevices(id={self.id}, fields={self.fields})" -def parse_csv(file_path: str, delimiter: str = ",") -> List[RawDevices]: +def parse_csv(file_path: Path, delimiter: str = ",") -> List[RawDevices]: logging.info(f"Parsing CSV file: {file_path}") with open(file_path, mode="r", encoding="utf-8") as csvfile: diff --git a/src/netdiag/parse/convert_raw.py b/src/netdiag/parse/convert_raw.py index 08b3dbe..bbb4476 100644 --- a/src/netdiag/parse/convert_raw.py +++ b/src/netdiag/parse/convert_raw.py @@ -1,16 +1,16 @@ +import logging + from ..domain.models import ( - Interface, - VirtualInterface, Device, Host, + Interface, + Network, Router, Switch, - Network, Topology, + VirtualInterface, ) - from . import RawDevices -import logging name_matching = { "DEVICE_TYPE": "Role",