#!/usr/bin/env python3
"""
Get access points from a MikroTik router by checking DHCP leases and ARP table
"""
import ssl
import sys
import json
import argparse

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:
            # SSL context for secure connection
            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 from the router"""
    try:
        leases = list(api.path('/ip/dhcp-server/lease'))
        return leases
    except Exception as e:
        print(f"Error getting DHCP leases: {e}")
        return []

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

def get_capsman_registrations(api):
    """Get Capsman registration table if available"""
    try:
        registrations = list(api.path('/caps-man/registration-table'))
        return registrations
    except Exception as e:
        # Capsman might not be configured
        return []

def is_likely_ap(hostname, mac, comment):
    """Determine if a device is likely an access point based on hostname/MAC/comment"""
    if not hostname:
        hostname = ""
    if not comment:
        comment = ""

    hostname_lower = hostname.lower()
    comment_lower = comment.lower()

    # Common AP identifiers
    ap_indicators = ['ap', 'access', 'ubiquiti', 'unifi', 'mikrotik', 'routerboard',
                     'airmax', 'litebeam', 'nanostation', 'powerbeam', 'rocket',
                     'hap', 'cap', 'sxt', 'lhg', 'wap']

    for indicator in ap_indicators:
        if indicator in hostname_lower or indicator in comment_lower:
            return True

    # Check for Ubiquiti MAC prefix (common for APs)
    if mac and mac.upper().startswith(('04:18:D6', 'F0:9F:C2', '24:A4:3C', '68:72:51', '80:2A:A8', 'FC:EC:DA')):
        return True

    # Check for MikroTik MAC prefix
    if mac and mac.upper().startswith(('00:0C:42', '4C:5E:0C', '6C:3B:6B', 'D4:CA:6D', 'E4:8D:8C', 'DC:2C:6E', 'B8:69:F4', '48:8F:5A')):
        return True

    return False

def main():
    parser = argparse.ArgumentParser(description='Get access points 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')

    args = parser.parse_args()

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

    print("\n=== Retrieving Access Point Information ===\n")

    # Get data from router
    dhcp_leases = get_dhcp_leases(api)
    arp_table = get_arp_table(api)
    capsman_regs = get_capsman_registrations(api)

    access_points = []

    # Check Capsman registrations first (most reliable for APs)
    if capsman_regs:
        print("=== Capsman Registered Access Points ===")
        for reg in capsman_regs:
            ap_info = {
                'name': reg.get('interface', 'Unknown'),
                'mac': reg.get('mac-address', ''),
                'ip': reg.get('address', ''),
                'source': 'capsman',
                'interface': reg.get('interface', ''),
                'rx_signal': reg.get('rx-signal', '')
            }
            access_points.append(ap_info)
            print(f"  {ap_info['name']}: {ap_info['ip']} ({ap_info['mac']})")

    # Process DHCP leases
    print("\n=== DHCP Leases (Potential Access Points) ===")
    for lease in dhcp_leases:
        hostname = lease.get('host-name', '')
        mac = lease.get('mac-address', '')
        ip = lease.get('address', '')
        comment = lease.get('comment', '')

        if is_likely_ap(hostname, mac, comment):
            ap_info = {
                'name': hostname or comment or mac,
                'ip': ip,
                'mac': mac,
                'source': 'dhcp',
                'comment': comment
            }

            # Check if already in list (from capsman)
            if not any(ap['mac'] == mac for ap in access_points):
                access_points.append(ap_info)
                print(f"  {ap_info['name']}: {ip} ({mac})")
                if comment:
                    print(f"    Comment: {comment}")

    # Check ARP table for any we might have missed
    print("\n=== ARP Table (Additional Devices) ===")
    for arp in arp_table:
        hostname = arp.get('interface', '')
        mac = arp.get('mac-address', '')
        ip = arp.get('address', '')

        if is_likely_ap(hostname, mac, ''):
            # Check if already in list
            if not any(ap['mac'] == mac for ap in access_points):
                ap_info = {
                    'name': hostname or mac,
                    'ip': ip,
                    'mac': mac,
                    'source': 'arp',
                    'interface': hostname
                }
                access_points.append(ap_info)
                print(f"  {ap_info['name']}: {ip} ({mac})")

    # Print summary
    print(f"\n=== Summary ===")
    print(f"Found {len(access_points)} access points\n")

    # Print clean list
    print("=== Access Point List ===")
    for ap in sorted(access_points, key=lambda x: x['ip']):
        print(f"{ap['name']:<40} {ap['ip']:<15} {ap['mac']}")

    # Save to file if requested
    if args.output:
        with open(args.output, 'w') as f:
            json.dump(access_points, f, indent=2)
        print(f"\nData saved to: {args.output}")

    api.close()

if __name__ == '__main__':
    main()
