#!/usr/bin/env python3
"""
Get all network devices from a MikroTik router (DHCP, ARP, interfaces)
to identify access points and other devices
"""
import ssl
import sys
import json
import argparse
from collections import defaultdict

try:
    from librouteros import connect
except ImportError:
    print("Error: librouteros module not found")
    print("Install with: pip install librouteros")
    sys.exit(1)

def connect_to_router(ip, username, password, use_ssl=True):
    """Connect to MikroTik router"""
    port = 8729 if use_ssl else 8728
    try:
        if use_ssl:
            ctx = ssl.create_default_context()
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE

            api = connect(
                username=username,
                password=password,
                host=ip,
                port=port,
                ssl_wrapper=ctx.wrap_socket
            )
        else:
            api = connect(
                username=username,
                password=password,
                host=ip,
                port=port
            )

        print(f"✓ Connected to {ip}")
        return api
    except Exception as e:
        print(f"✗ Connection failed: {e}")
        return None

def get_dhcp_leases(api):
    """Get DHCP leases"""
    try:
        return list(api.path('/ip/dhcp-server/lease'))
    except Exception as e:
        print(f"Error getting DHCP leases: {e}")
        return []

def get_arp_table(api):
    """Get ARP table"""
    try:
        return list(api.path('/ip/arp'))
    except Exception as e:
        print(f"Error getting ARP table: {e}")
        return []

def get_device_type(mac, hostname, comment):
    """Determine device type based on MAC prefix and other identifiers"""
    if not mac:
        return "Unknown"

    mac_upper = mac.upper()

    # Ubiquiti prefixes
    ubiquiti_prefixes = ['04:18:D6', 'F0:9F:C2', '24:A4:3C', '68:72:51', '80:2A:A8', 'FC:EC:DA']
    if any(mac_upper.startswith(prefix) for prefix in ubiquiti_prefixes):
        # Check if it's likely an access point vs customer device
        if hostname or comment:
            name = (hostname or comment).lower()
            if any(x in name for x in ['ap', 'access', 'bridge', 'station', 'loco', 'nano', 'powerbeam', 'rocket']):
                return "Ubiquiti AP"
        return "Ubiquiti Device"

    # MikroTik prefixes
    mikrotik_prefixes = ['00:0C:42', '4C:5E:0C', '6C:3B:6B', 'D4:CA:6D', 'E4:8D:8C', 'DC:2C:6E', 'B8:69:F4', '48:8F:5A']
    if any(mac_upper.startswith(prefix) for prefix in mikrotik_prefixes):
        if hostname or comment:
            name = (hostname or comment).lower()
            if any(x in name for x in ['ap', 'cap', 'hap', 'wap', 'sxt', 'lhg']):
                return "MikroTik AP"
        return "MikroTik Device"

    return "Other Device"

def main():
    parser = argparse.ArgumentParser(description='Get all network devices from MikroTik router')
    parser.add_argument('router_ip', help='Router IP address')
    parser.add_argument('-u', '--username', default='grahamro', help='Username (default: grahamro)')
    parser.add_argument('-p', '--password', default='cFKhz8q5gPLoucMbcT1Iy58r3IXgc3', help='Password')
    parser.add_argument('-o', '--output', help='Output file (JSON)')
    parser.add_argument('--no-ssl', action='store_true', help='Use plain API instead of API-SSL')
    parser.add_argument('--show-all', action='store_true', help='Show all devices, not just likely APs')

    args = parser.parse_args()

    api = connect_to_router(args.router_ip, args.username, args.password, not args.no_ssl)
    if not api:
        sys.exit(1)

    print("\n=== Retrieving Network Device Information ===\n")

    dhcp_leases = get_dhcp_leases(api)
    arp_table = get_arp_table(api)

    # Merge data by MAC address
    devices = defaultdict(lambda: {
        'ip': '',
        'mac': '',
        'hostname': '',
        'comment': '',
        'status': '',
        'type': 'Unknown',
        'server': '',
        'interface': ''
    })

    # Process DHCP leases
    for lease in dhcp_leases:
        mac = lease.get('mac-address', '')
        if mac:
            dev = devices[mac]
            dev['mac'] = mac
            dev['ip'] = lease.get('address', dev['ip'])
            dev['hostname'] = lease.get('host-name', dev['hostname'])
            dev['comment'] = lease.get('comment', dev['comment'])
            dev['status'] = lease.get('status', dev['status'])
            dev['server'] = lease.get('server', dev['server'])

    # Process ARP table
    for arp in arp_table:
        mac = arp.get('mac-address', '')
        if mac:
            dev = devices[mac]
            dev['mac'] = mac
            if not dev['ip']:
                dev['ip'] = arp.get('address', '')
            if not dev['interface']:
                dev['interface'] = arp.get('interface', '')

    # Determine device types
    for mac, dev in devices.items():
        dev['type'] = get_device_type(mac, dev['hostname'], dev['comment'])

    # Group by type
    by_type = defaultdict(list)
    for dev in devices.values():
        by_type[dev['type']].append(dev)

    # Display results
    print("=== Network Devices by Type ===\n")

    for device_type in ['Ubiquiti AP', 'MikroTik AP', 'Ubiquiti Device', 'MikroTik Device', 'Other Device']:
        if device_type in by_type:
            devices_of_type = sorted(by_type[device_type], key=lambda x: x['ip'])

            if not args.show_all and device_type == 'Other Device':
                print(f"\n{device_type}s: {len(devices_of_type)} (use --show-all to display)")
                continue

            print(f"\n{device_type}s: {len(devices_of_type)}")
            print("-" * 100)
            for dev in devices_of_type:
                name = dev['hostname'] or dev['comment'] or dev['mac']
                print(f"  {name:<35} {dev['ip']:<15} {dev['mac']:<17} {dev['interface']:<20}")
                if dev['comment'] and dev['comment'] != name:
                    print(f"    Comment: {dev['comment']}")

    # Summary
    print(f"\n=== Summary ===")
    print(f"Total devices found: {len(devices)}")
    for device_type, devs in sorted(by_type.items()):
        print(f"  {device_type}: {len(devs)}")

    # Save to file if requested
    if args.output:
        output_data = {
            'router': args.router_ip,
            'devices': [dev for dev in devices.values()],
            'by_type': {k: v for k, v in by_type.items()}
        }
        with open(args.output, 'w') as f:
            json.dump(output_data, f, indent=2)
        print(f"\nData saved to: {args.output}")

    api.close()

if __name__ == '__main__':
    main()
