import nmap
import psutil
import socket
import ipaddress
import pandas as pd
from datetime import datetime
import os

# Set up data directory
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.abspath(os.path.join(BASE_DIR, "..", "data"))
os.makedirs(DATA_DIR, exist_ok=True)
excluded_prefixes = ("lo", "docker", "br-", "veth", "vmnet", "virbr")

def get_ipv4_interfaces():
    ipv4_interfaces = []
    for name, addrs in psutil.net_if_addrs().items():
        for addr in addrs:
            if addr.family == socket.AF_INET:
                # Skip loopback or docker interfaces
                if name.startswith(excluded_prefixes):
                    continue
                ipv4_interfaces.append((name, addr.address, addr.netmask))
    return ipv4_interfaces

def network_discovery():
    ipv4_interfaces = get_ipv4_interfaces()
    if not ipv4_interfaces:
        print("No active IPv4 network interfaces found.")
        return

    nm = nmap.PortScanner()
    active_hosts = []

    for name, ip, netmask in ipv4_interfaces:
        network = ipaddress.IPv4Network((ip, netmask), strict=False)
        print(f"Discovering hosts in IP range {network} using interface {name}...")

        try:
            nm.scan(hosts=str(network), arguments='-sn -T4')
        except Exception as e:
            print(f"Error scanning {network}: {e}")
            continue

        discovered = nm.all_hosts()
        print(f"Found {len(discovered)} hosts on {network}")

        for host in discovered:
            active_hosts.append([host, datetime.now().strftime("%Y-%m-%d %H:%M:%S")])

    if active_hosts:
        columns = ['IP', 'Discovery Timestamp']
        df = pd.DataFrame(active_hosts, columns=columns)
        filename = "active_hosts.csv"
        output_file = os.path.join(DATA_DIR, filename)
        df.to_csv(output_file, index=False)
        print(f"Discovery complete. Results saved to {filename}")
        return df
    else:
        print("No active hosts found.")
        return None

if __name__ == "__main__":
    network_discovery()