refactor code

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:
-    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,
+        )

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