#!/usr/bin/env python3
"""
MITM proxy for btest protocol.
Sits between MikroTik client and MikroTik server, logs all bytes exchanged.

Usage:
    python3 btest_mitm.py --listen 2000 --target 172.16.81.1:2000

Then point MikroTik client at THIS machine's IP on port 2000.
"""
import socket
import select
import sys
import argparse
import time
import threading


def hexdump_line(data, offset=0):
    """Format a single line of hex dump."""
    hex_part = " ".join(f"{b:02x}" for b in data)
    ascii_part = "".join(chr(b) if 32 <= b < 127 else "." for b in data)
    return f"  {offset:04x}  {hex_part:<48s}  {ascii_part}"


def hexdump(label, data):
    """Pretty hex dump with ASCII."""
    ts = time.strftime("%H:%M:%S", time.localtime())
    print(f"\n[{ts}] {label} ({len(data)} bytes)")
    for i in range(0, len(data), 16):
        print(hexdump_line(data[i : i + 16], i))
    sys.stdout.flush()


def annotate_btest(direction, data, state):
    """Try to annotate known btest protocol messages."""
    if len(data) == 4:
        val = data.hex()
        if val == "01000000":
            return "HELLO / AUTH_OK"
        elif val == "02000000":
            return "AUTH_REQUIRED (MD5)"
        elif val == "03000000":
            state["auth_type"] = "ecsrp5"
            return "AUTH_REQUIRED (EC-SRP5)"
        elif val == "00000000":
            return "AUTH_FAILED"

    if len(data) == 16 and state.get("stage") == "command":
        proto = "UDP" if data[0] == 0 else "TCP"
        dirs = {1: "RX", 2: "TX", 3: "BOTH"}
        d = dirs.get(data[1], f"0x{data[1]:02x}")
        conn = data[3]
        return f"COMMAND: proto={proto} dir={d} conn_count={conn}"

    if len(data) == 2 and state.get("auth_type") == "ecsrp5":
        return f"UDP port assignment: {int.from_bytes(data, 'big')}"

    if state.get("auth_type") == "ecsrp5" and state.get("stage") == "ecsrp5":
        if len(data) >= 33:
            return f"EC-SRP5 message (likely pubkey + parity/salt)"

    return None


def proxy_connection(client_sock, client_addr, target_host, target_port):
    """Handle one proxied connection."""
    print(f"\n{'='*60}")
    print(f"New connection from {client_addr[0]}:{client_addr[1]}")
    print(f"Proxying to {target_host}:{target_port}")
    print(f"{'='*60}")

    try:
        server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_sock.settimeout(30)
        server_sock.connect((target_host, target_port))
        server_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    except Exception as e:
        print(f"Failed to connect to target: {e}")
        client_sock.close()
        return

    state = {"stage": "hello", "msg_count": 0}

    try:
        while True:
            readable, _, _ = select.select(
                [client_sock, server_sock], [], [], 30
            )

            if not readable:
                print("\n[TIMEOUT] No data for 30 seconds")
                break

            for sock in readable:
                if sock is server_sock:
                    data = server_sock.recv(65536)
                    if not data:
                        print("\n[SERVER CLOSED]")
                        return
                    direction = "SERVER → CLIENT"
                    annotation = annotate_btest("s2c", data, state)
                    hexdump(direction, data)
                    if annotation:
                        print(f"  >>> {annotation}")

                    # Track state
                    if state["stage"] == "hello" and len(data) == 4:
                        state["stage"] = "command"
                    elif state["stage"] == "auth_resp":
                        state["stage"] = "ecsrp5"

                    # Forward to client
                    client_sock.sendall(data)

                elif sock is client_sock:
                    data = client_sock.recv(65536)
                    if not data:
                        print("\n[CLIENT CLOSED]")
                        return
                    direction = "CLIENT → SERVER"
                    annotation = annotate_btest("c2s", data, state)
                    hexdump(direction, data)
                    if annotation:
                        print(f"  >>> {annotation}")

                    # Track state
                    if state["stage"] == "command" and len(data) == 16:
                        state["stage"] = "auth_resp"
                    state["msg_count"] += 1

                    # Forward to server
                    server_sock.sendall(data)

    except (ConnectionResetError, BrokenPipeError) as e:
        print(f"\n[CONNECTION ERROR] {e}")
    except Exception as e:
        print(f"\n[ERROR] {e}")
        import traceback
        traceback.print_exc()
    finally:
        client_sock.close()
        server_sock.close()
        print(f"\n{'='*60}")
        print(f"Connection closed ({state['msg_count']} client messages)")
        print(f"{'='*60}")


def main():
    parser = argparse.ArgumentParser(description="btest MITM proxy")
    parser.add_argument(
        "-l", "--listen", type=int, default=2000, help="Listen port (default: 2000)"
    )
    parser.add_argument(
        "-t",
        "--target",
        required=True,
        help="Target MikroTik address:port (e.g., 172.16.81.1:2000)",
    )
    args = parser.parse_args()

    target_parts = args.target.split(":")
    target_host = target_parts[0]
    target_port = int(target_parts[1]) if len(target_parts) > 1 else 2000

    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind(("0.0.0.0", args.listen))
    listener.listen(5)

    print(f"MITM proxy listening on 0.0.0.0:{args.listen}")
    print(f"Forwarding to {target_host}:{target_port}")
    print(f"Point MikroTik client at this machine's IP, port {args.listen}")
    print()

    while True:
        client_sock, client_addr = listener.accept()
        client_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        t = threading.Thread(
            target=proxy_connection,
            args=(client_sock, client_addr, target_host, target_port),
            daemon=True,
        )
        t.start()


if __name__ == "__main__":
    main()