prod ready

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 David Romero
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

bin/d2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:470c99eedfc01adf3a9162fbc24e451dc0e5ea91b687d38ebe3ac6547e0cf412
-size 46341847

ipmentor/core.py

@@ -1,259 +0,0 @@
-"""
-Core IPv4 networking functions for IPMentor.
-"""
-
-import ipaddress
-import math
-from typing import Dict, List, Tuple
-
-
-def ip_to_binary(ip_str: str) -> str:
-    """Convert IP address to binary format with dots."""
-    try:
-        ip = ipaddress.IPv4Address(ip_str)
-        binary = format(int(ip), '032b')
-        return f"{binary[:8]}.{binary[8:16]}.{binary[16:24]}.{binary[24:]}"
-    except:
-        return "Invalid IP"
-
-
-def binary_to_ip(binary_str: str) -> str:
-    """Convert binary IP to decimal format."""
-    try:
-        binary_clean = binary_str.replace('.', '').replace(' ', '')
-        if len(binary_clean) != 32 or not all(c in '01' for c in binary_clean):
-            return "Invalid Binary"
-        ip_int = int(binary_clean, 2)
-        return str(ipaddress.IPv4Address(ip_int))
-    except:
-        return "Invalid Binary"
-
-
-def parse_subnet_mask(mask_str: str) -> Tuple[str, int]:
-    """Parse subnet mask from various formats."""
-    mask_str = mask_str.strip()
-    
-    if mask_str.startswith('/'):
-        cidr = int(mask_str[1:])
-    elif '.' in mask_str:
-        mask_ip = ipaddress.IPv4Address(mask_str)
-        cidr = bin(int(mask_ip)).count('1')
-    else:
-        cidr = int(mask_str)
-    
-    if not 0 <= cidr <= 32:
-        raise ValueError("Invalid CIDR")
-    
-    mask_ip = ipaddress.IPv4Network(f"0.0.0.0/{cidr}").netmask
-    return str(mask_ip), cidr
-
-
-def analyze_ip(ip: str, subnet_mask: str) -> Dict:
-    """Analyze IP address with subnet mask."""
-    try:
-        # Handle binary IP
-        if '.' in ip and all(c in '01.' for c in ip.replace('.', '')):
-            ip = binary_to_ip(ip)
-            if ip == "Invalid Binary":
-                raise ValueError("Invalid binary IP")
-        
-        # Parse mask
-        mask_decimal, cidr = parse_subnet_mask(subnet_mask)
-        
-        # Create network
-        network = ipaddress.IPv4Network(f"{ip}/{cidr}", strict=False)
-        
-        # Calculate hosts
-        if cidr < 31:
-            total_hosts = 2 ** (32 - cidr) - 2
-            first_host = str(network.network_address + 1)
-            last_host = str(network.broadcast_address - 1)
-        elif cidr == 31:
-            total_hosts = 2
-            first_host = str(network.network_address)
-            last_host = str(network.broadcast_address)
-        else:
-            total_hosts = 1
-            first_host = str(network.network_address)
-            last_host = str(network.network_address)
-        
-        return {
-            "ip_decimal": ip,
-            "ip_binary": ip_to_binary(ip),
-            "subnet_mask_decimal": mask_decimal,
-            "subnet_mask_binary": ip_to_binary(mask_decimal),
-            "subnet_mask_cidr": f"/{cidr}",
-            "network_address": str(network.network_address),
-            "broadcast_address": str(network.broadcast_address),
-            "first_host": first_host,
-            "last_host": last_host,
-            "total_hosts": total_hosts
-        }
-    
-    except Exception as e:
-        return {"error": str(e)}
-
-
-def calculate_subnets(network: str, number: int, method: str, hosts_list: str = "") -> Dict:
-    """Calculate subnets using different methods."""
-    try:
-        base_network = ipaddress.IPv4Network(network, strict=False)
-        base_cidr = base_network.prefixlen
-        
-        if method == "max_subnets":
-            # Calculate subnets needed
-            bits_needed = math.ceil(math.log2(number))
-            new_cidr = base_cidr + bits_needed
-            
-            if new_cidr > 32:
-                raise ValueError("Too many subnets requested")
-            
-            subnets = list(base_network.subnets(new_prefix=new_cidr))
-            hosts_per_subnet = 2 ** (32 - new_cidr) - 2 if new_cidr < 31 else (2 if new_cidr == 31 else 1)
-            
-            subnet_list = []
-            for i, subnet in enumerate(subnets[:number]):
-                subnet_list.append({
-                    "subnet": str(subnet),
-                    "network": str(subnet.network_address),
-                    "broadcast": str(subnet.broadcast_address),
-                    "first_host": str(subnet.network_address + 1) if hosts_per_subnet > 1 else str(subnet.network_address),
-                    "last_host": str(subnet.broadcast_address - 1) if hosts_per_subnet > 1 else str(subnet.broadcast_address),
-                    "hosts": hosts_per_subnet
-                })
-            
-            return {
-                "method": "Max Subnets",
-                "subnets": subnet_list,
-                "bits_borrowed": bits_needed,
-                "hosts_per_subnet": hosts_per_subnet,
-                "total_subnets": len(subnets)
-            }
-        
-        elif method == "max_hosts_per_subnet":
-            # Calculate CIDR for hosts
-            if number <= 2:
-                bits_for_hosts = 1 if number == 2 else 0
-            else:
-                bits_for_hosts = math.ceil(math.log2(number + 2))
-            
-            new_cidr = 32 - bits_for_hosts
-            
-            if new_cidr < base_cidr:
-                raise ValueError("Too many hosts requested")
-            
-            subnets = list(base_network.subnets(new_prefix=new_cidr))
-            actual_hosts = 2 ** bits_for_hosts - 2 if new_cidr < 31 else (2 if new_cidr == 31 else 1)
-            
-            subnet_list = []
-            for subnet in subnets:
-                subnet_list.append({
-                    "subnet": str(subnet),
-                    "network": str(subnet.network_address),
-                    "broadcast": str(subnet.broadcast_address),
-                    "first_host": str(subnet.network_address + 1) if actual_hosts > 1 else str(subnet.network_address),
-                    "last_host": str(subnet.broadcast_address - 1) if actual_hosts > 1 else str(subnet.broadcast_address),
-                    "hosts": actual_hosts
-                })
-            
-            return {
-                "method": "Max Hosts per Subnet",
-                "subnets": subnet_list,
-                "hosts_per_subnet": actual_hosts,
-                "total_subnets": len(subnets)
-            }
-        
-        elif method == "vlsm":
-            # Parse hosts requirements
-            hosts_requirements = [int(x.strip()) for x in hosts_list.split(',')]
-            
-            if len(hosts_requirements) != number:
-                raise ValueError(f"Need exactly {number} host values")
-            
-            # Sort largest first for optimal allocation
-            sorted_reqs = sorted(enumerate(hosts_requirements), key=lambda x: x[1], reverse=True)
-            
-            subnets = []
-            available_networks = [base_network]
-            
-            for original_idx, hosts_needed in sorted_reqs:
-                # Find required CIDR
-                if hosts_needed == 1:
-                    required_cidr = 32
-                elif hosts_needed == 2:
-                    required_cidr = 31
-                else:
-                    # Calculate bits needed for hosts + network + broadcast
-                    bits_for_hosts = math.ceil(math.log2(hosts_needed + 2))
-                    required_cidr = 32 - bits_for_hosts
-                
-                # Find a suitable available network
-                allocated = False
-                for i, avail_net in enumerate(available_networks):
-                    # Check if we can fit the required subnet in this available network
-                    if required_cidr >= avail_net.prefixlen:
-                        # Allocate the first subnet of the required size
-                        try:
-                            allocated_subnet = list(avail_net.subnets(new_prefix=required_cidr))[0]
-                        except ValueError:
-                            # Cannot create subnet of this size
-                            continue
-                        
-                        # Calculate actual host capacity
-                        if required_cidr == 32:
-                            actual_hosts = 1
-                            first_host = str(allocated_subnet.network_address)
-                            last_host = str(allocated_subnet.network_address)
-                        elif required_cidr == 31:
-                            actual_hosts = 2
-                            first_host = str(allocated_subnet.network_address)
-                            last_host = str(allocated_subnet.broadcast_address)
-                        else:
-                            actual_hosts = 2 ** (32 - required_cidr) - 2
-                            first_host = str(allocated_subnet.network_address + 1)
-                            last_host = str(allocated_subnet.broadcast_address - 1)
-                        
-                        subnets.append({
-                            "subnet": str(allocated_subnet),
-                            "network": str(allocated_subnet.network_address),
-                            "broadcast": str(allocated_subnet.broadcast_address),
-                            "first_host": first_host,
-                            "last_host": last_host,
-                            "hosts": actual_hosts,
-                            "hosts_requested": hosts_needed,
-                            "original_order": original_idx + 1
-                        })
-                        
-                        # Remove the used network
-                        available_networks.pop(i)
-                        
-                        # Add remaining available spaces by splitting the original network
-                        remaining_subnets = []
-                        for subnet in avail_net.subnets(new_prefix=required_cidr):
-                            if subnet != allocated_subnet:
-                                remaining_subnets.append(subnet)
-                        
-                        # Add the remaining subnets back to available networks
-                        available_networks.extend(remaining_subnets)
-                        
-                        allocated = True
-                        break
-                
-                if not allocated:
-                    raise ValueError(f"Cannot allocate subnet for {hosts_needed} hosts")
-            
-            # Sort back to original order
-            subnets.sort(key=lambda x: x["original_order"])
-            
-            return {
-                "method": "VLSM",
-                "subnets": subnets,
-                "total_hosts_requested": sum(hosts_requirements),
-                "total_hosts_allocated": sum(s["hosts"] for s in subnets)
-            }
-        
-        else:
-            raise ValueError("Invalid method")
-    
-    except Exception as e:
-        return {"error": str(e)}

ipmentor/tools.py

@@ -1,5 +1,5 @@
 """
-MCP tools for IPMentor.
+Tools for IPMentor.
 """
 
 import json
@@ -8,9 +8,262 @@ import subprocess
 import tempfile
 import shutil
 import os
+import ipaddress
+import math
 from pathlib import Path
-from typing import List
-from .core import analyze_ip, calculate_subnets
+from typing import List, Dict, Tuple
+
+
+def ip_to_binary(ip_str: str) -> str:
+    """Convert IP address to binary format with dots."""
+    try:
+        ip = ipaddress.IPv4Address(ip_str)
+        binary = format(int(ip), '032b')
+        return f"{binary[:8]}.{binary[8:16]}.{binary[16:24]}.{binary[24:]}"
+    except:
+        return "Invalid IP"
+
+
+def binary_to_ip(binary_str: str) -> str:
+    """Convert binary IP to decimal format."""
+    try:
+        binary_clean = binary_str.replace('.', '').replace(' ', '')
+        if len(binary_clean) != 32 or not all(c in '01' for c in binary_clean):
+            return "Invalid Binary"
+        ip_int = int(binary_clean, 2)
+        return str(ipaddress.IPv4Address(ip_int))
+    except:
+        return "Invalid Binary"
+
+
+def parse_subnet_mask(mask_str: str) -> Tuple[str, int]:
+    """Parse subnet mask from various formats."""
+    mask_str = mask_str.strip()
+    
+    if mask_str.startswith('/'):
+        cidr = int(mask_str[1:])
+    elif '.' in mask_str:
+        mask_ip = ipaddress.IPv4Address(mask_str)
+        cidr = bin(int(mask_ip)).count('1')
+    else:
+        cidr = int(mask_str)
+    
+    if not 0 <= cidr <= 32:
+        raise ValueError("Invalid CIDR")
+    
+    mask_ip = ipaddress.IPv4Network(f"0.0.0.0/{cidr}").netmask
+    return str(mask_ip), cidr
+
+
+def analyze_ip(ip: str, subnet_mask: str) -> Dict:
+    """Analyze IP address with subnet mask."""
+    try:
+        # Handle binary IP
+        if '.' in ip and all(c in '01.' for c in ip.replace('.', '')):
+            ip = binary_to_ip(ip)
+            if ip == "Invalid Binary":
+                raise ValueError("Invalid binary IP")
+        
+        # Parse mask
+        mask_decimal, cidr = parse_subnet_mask(subnet_mask)
+        
+        # Create network
+        network = ipaddress.IPv4Network(f"{ip}/{cidr}", strict=False)
+        
+        # Calculate hosts
+        if cidr < 31:
+            total_hosts = 2 ** (32 - cidr) - 2
+            first_host = str(network.network_address + 1)
+            last_host = str(network.broadcast_address - 1)
+        elif cidr == 31:
+            total_hosts = 2
+            first_host = str(network.network_address)
+            last_host = str(network.broadcast_address)
+        else:
+            total_hosts = 1
+            first_host = str(network.network_address)
+            last_host = str(network.network_address)
+        
+        return {
+            "ip_decimal": ip,
+            "ip_binary": ip_to_binary(ip),
+            "subnet_mask_decimal": mask_decimal,
+            "subnet_mask_binary": ip_to_binary(mask_decimal),
+            "subnet_mask_cidr": f"/{cidr}",
+            "network_address": str(network.network_address),
+            "broadcast_address": str(network.broadcast_address),
+            "first_host": first_host,
+            "last_host": last_host,
+            "total_hosts": total_hosts
+        }
+    
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def calculate_subnets(network: str, number: int, method: str, hosts_list: str = "") -> Dict:
+    """Calculate subnets using different methods."""
+    try:
+        base_network = ipaddress.IPv4Network(network, strict=False)
+        base_cidr = base_network.prefixlen
+        
+        if method == "max_subnets":
+            # Calculate subnets needed
+            bits_needed = math.ceil(math.log2(number))
+            new_cidr = base_cidr + bits_needed
+            
+            if new_cidr > 32:
+                raise ValueError("Too many subnets requested")
+            
+            subnets = list(base_network.subnets(new_prefix=new_cidr))
+            hosts_per_subnet = 2 ** (32 - new_cidr) - 2 if new_cidr < 31 else (2 if new_cidr == 31 else 1)
+            
+            subnet_list = []
+            for i, subnet in enumerate(subnets[:number]):
+                subnet_list.append({
+                    "subnet": str(subnet),
+                    "network": str(subnet.network_address),
+                    "broadcast": str(subnet.broadcast_address),
+                    "first_host": str(subnet.network_address + 1) if hosts_per_subnet > 1 else str(subnet.network_address),
+                    "last_host": str(subnet.broadcast_address - 1) if hosts_per_subnet > 1 else str(subnet.broadcast_address),
+                    "hosts": hosts_per_subnet
+                })
+            
+            return {
+                "method": "Max Subnets",
+                "subnets": subnet_list,
+                "bits_borrowed": bits_needed,
+                "hosts_per_subnet": hosts_per_subnet,
+                "total_subnets": len(subnets)
+            }
+        
+        elif method == "max_hosts_per_subnet":
+            # Calculate CIDR for hosts
+            if number <= 2:
+                bits_for_hosts = 1 if number == 2 else 0
+            else:
+                bits_for_hosts = math.ceil(math.log2(number + 2))
+            
+            new_cidr = 32 - bits_for_hosts
+            
+            if new_cidr < base_cidr:
+                raise ValueError("Too many hosts requested")
+            
+            subnets = list(base_network.subnets(new_prefix=new_cidr))
+            actual_hosts = 2 ** bits_for_hosts - 2 if new_cidr < 31 else (2 if new_cidr == 31 else 1)
+            
+            subnet_list = []
+            for subnet in subnets:
+                subnet_list.append({
+                    "subnet": str(subnet),
+                    "network": str(subnet.network_address),
+                    "broadcast": str(subnet.broadcast_address),
+                    "first_host": str(subnet.network_address + 1) if actual_hosts > 1 else str(subnet.network_address),
+                    "last_host": str(subnet.broadcast_address - 1) if actual_hosts > 1 else str(subnet.broadcast_address),
+                    "hosts": actual_hosts
+                })
+            
+            return {
+                "method": "Max Hosts per Subnet",
+                "subnets": subnet_list,
+                "hosts_per_subnet": actual_hosts,
+                "total_subnets": len(subnets)
+            }
+        
+        elif method == "vlsm":
+            # Parse hosts requirements
+            hosts_requirements = [int(x.strip()) for x in hosts_list.split(',')]
+            
+            if len(hosts_requirements) != number:
+                raise ValueError(f"Need exactly {number} host values")
+            
+            # Sort largest first for optimal allocation
+            sorted_reqs = sorted(enumerate(hosts_requirements), key=lambda x: x[1], reverse=True)
+            
+            subnets = []
+            available_networks = [base_network]
+            
+            for original_idx, hosts_needed in sorted_reqs:
+                # Find required CIDR
+                if hosts_needed == 1:
+                    required_cidr = 32
+                elif hosts_needed == 2:
+                    required_cidr = 31
+                else:
+                    # Calculate bits needed for hosts + network + broadcast
+                    bits_for_hosts = math.ceil(math.log2(hosts_needed + 2))
+                    required_cidr = 32 - bits_for_hosts
+                
+                # Find a suitable available network
+                allocated = False
+                for i, avail_net in enumerate(available_networks):
+                    # Check if we can fit the required subnet in this available network
+                    if required_cidr >= avail_net.prefixlen:
+                        # Allocate the first subnet of the required size
+                        try:
+                            allocated_subnet = list(avail_net.subnets(new_prefix=required_cidr))[0]
+                        except ValueError:
+                            # Cannot create subnet of this size
+                            continue
+                        
+                        # Calculate actual host capacity
+                        if required_cidr == 32:
+                            actual_hosts = 1
+                            first_host = str(allocated_subnet.network_address)
+                            last_host = str(allocated_subnet.network_address)
+                        elif required_cidr == 31:
+                            actual_hosts = 2
+                            first_host = str(allocated_subnet.network_address)
+                            last_host = str(allocated_subnet.broadcast_address)
+                        else:
+                            actual_hosts = 2 ** (32 - required_cidr) - 2
+                            first_host = str(allocated_subnet.network_address + 1)
+                            last_host = str(allocated_subnet.broadcast_address - 1)
+                        
+                        subnets.append({
+                            "subnet": str(allocated_subnet),
+                            "network": str(allocated_subnet.network_address),
+                            "broadcast": str(allocated_subnet.broadcast_address),
+                            "first_host": first_host,
+                            "last_host": last_host,
+                            "hosts": actual_hosts,
+                            "hosts_requested": hosts_needed,
+                            "original_order": original_idx + 1
+                        })
+                        
+                        # Remove the used network
+                        available_networks.pop(i)
+                        
+                        # Add remaining available spaces by splitting the original network
+                        remaining_subnets = []
+                        for subnet in avail_net.subnets(new_prefix=required_cidr):
+                            if subnet != allocated_subnet:
+                                remaining_subnets.append(subnet)
+                        
+                        # Add the remaining subnets back to available networks
+                        available_networks.extend(remaining_subnets)
+                        
+                        allocated = True
+                        break
+                
+                if not allocated:
+                    raise ValueError(f"Cannot allocate subnet for {hosts_needed} hosts")
+            
+            # Sort back to original order
+            subnets.sort(key=lambda x: x["original_order"])
+            
+            return {
+                "method": "VLSM",
+                "subnets": subnets,
+                "total_hosts_requested": sum(hosts_requirements),
+                "total_hosts_allocated": sum(s["hosts"] for s in subnets)
+            }
+        
+        else:
+            raise ValueError("Invalid method")
+    
+    except Exception as e:
+        return {"error": str(e)}
 
 
 def ip_info(ip: str, subnet_mask: str) -> str:
@@ -150,29 +403,21 @@ def _style_d2_diagram(diagram: str) -> str:
     return "\n".join(styled)
 
 
-def _export_to_svg(d2_diagram: str, output_path: str | Path = "diagram.svg") -> Path:
-    """Export D2 diagram to SVG using d2 CLI."""
-    # Try to find d2 binary in different locations
-    d2_binary = None
-    
-    # First try local binary relative to project root
-    script_dir = Path(__file__).parent.parent  # Go up from ipmentor/tools.py to project root
-    local_d2 = script_dir / "bin" / "d2"
-    if local_d2.exists() and local_d2.is_file():
-        d2_binary = str(local_d2)
-    
-    # Fallback to system d2
-    elif shutil.which("d2") is not None:
-        d2_binary = "d2"
-    
-    if d2_binary is None:
+def _export_to_image(d2_diagram: str, output_path: str | Path = "diagram.png", format: str = "png") -> Path:
+    """Export D2 diagram to PNG or SVG using d2 CLI."""
+    if shutil.which("d2") is None:
         raise RuntimeError(
             "D2 executable not found.\n"
-            "Expected at bin/d2 or in system PATH.\n"
-            "Install from https://d2lang.com/tour/installation"
+            "Install it from https://d2lang.com/tour/installation "
+            "or with Homebrew / chocolatey / scoop."
         )
 
+    # Ensure output path has correct extension
     output_path = Path(output_path).expanduser().resolve()
+    if format == "svg" and not output_path.suffix == ".svg":
+        output_path = output_path.with_suffix(".svg")
+    elif format == "png" and not output_path.suffix == ".png":
+        output_path = output_path.with_suffix(".png")
 
     with tempfile.NamedTemporaryFile("w+", suffix=".d2", delete=False) as tmp:
         tmp.write(d2_diagram)
@@ -180,7 +425,7 @@ def _export_to_svg(d2_diagram: str, output_path: str | Path = "diagram.svg") ->
         tmp_path = Path(tmp.name)
 
     try:
-        cmd = [d2_binary, str(tmp_path), str(output_path)]
+        cmd = ["d2", str(tmp_path), str(output_path)]
         subprocess.run(cmd, check=True)
     finally:
         if tmp_path.exists():
@@ -189,78 +434,37 @@ def _export_to_svg(d2_diagram: str, output_path: str | Path = "diagram.svg") ->
     return output_path
 
 
-def generate_diagram(ip_network: str, hosts_list: str) -> str:
+def generate_diagram(ip_network: str, hosts_list: str, use_svg: bool = False) -> str:
     """
-    Generate a network diagram in SVG format for Gradio UI.
+    Generate a network diagram in PNG or SVG format.
 
     Args:
         ip_network (str): Network IP with mask in CIDR format (e.g., "192.168.1.0/24")
         hosts_list (str): Comma-separated list of host counts per subnet (e.g., "50,20,10,5")
+        use_svg (bool): Whether to generate SVG format (default: PNG)
 
     Returns:
-        str: Path to the generated SVG file for Gradio Image component
+        str: Result information in JSON format including image path or error
     """
     try:
         # Parse hosts list
         hosts_per_subnet = [int(h.strip()) for h in hosts_list.split(",") if h.strip()]
         
         if not hosts_per_subnet:
-            raise ValueError("No valid host counts provided")
-
-        # Generate basic diagram
-        base_diagram = _generate_basic_d2_diagram(ip_network.strip(), hosts_per_subnet)
-        
-        # Add styling
-        styled_diagram = _style_d2_diagram(base_diagram)
-        
-        # Export to SVG
-        output_path = _export_to_svg(styled_diagram, "network_diagram.svg")
-        
-        return str(output_path)
-        
-    except Exception as e:
-        # Create a simple error diagram for Gradio to display
-        error_diagram = f"""
-Internet: Cloud {{
-  style: {{
-    font-color: red
-  }}
-}}
-Error: Error occurred: {str(e)[:50]}... {{
-  style: {{
-    font-color: red
-    font-size: 20
-  }}
-}}
-Internet -> Error
-"""
-        try:
-            error_path = _export_to_svg(error_diagram, "error_diagram.svg")
-            return str(error_path)
-        except Exception as svg_error:
-            # Log both errors for debugging
-            print(f"Original error: {e}", file=__import__('sys').stderr)
-            print(f"SVG error: {svg_error}", file=__import__('sys').stderr)
-            return None
-
-
-def generate_diagram_mcp(ip_network: str, hosts_list: str) -> str:
-    """
-    Generate a network diagram in SVG format for MCP API.
-
-    Args:
-        ip_network (str): Network IP with mask in CIDR format (e.g., "192.168.1.0/24")
-        hosts_list (str): Comma-separated list of host counts per subnet (e.g., "50,20,10,5")
+            return json.dumps({"error": "No valid host counts provided"}, indent=2)
 
-    Returns:
-        str: Result information in JSON format including SVG path or error
-    """
-    try:
-        # Parse hosts list
-        hosts_per_subnet = [int(h.strip()) for h in hosts_list.split(",") if h.strip()]
+        # Validate that the subnet distribution is possible using VLSM
+        validation_result = calculate_subnets(
+            ip_network.strip(), 
+            len(hosts_per_subnet), 
+            "vlsm", 
+            hosts_list.strip()
+        )
         
-        if not hosts_per_subnet:
-            return json.dumps({"error": "No valid host counts provided"}, indent=2)
+        if "error" in validation_result:
+            return json.dumps({
+                "error": f"Invalid subnet distribution: {validation_result['error']}"
+            }, indent=2)
 
         # Generate basic diagram
         base_diagram = _generate_basic_d2_diagram(ip_network.strip(), hosts_per_subnet)
@@ -268,12 +472,15 @@ def generate_diagram_mcp(ip_network: str, hosts_list: str) -> str:
         # Add styling
         styled_diagram = _style_d2_diagram(base_diagram)
         
-        # Export to SVG
-        output_path = _export_to_svg(styled_diagram, "network_diagram.svg")
+        # Export to image
+        format = "svg" if use_svg else "png"
+        filename = f"network_diagram.{format}"
+        output_path = _export_to_image(styled_diagram, filename, format)
         
         result = {
             "success": True,
-            "svg_path": str(output_path),
+            "image_path": str(output_path),
+            "format": format,
             "network": ip_network.strip(),
             "hosts_per_subnet": hosts_per_subnet,
             "d2_diagram": styled_diagram

ipmentor/ui.py

@@ -4,75 +4,36 @@ Gradio UI for IPMentor.
 
 import gradio as gr
 import json
-from .core import analyze_ip, calculate_subnets
-from .tools import generate_diagram, generate_diagram_mcp
-from .config import APP_NAME
+from .tools import generate_diagram as generate_diagram_core, ip_info, subnet_calculator
 
-
-
-
-# MCP Tool Functions
-def ip_info(ip: str, subnet_mask: str) -> str:
+def generate_diagram(ip_network: str, hosts_list: str, use_svg: bool = False):
     """
-    Analyze a complete IPv4 address with its subnet mask.
+    Generate a network diagram in PNG or SVG format.
     
     Args:
-        ip (str): IP address in decimal (192.168.1.10) or binary format
-        subnet_mask (str): Subnet mask in decimal (255.255.255.0), CIDR (/24), or number (24) format
+        ip_network (str): Network IP with mask in CIDR format (e.g., "192.168.1.0/24")
+        hosts_list (str): Comma-separated list of host counts per subnet (e.g., "50,20,10,5")
+        use_svg (bool): Whether to generate SVG format (default: PNG)
     
     Returns:
-        str: Complete IP and network information in JSON format
+        tuple: (image_path, status_message) for Gradio outputs
     """
-    try:
-        result = analyze_ip(ip.strip(), subnet_mask.strip())
-        return json.dumps(result, indent=2)
-    except Exception as e:
-        return json.dumps({"error": str(e)}, indent=2)
-
-
-def subnet_calculator(network: str, number: str, division_type: str, hosts_per_subnet: str = "") -> str:
-    """
-    Calculate subnets using different division methods.
-    
-    Args:
-        network (str): Main network in CIDR format (e.g., "192.168.1.0/24")
-        number (str): Number for division calculation
-        division_type (str): Division method - "max_subnets", "max_hosts_per_subnet", or "vlsm"
-        hosts_per_subnet (str): Comma-separated host counts per subnet (VLSM only)
-    
-    Returns:
-        str: Calculated subnet information in JSON format
-    """
-    try:
-        number_int = int(number.strip())
-        result = calculate_subnets(
-            network.strip(), 
-            number_int, 
-            division_type.strip(), 
-            hosts_per_subnet.strip()
-        )
-        return json.dumps(result, indent=2)
-    except Exception as e:
-        return json.dumps({"error": str(e)}, indent=2)
-
-
-# MCP API function for diagram generation  
-def generate_diagram_api(ip_network: str, hosts_list: str) -> str:
-    """MCP API wrapper for diagram generation that returns JSON."""
-    return generate_diagram_mcp(ip_network, hosts_list)
-
-
-# Gradio UI wrapper for diagram generation
-def generate_diagram_ui(ip_network: str, hosts_list: str):
-    """Gradio UI wrapper for diagram generation that handles errors gracefully."""
     try:
         if not ip_network.strip() or not hosts_list.strip():
-            return None
-        result = generate_diagram(ip_network, hosts_list)
-        return result
+            return None, "❌ Error: Please provide both network and hosts list"
+        
+        result_json = generate_diagram_core(ip_network.strip(), hosts_list.strip(), use_svg)
+        result = json.loads(result_json)
+        
+        if "error" in result:
+            return None, f"❌ Error: {result['error']}"
+        
+        format_type = "SVG" if use_svg else "PNG"
+        hosts_count = len(result.get("hosts_per_subnet", []))
+        return result.get("image_path"), f"✅ Success: {format_type} diagram generated for {hosts_count} subnets"
+        
     except Exception as e:
-        # Return None if there's any error - Gradio will handle this gracefully
-        return None
+        return None, f"❌ Error: {str(e)}"
 
 
 def create_interface():
@@ -98,7 +59,7 @@ def create_interface():
             gr.Textbox(label="Network", placeholder="192.168.1.0/24"),
             gr.Textbox(label="Number", placeholder="4"),
             gr.Dropdown(label="Division Type", choices=["max_subnets","max_hosts_per_subnet","vlsm"], value="max_subnets"),
-            gr.Textbox(label="Hosts per Subnet", placeholder="100,50,25,10")
+            gr.Textbox(label="Hosts per Subnet", placeholder="100,50,25,10", value="")
         ],
         outputs=gr.Textbox(label="Calculation Result"),
         title="Subnet Calculator",
@@ -106,22 +67,48 @@ def create_interface():
     )
     
     diagram_interface = gr.Interface(
-        fn=generate_diagram_ui,
+        fn=generate_diagram,
         api_name="generate_diagram",
         inputs=[
             gr.Textbox(label="Network", placeholder="192.168.1.0/24"),
-            gr.Textbox(label="Hosts per Subnet", placeholder="50,20,10,5")
+            gr.Textbox(label="Hosts per Subnet", placeholder="50,20,10,5"),
+            gr.Checkbox(label="Generate as SVG", value=False)
+        ],
+        outputs=[
+            gr.Image(label="Network Diagram", type="filepath"),
+            gr.Textbox(label="Status", lines=2, interactive=False)
         ],
-        outputs=gr.Image(label="Network Diagram", type="filepath"),
         title="Network Diagram Generator",
-        description="Generate network diagrams with SVG output"
+        description="Generate network diagrams (PNG by default, SVG optional)"
     )
     
-    # Combine all the MCP tool interfaces
-    combined_app = gr.TabbedInterface(
-        [ip_interface, subnet_interface, diagram_interface],
-        ["IP Info", "Subnet Calculator", "Network Diagram"],
-        title=f"{APP_NAME} - MCP Tools"
-    )
+    # Create main interface with custom header and description
+    with gr.Blocks() as combined_app:
+        # Header with logo
+        gr.Image("assets/header.png", show_label=False, interactive=False, container=False, height=120)
+        
+        # Description
+        gr.Markdown("""
+        **IPMentor** is a comprehensive IPv4 networking toolkit that provides three powerful tools:
+        
+        - **IP Info**: Analyze IPv4 addresses with subnet masks, supporting decimal, binary, and CIDR formats
+        - **Subnet Calculator**: Calculate subnets using different methods (max subnets, max hosts per subnet, and VLSM)
+        - **Network Diagram**: Generate visual network diagrams with automatic subnet validation
+                    
+        <a href="https://github.com/DavidLMS/ipmentor" target="_blank">Project on GitHub</a>
+                    
+        <a href="https://huggingface.co/spaces/DavidLMS/ipmentor-demo" target="_blank">Client demo</a>
+        
+        Choose a tab below to get started with your networking calculations and visualizations.
+        """)
+        
+        # Tabbed interface
+        with gr.Tabs():
+            with gr.Tab("IP Info"):
+                ip_interface.render()
+            with gr.Tab("Subnet Calculator"):
+                subnet_interface.render()
+            with gr.Tab("Network Diagram"):
+                diagram_interface.render()
     
     return combined_app

requirements.txt

@@ -1,3 +1,2 @@
 gradio[mcp]>=4.44.0
-python-dotenv>=1.0.0
-playwright
+python-dotenv>=1.0.0