from prometheus_client import start_http_server, Gauge
import time
import subprocess
import re

pcs = [
    {"ip": "10.10.20.28", "pc":13, "location": "HUB1"},
    {"ip": "10.10.20.82", "pc":12, "location": "HUB1"},
    {"ip": "10.10.20.76", "pc":11, "location": "HUB1"},
    {"ip": "10.10.20.31", "pc":16, "location": "HUB1"},
    {"ip": "10.10.20.27", "pc":15, "location": "HUB1"},
    {"ip": "10.10.20.26", "pc":14, "location": "HUB1"},

    {"ip": "10.10.20.40", "pc":23, "location": "HUB2"},
    {"ip": "10.10.20.30", "pc":22, "location": "HUB2"},
    {"ip": "10.10.20.43", "pc":21, "location": "HUB2"},
    {"ip": "10.10.20.66", "pc":26, "location": "HUB2"},
    {"ip": "10.10.20.52", "pc":25, "location": "HUB2"},
    {"ip": "10.10.20.34", "pc":24, "location": "HUB2"},

    {"ip": "10.10.20.104", "pc":33, "location": "HUB3"},
    {"ip": "10.10.20.60", "pc":32, "location": "HUB3"},
    {"ip": "10.10.20.4", "pc":31, "location": "HUB3"},
    {"ip": "10.10.20.103", "pc":36, "location": "HUB3"},
    {"ip": "10.10.20.65", "pc":35, "location": "HUB3"},
    {"ip": "10.10.20.5", "pc":34, "location": "HUB3"},

    {"ip": "10.10.20.33", "pc":43, "location": "HUB4"},
    {"ip": "10.10.20.16", "pc":42, "location": "HUB4"},
    {"ip": "10.10.20.22", "pc":41, "location": "HUB4"},
    {"ip": "10.10.20.8", "pc":46, "location": "HUB4"},
    {"ip": "10.10.20.70", "pc":45, "location": "HUB4"},
    {"ip": "10.10.20.29", "pc":44, "location": "HUB4"},

    {"ip": "10.10.20.18", "pc":53, "location": "HUB5"},
    {"ip": "10.10.20.9", "pc":52, "location": "HUB5"},
    {"ip": "10.10.20.46", "pc":51, "location": "HUB5"},
    {"ip": "10.10.20.10", "pc":56, "location": "HUB5"},
    {"ip": "10.10.20.86", "pc":55, "location": "HUB5"},
    {"ip": "10.10.20.127", "pc":54, "location": "HUB5"},

    {"ip": "10.10.20.98", "pc":63, "location": "HUB6"},
    {"ip": "10.10.20.137", "pc":62, "location": "HUB6"},
    {"ip": "10.10.20.15", "pc":61, "location": "HUB6"},
    {"ip": "10.10.20.72", "pc":66, "location": "HUB6"},
    {"ip": "10.10.20.21", "pc":65, "location": "HUB6"},
    {"ip": "10.10.20.90", "pc":64, "location": "HUB6"},

    {"ip": "10.10.20.88", "pc":73, "location": "HUB7"},
    {"ip": "10.10.20.58", "pc":72, "location": "HUB7"},
    {"ip": "10.10.20.132", "pc":71, "location": "HUB7"},
    {"ip": "10.10.20.25", "pc":76, "location": "HUB7"},
    {"ip": "10.10.20.74", "pc":75, "location": "HUB7"},
    {"ip": "10.10.20.87", "pc":74, "location": "HUB7"},
]

pc_status = Gauge('pc_status', 'Status of PCs (1 = online, 0 = offline)', ['ip', 'location', 'pc'])

def run_arp_scan(network):
    """
    Run arp-scan command on the network.
    Returns a set of IPs that responded.
    """
    try:
        # Check if arp-scan is installed
        try:
            subprocess.run(['which', 'arp-scan'], check=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        except:
            print("arp-scan not found. Install with: sudo apt-get install arp-scan")
            return set()
        
        # Run the scan on the specified network
        print(f"Running arp-scan on {network}...")
        result = subprocess.run(
            ['sudo', 'arp-scan', '--interface=wlan0', network],
            check=False,
            stderr=subprocess.PIPE,
            stdout=subprocess.PIPE,
            text=True,
            timeout=10
        )
        
        # Parse the output to extract IPs
        ip_pattern = r'(\d+\.\d+\.\d+\.\d+)'
        ips = re.findall(ip_pattern, result.stdout)
        
        print(f"Found {len(ips)} devices with arp-scan")
        return set(ips)
    except Exception as e:
        print(f"Error running arp-scan: {e}")
        return set()

def update_metrics():
    """
    Update PC status metrics using arp-scan.
    """
    start_time = time.time()
    
    # Run arp-scan on the network
    online_ips = run_arp_scan("10.10.20.0/24")
    
    # Update metrics for each PC
    online_count = 0
    for pc in pcs:
        ip = pc["ip"]
        location = pc["location"]
        pc_num = str(pc["pc"])
        
        # Check if the IP is in the set of online IPs
        status = ip in online_ips
        
        # Update Prometheus metric
        pc_status.labels(ip=ip, location=location, pc=pc_num).set(1 if status else 0)
        
        # Print status
        if status:
            print(f"PC {pc_num} at {ip} is ONLINE")
            online_count += 1
        else:
            print(f"PC {pc_num} at {ip} is offline")
    
    scan_time = time.time() - start_time
    print(f"Scan completed in {scan_time:.2f} seconds, found {online_count} online PCs out of {len(pcs)}")

def check_network_interfaces():
    """
    Print available network interfaces
    """
    try:
        result = subprocess.run(['ip', 'addr'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        print("Available network interfaces:")
        print(result.stdout)
    except Exception as e:
        print(f"Error checking network interfaces: {e}")

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='PC Status Monitor for Prometheus')
    parser.add_argument('--test', help='Test scan without starting Prometheus server', action='store_true')
    parser.add_argument('--check-interfaces', help='Check available network interfaces', action='store_true')
    parser.add_argument('--port', help='Port for Prometheus metrics', type=int, default=8003)
    args = parser.parse_args()
    
    # Check network interfaces if requested
    if args.check_interfaces:
        check_network_interfaces()
        exit(0)
    
    # Test scan if requested
    if args.test:
        update_metrics()
        exit(0)
    
    # Start the Prometheus HTTP server
    port = args.port
    start_http_server(port)
    print(f"Prometheus metrics server started on port {port}")
    
    # Continuously update metrics
    try:
        while True:
            print("\n--- Updating PC status ---")
            update_metrics()
            print(f"Waiting 15 seconds before next check...")
            time.sleep(15)  # Update every 15 seconds
    except KeyboardInterrupt:
        print("\nMonitoring stopped by user")