import xml.etree.ElementTree as ET
import re
import base64
from ciscoconfparse import CiscoConfParse
from Decipher.cisco_pwd import decrypt_cisco_type7

def sanitize_xml_content(xml_bytes):
    """
    Sanitizes XML content by properly handling problematic binary data in attributes
    instead of just deleting them.
    """
    # Robust decoding
    xml_content = xml_bytes.decode('utf-8', errors='ignore')
    
    def encode_binary_attr(match):
        attr_name = match.group(1)
        attr_value = match.group(2)
        # Check if value contains non-XML friendly characters
        if any(ord(c) < 32 and c not in '\n\r\t' for c in attr_value):
            encoded = base64.b64encode(attr_value.encode('utf-8', errors='ignore')).decode('utf-8')
            return f'{attr_name}="base64:{encoded}"'
        return match.group(0)

    # Find attributes that might contain binary data (like VALUE="...")
    xml_content = re.sub(r'(\b\w+)="([^"]*)"', encode_binary_attr, xml_content)
    
    # Ensure root tag is closed if it's a snippet
    if "</PACKETTRACER5>" not in xml_content and "<PACKETTRACER5>" in xml_content:
        xml_content += "</PACKETTRACER5>"
    elif "</PACKETTRACER5_ACTIVITY>" not in xml_content and "<PACKETTRACER5_ACTIVITY>" in xml_content:
        xml_content += "</PACKETTRACER5_ACTIVITY>"
        
    return xml_content

def extract_structured_data(xml_path):
    try:
        with open(xml_path, 'rb') as f:
            xml_bytes = f.read()
        
        xml_content = sanitize_xml_content(xml_bytes)
        root = ET.fromstring(xml_content)
    except Exception as e:
        print(f"[-] Error parsing {xml_path}: {e}")
        return None

    data = {
        "version": root.findtext(".//VERSION"),
        "devices": [],
        "physical_links": [],
        "wireless_connections": [],
        "vlans": {},
        "instructions": ""
    }

    # 1. Extract Devices and their properties
    devices_by_id = {}
    ssids_servers = {} # ssid -> device_name

    for device in root.findall(".//DEVICE"):
        engine = device.find("ENGINE")
        if engine is None:
            continue
            
        name = engine.findtext("NAME")
        dev_type = engine.findtext("TYPE")
        model = engine.find("TYPE").get("model") if engine.find("TYPE") is not None else ""
        ref_id = engine.findtext("SAVE_REF_ID")
        
        dev_info = {
            "name": name,
            "type": dev_type,
            "model": model,
            "interfaces": [],
            "wireless_info": None,
            "vlans": []
        }
        
        if ref_id:
            devices_by_id[ref_id] = name

        # Extract Interfaces & IP info
        for port in engine.findall(".//PORT"):
            port_name = port.findtext("PORT_NAME") or port.findtext("NAME")
            if not port_name:
                # Try to find in parent SLOT/MODULE
                parent_module = port.find("..")
                if parent_module is not None:
                    port_type = port.findtext("TYPE")
                    if port_type:
                        port_name = port_type
            
            if port_name:
                ip = port.findtext("IP")
                subnet = port.findtext("SUBNET")
                ipv6 = port.findtext("IPV6_LINK_LOCAL")
                
                if ip or ipv6 or port.findtext("PORT_DHCP_ENABLE") == "true":
                    dev_info["interfaces"].append({
                        "name": port_name,
                        "ip": ip if ip else "DHCP" if port.findtext("PORT_DHCP_ENABLE") == "true" else None,
                        "subnet": subnet,
                        "ipv6": ipv6
                    })

        # Extract Wireless Server info (Access Points)
        wireless_server = engine.find("WIRELESS_SERVER")
        if wireless_server is not None:
            common = wireless_server.find("WIRELESS_COMMON")
            if common is not None:
                ssid = common.findtext("SSID")
                dev_info["wireless_info"] = {
                    "role": "server",
                    "ssid": ssid,
                    "encryption": common.findtext("ENCRYPT_TYPE"),
                    "key": common.find(".//KEY").text if common.find(".//KEY") is not None else None
                }
                if ssid:
                    ssids_servers[ssid] = name

        # Extract Wireless Client info (PC/Laptop/Phone)
        wireless_client = engine.find("WIRELESS_CLIENT")
        if wireless_client is not None:
            common = wireless_client.find("WIRELESS_COMMON")
            if common is not None:
                ssid = common.findtext("SSID")
                dev_info["wireless_info"] = {
                    "role": "client",
                    "ssid": ssid,
                    "encryption": common.findtext("ENCRYPT_TYPE")
                }

        # Extract VLANs from device-specific storage (Switches)
        vlan_storage = engine.findall(".//VLANS/VLAN")
        for v in vlan_storage:
            v_num = v.get("number")
            v_name = v.get("name")
            if v_num and v_num not in data["vlans"]:
                data["vlans"][v_num] = v_name
            dev_info["vlans"].append({"id": v_num, "name": v_name})

        # Parse RUNNINGCONFIG using CiscoConfParse for deeper analysis
        config = engine.findtext("RUNNINGCONFIG")
        if config:
            parse = CiscoConfParse(config.splitlines(), factory=True)
            
            # 1. Improved VLAN assignments
            for intf_obj in parse.find_objects(r'^interface'):
                intf_name = intf_obj.text.replace('interface ', '').strip()
                vlan_id = None
                
                # Check for access vlan
                access_vlan = intf_obj.re_search_children(r'switchport access vlan (\d+)')
                if access_vlan:
                    vlan_id = access_vlan[0].re_match_typed(r'switchport access vlan (\d+)')
                
                if vlan_id:
                    # Update or add interface info
                    found = False
                    for i in dev_info["interfaces"]:
                        if i["name"] == intf_name:
                            i["vlan"] = vlan_id
                            found = True
                            break
                    if not found:
                        dev_info["interfaces"].append({"name": intf_name, "vlan": vlan_id})

            # 2. L3/L4 Support: Routing Protocols (Static, OSPF)
            dev_info["routing"] = {
                "static_routes": [],
                "ospf": [],
                "nat": [],
                "dhcp_pools": []
            }
            
            # Static Routes
            static_routes = parse.find_objects(r'^ip route')
            for route in static_routes:
                dev_info["routing"]["static_routes"].append(route.text)
                
            # OSPF
            ospf_processes = parse.find_objects(r'^router ospf')
            for ospf in ospf_processes:
                ospf_info = {"process_id": ospf.text.split()[-1], "networks": []}
                for network in ospf.re_search_children(r'network'):
                    ospf_info["networks"].append(network.text.strip())
                dev_info["routing"]["ospf"].append(ospf_info)

            # NAT
            nat_rules = parse.find_objects(r'^ip nat')
            for rule in nat_rules:
                dev_info["routing"]["nat"].append(rule.text)

            # DHCP Pools
            dhcp_pools = parse.find_objects(r'^ip dhcp pool')
            for pool in dhcp_pools:
                pool_info = {
                    "name": pool.text.replace('ip dhcp pool ', '').strip(),
                    "details": [child.text.strip() for child in pool.children]
                }
                dev_info["routing"]["dhcp_pools"].append(pool_info)

            # 3. Security Audit & ACLs
            dev_info["security"] = {
                "acls": [],
                "vulnerabilities": [],
                "decrypted_passwords": []
            }

            # Decrypt Type 7 passwords
            type7_pwds = re.findall(r'password 7 ([0-9A-Fa-f]+)|enable password 7 ([0-9A-Fa-f]+)', config)
            for pwd_match in type7_pwds:
                encrypted = pwd_match[0] or pwd_match[1]
                decrypted = decrypt_cisco_type7(encrypted)
                if decrypted:
                    dev_info["security"]["decrypted_passwords"].append({
                        "encrypted": encrypted,
                        "decrypted": decrypted,
                        "type": "Cisco Type 7"
                    })
                    dev_info["security"]["vulnerabilities"].append(f"Weak encryption found: Decrypted password '{decrypted}'")

            # Standard and Extended ACLs
            acls = parse.find_objects(r'^access-list|ip access-list')
            for acl in acls:
                acl_info = {"name": acl.text.strip(), "rules": []}
                if acl.children:
                    acl_info["rules"] = [child.text.strip() for child in acl.children]
                dev_info["security"]["acls"].append(acl_info)

            # Security Audit: Check for plaintext passwords or weak services
            if parse.find_objects(r'^no service password-encryption'):
                dev_info["security"]["vulnerabilities"].append("Plaintext password storage enabled (no service password-encryption)")
            
            if parse.find_objects(r'^line vty'):
                vty_lines = parse.find_objects(r'^line vty')
                for vty in vty_lines:
                    if not vty.re_search_children(r'transport input ssh'):
                        dev_info["security"]["vulnerabilities"].append(f"Insecure remote access (Telnet allowed on {vty.text.strip()})")
            
            # 4. Server Services (if applicable)
            dev_info["services"] = []
            services_node = engine.find("SERVICES")
            if services_node is not None:
                # HTTP Service
                http = services_node.find("HTTP")
                if http is not None and http.get("enabled") == "true":
                    dev_info["services"].append({"type": "HTTP", "port": 80})
                
                # DNS Service
                dns = services_node.find("DNS")
                if dns is not None and dns.get("enabled") == "true":
                    dns_info = {"type": "DNS", "records": []}
                    for record in dns.findall("RECORD"):
                        dns_info["records"].append({
                            "name": record.get("name"),
                            "type": record.get("type"),
                            "value": record.get("value")
                        })
                    dev_info["services"].append(dns_info)

        data["devices"].append(dev_info)

    # 2. Extract Physical Links
    for link in root.findall(".//LINK"):
        cable = link.find("CABLE")
        if cable is not None:
            ports = cable.findall("PORT")
            from_id = cable.findtext("FROM")
            to_id = cable.findtext("TO")
            
            link_info = {
                "type": link.findtext("TYPE"),
                "from_device": devices_by_id.get(from_id, from_id),
                "to_device": devices_by_id.get(to_id, to_id),
            }
            if len(ports) >= 2:
                link_info["from_port"] = ports[0].text
                link_info["to_port"] = ports[1].text
            
            data["physical_links"].append(link_info)

    # 3. Establish Wireless Connections
    for dev in data["devices"]:
        if dev["wireless_info"] and dev["wireless_info"]["role"] == "client":
            ssid = dev["wireless_info"]["ssid"]
            if ssid in ssids_servers:
                data["wireless_connections"].append({
                    "from_device": dev["name"],
                    "to_device": ssids_servers[ssid],
                    "ssid": ssid,
                    "type": "wireless"
                })

    # 4. Extract Instructions
    instr_page = root.find(".//INSTRUCTIONS/PAGE")
    if instr_page is not None:
        data["instructions"] = instr_page.text

    return data