wif2ansible/wif2ansible/inventory.py

from typing import List, Dict, Any
from .models import Server, Flow
from .models import Server, Flow
from .network import to_mgt_ip, is_valid_hostname, get_hostname
from .parsers import clean_reference

def generate_inventory(servers: Dict[str, Server], flows: List[Flow]) -> Dict[str, Any]:
    """
    Generates the Ansible inventory dictionary.
    servers: Dict[Reference, Server]
    flows: List[Flow]
    """
    
    # Build Lookup Map: IP -> Server
    # Note: A server might have multiple IPs (e.g. Mgt, Public, Private).
    # The 'Server' object mainly captures the Management IP or the one listed in the "IP Address" column.
    # If the WIF has "Source Public IP" and that differs from "IP Address" in Servers tab, 
    # we might miss it if we only index the primary IP.
    # However, strict filtering means we trust the 'Servers' tab.
    
    ip_to_server = {}
    for s in servers.values():
        # Index all Management IPs
        for ip in s.ip_addresses:
            ip_to_server[ip] = s
        
        # Index all Production IPs
        for ip in s.production_ips:
            ip_to_server[ip] = s
            
        # Also index by reference/hostname for DNS matches
        if s.reference:
             ip_to_server[s.reference.lower()] = s
        if s.hostname:
             ip_to_server[s.hostname.lower()] = s
            
    inventory_hosts = {}
    
    # Process flows
    match_count = 0
    drop_count = 0
    total_flows = len(flows)
    
    print(f"Starting inventory generation for {total_flows} flows...")
    
    for idx, flow in enumerate(flows, 1):
        if idx % 10 == 0:
             print(f"Processing flow {idx}/{total_flows}...")
             
        # Find source server
        server = ip_to_server.get(flow.source_ip)
        
        if not server:
            # Try DNS resolution (Public IP -> Management FQDN)
            print(f"Flow {idx}: Source {flow.source_ip} not found in map. Attempting DNS resolution...")
            mgt_dns = to_mgt_ip(flow.source_ip)
            if mgt_dns:
                # mgt_dns might be "server.ds.gc.ca".
                # Our keys might be "server" or "server.ds.gc.ca" or IPs
                # Try exact match
                server = ip_to_server.get(mgt_dns.lower())
                
                # If not found, try shortname?
                if not server:
                    short = mgt_dns.split('.')[0]
                    server = ip_to_server.get(short.lower())
                    
            if not server:
                drop_count += 1
                if drop_count <= 10: # Increased debug spam limit
                     print(f"Dropping flow {flow.flow_id} ({idx}/{total_flows}): Source {flow.source_ip} (Mgt: {mgt_dns}) resolved but not found in Servers tab.")
                continue
            else:
                 print(f"Flow {idx}: Resolved {flow.source_ip} -> {server.hostname or server.reference}")
            
        match_count += 1
        
        # Prepare host entry if new
        
        # Candidate Resolution Logic
        # User Requirement: "gather all potential names ... check to see what actually resolves"
        candidates = []
        
        # 1. Server Name Column (Highest priority from Excel)
        if server.hostname:
            candidates.append(server.hostname)
            
        # 2. Cleaned Reference (Fallback from Excel)
        if server.reference:
            candidates.append(clean_reference(server.reference))
            
        # 3. Reverse DNS of Primary IP?
        # If the Excel names are garbage, maybe the IP resolves to the "Real" DNS name.
        if server.primary_ip:
             # Try simple reverse lookup
             rev_name = get_hostname(server.primary_ip)
             if rev_name:
                 candidates.append(rev_name)

        # Select the first candidate that resolves
        
        final_host_key = None
        for cand in candidates:
            if not cand: continue
            if is_valid_hostname(cand):
                final_host_key = cand
                break
                
        # Fallback: strict fallback to IP if nothing resolves? 
        # Or best effort (first candidate)?
        # User said: "You are getting it incorrect every time" -> likely implying the garbage name was used.
        # But if *nothing* resolves, we must output something. The IP is safe connectivity-wise, but user wants Names.
        # Let's fallback to the IP if NO name works, to ensure ansible works.
        if not final_host_key:
            if candidates:
                 # Warn?
                 print(f"Warning: No resolvable name found for {server.primary_ip} (Candidates: {candidates}). Using IP.")
            final_host_key = server.primary_ip
            
        host_key = final_host_key
        
        if host_key not in inventory_hosts:
            host_vars = server.get_ansible_vars()
            host_vars['flows'] = []
            inventory_hosts[host_key] = host_vars
            
        # Add flow
        flow_entry = {
            'flow_id': flow.flow_id,
            'dest': flow.destination_ip,
            'ports': flow.ports,
            'protocol': flow.protocol
        }
        
        # Dedup check? 
        # Ideally we shouldn't have exact duplicates, but appending is safe.
        inventory_hosts[host_key]['flows'].append(flow_entry)

    print(f"Inventory Generation Report: Matches={match_count}, Dropped={drop_count}")

    return {
        'all': {
            'hosts': inventory_hosts
        }
    }