#!/usr/bin/env python3
"""MQTT Broker Scanner & Auth Tester.

Scans VLANs for MQTT brokers (port 1883/8883) and tests connectivity
by sending a raw MQTT v3.1.1 CONNECT packet and reading the CONNACK.

Usage:
    python3 mqtt-broker-scan.py                    # Scan known homelab VLANs
    python3 mqtt-broker-scan.py 10.0.30.0/24       # Scan specific subnet
    python3 mqtt-broker-scan.py 10.0.30.10 1883    # Test specific broker
"""

import socket
import struct
import sys
import concurrent.futures


KNOWN_SUBNETS = [
    "10.0.30.",   # Server/HA VLAN
    "10.0.40.",   # Guest/IoT VLAN
    "10.0.50.",   # IoT VLAN
]


def check_port(ip, port, timeout=0.3):
    """Quick TCP connect check."""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(timeout)
        result = s.connect_ex((ip, port))
        s.close()
        return result == 0
    except Exception:
        return False


def encode_remaining_length(length):
    """MQTT variable-length encoding for remaining length field."""
    bytes_rl = []
    while True:
        byte = length % 128
        length = length // 128
        if length > 0:
            byte += 128
        bytes_rl.append(byte)
        if length == 0:
            break
    return bytes(bytes_rl)


def mqtt_connect_test(host, port=1883, client_id="hermes-scan", timeout=5):
    """Send minimal MQTT v3.1.1 CONNECT, read CONNACK, return status string."""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(timeout)
        s.connect((host, port))

        # Variable header: Protocol Name + Level + Flags + Keep Alive
        var_header = (
            b'\x00\x04MQTT'   # Protocol name "MQTT"
            b'\x04'            # Protocol level (3.1.1)
            b'\x02'            # Flags: Clean Session
            b'\x00\x3c'        # Keep alive: 60s
        )

        # Payload: Client ID (length-prefixed)
        cid = client_id.encode('utf-8')
        payload = struct.pack('>H', len(cid)) + cid

        # Fixed header: CONNECT (0x10) + remaining length
        remaining = len(var_header) + len(payload)
        fixed_header = b'\x10' + encode_remaining_length(remaining)

        s.send(fixed_header + var_header + payload)

        # Read CONNACK: 4 bytes (0x20, 0x02, session_present, return_code)
        connack = s.recv(4)
        s.close()

        if len(connack) >= 4 and connack[0] == 0x20:
            rc = connack[3]
            codes = {
                0: 'Accepted',
                1: 'Bad protocol version',
                2: 'Client ID rejected',
                3: 'Server unavailable',
                4: 'Bad username/password',
                5: 'Not authorized',
            }
            return codes.get(rc, f'Unknown ({rc})')
        else:
            return f'Unexpected response: {connack.hex()}'
    except Exception as e:
        return f'Connection failed: {e}'


def scan_subnet(prefix, ports=(1883, 8883)):
    """Scan a /24 subnet for open MQTT ports. Returns list of (ip, port) tuples."""
    found = []

    def scan(ip):
        hits = []
        for p in ports:
            if check_port(ip, p):
                hits.append(p)
        return (ip, hits)

    with concurrent.futures.ThreadPoolExecutor(max_workers=128) as ex:
        futures = {
            ex.submit(scan, f"{prefix}{i}"): i
            for i in range(1, 255)
        }
        for f in concurrent.futures.as_completed(futures):
            ip, hits = f.result()
            if hits:
                found.append((ip, hits))

    return sorted(found)


def main():
    args = sys.argv[1:]

    if len(args) >= 2:
        # Test specific broker: mqtt-broker-scan.py <ip> <port>
        host, port = args[0], int(args[1])
        print(f"Testing {host}:{port}...")
        result = mqtt_connect_test(host, port)
        print(f"  Result: {result}")
        return

    if len(args) == 1 and '/' in args[0]:
        # Scan specific subnet: mqtt-broker-scan.py 10.0.30.0/24
        prefix = args[0].rsplit('.', 1)[0] + '.'
        print(f"Scanning {prefix}0/24 for MQTT brokers...")
        results = scan_subnet(prefix)
    else:
        # Scan known homelab subnets
        print("Scanning known homelab VLANs for MQTT brokers...")
        results = []
        for prefix in KNOWN_SUBNETS:
            print(f"  Scanning {prefix}0/24...", end=' ', flush=True)
            found = scan_subnet(prefix)
            if found:
                print(f"found {len(found)}")
                results.extend(found)
            else:
                print("none")

    print(f"\n=== Results: {len(results)} broker(s) found ===")
    for ip, ports in results:
        for port in ports:
            proto = "MQTT" if port == 1883 else f"MQTT-TLS({port})"
            auth = mqtt_connect_test(ip, port)
            print(f"  {proto} {ip}:{port} -> {auth}")


if __name__ == '__main__':
    main()