refactor code

.gitignore

@@ -1,4 +1,4 @@
-tmp
+data/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

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

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)

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():
+def run(argv: list[str] | None = None) -> None:
-    raw_devices = parse_csv("data/input/table.csv")
+    args = parse_args(argv)
-    # print(raw_devices)
-    topology = convert_raw_topology(raw_devices)
-    # print(topology)
 
-    generate_diagram(topology, "data/output/diagram.png")
+    raw_devices = parse_csv(Path(args.input))
-    make_yaml(topology, "data/output/topology.yaml")
+    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__":

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: "Device"  # set by Device.add_interface() when the interface is added to a 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"
 
 
 # ===========================

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()}"
+        )

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:
+    data["meta"] = {
-    pass
+        "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,
+        )

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
+    output_path = (
-    dot.render(output_path, cleanup=True)
+        output_path.with_suffix("") if output_path.suffix == ".png" else output_path
+    )
+
+    dot.render(str(output_path), cleanup=True)
 
     return

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:

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",