Update modules/visual_route.py

modules/visual_route.py

@@ -19,6 +19,36 @@ AS_NAME_WRAP_WIDTH = 25
 ROUTER_NAME = os.getenv('BGP_VIS_ROUTER_NAME', 'My Router')
 MAX_IP_CIDR_LENGTH = 45
 
+def _is_valid_ip_or_prefix(target: str) -> bool:
+    try:
+        if '/' in target:
+            ipaddress.ip_network(target, strict=False)
+        else:
+            ipaddress.ip_address(target)
+        return True
+    except ValueError:
+        return False
+
+def _validate_ip_prefix_input(ip_address_str: str):
+    if not ip_address_str:
+        return None, "An IP address or prefix is required."
+    if len(ip_address_str) > MAX_IP_CIDR_LENGTH:
+        return None, f"Input exceeds maximum length of {MAX_IP_CIDR_LENGTH} characters."
+    if not _is_valid_ip_or_prefix(ip_address_str):
+        return None, f"Invalid input '{ip_address_str}'. Please provide a valid IPv4/IPv6 address or prefix."
+    return ip_address_str, None
+
+def _get_lookup_params(ip_address_str: str):
+    try:
+        if '/' in ip_address_str:
+            net_obj = ipaddress.ip_network(ip_address_str, strict=False)
+            return f"ipv{net_obj.version}", net_obj.with_prefixlen
+        else:
+            ip_obj = ipaddress.ip_address(ip_address_str)
+            return f"ipv{ip_obj.version}", str(ip_obj)
+    except ValueError:
+        return None, None
+
 def _bulk_get_as_names(asn_numbers: list[str]):
     lookup_list = [asn for asn in asn_numbers if asn not in AS_NAME_CACHE]
     if not lookup_list: return
@@ -171,24 +201,31 @@ def _parse_bgp_paths_to_graph(bgp_data: str) -> dict:
     return {"nodes": nodes, "edges": edges, "path_count": len(ordered_paths)}
 
 
-def generate_visual_route_graph(ip_address_str: str) -> dict:
-    if not ip_address_str:
-        return {"error": "An IP address or prefix is required."}
+def get_raw_bgp_route(ip_address_str: str) -> tuple[str | None, str | None]:
+    ip_address_str, error = _validate_ip_prefix_input(ip_address_str)
+    if error:
+        return None, error
     
-    if len(ip_address_str) > MAX_IP_CIDR_LENGTH:
-        return {"error": f"Input exceeds maximum length of {MAX_IP_CIDR_LENGTH} characters."}
+    ip_version, address_to_lookup = _get_lookup_params(ip_address_str)
+    if not address_to_lookup:
+        return None, f"Invalid input '{ip_address_str}'."
 
-    address_to_lookup, ip_version = "", ""
-    try:
-        if '/' in ip_address_str:
-            net_obj = ipaddress.ip_network(ip_address_str, strict=False)
-            ip_version, address_to_lookup = f"ipv{net_obj.version}", net_obj.with_prefixlen
-        else:
-            ip_obj = ipaddress.ip_address(ip_address_str)
-            ip_version, address_to_lookup = f"ipv{ip_obj.version}", str(ip_obj)
-    except ValueError:
-        return {"error": f"Invalid input '{ip_address_str}'. Please provide a valid IPv4/IPv6 address or prefix."}
+    bgp_data = _get_bgp_data(address_to_lookup, ip_version)
+    if not bgp_data:
+        return None, f"Route information not found for: {address_to_lookup}"
+        
+    return bgp_data, None
+
+
+def generate_visual_route_graph(ip_address_str: str) -> dict:
+    ip_address_str, error = _validate_ip_prefix_input(ip_address_str)
+    if error:
+        return {"error": error}
     
+    ip_version, address_to_lookup = _get_lookup_params(ip_address_str)
+    if not address_to_lookup:
+        return {"error": f"Invalid input '{ip_address_str}'."}
+
     bgp_data = _get_bgp_data(address_to_lookup, ip_version)
     if not bgp_data:
         return {"not_found": True, "target": address_to_lookup}