import socket
import threading
import argparse
import sys

LISTEN_HOST = '0.0.0.0'
LISTEN_PORT = 4222
REAL_PORT = 4222  # Fixed NATS port

def handle_client(client_sock, real_host):
    remote_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        remote_sock.connect((real_host, REAL_PORT))
    except Exception as e:
        print(f"[!] Failed to connect to real server {real_host}:{REAL_PORT} - {e}")
        client_sock.close()
        return

    def forward(src, dst):
        try:
            while True:
                data = src.recv(4096)
                if not data:
                    break
                print(f"[DATA] {data.decode(errors='ignore')}")
                dst.sendall(data)
        except Exception:
            pass
        finally:
            src.close()
            dst.close()

    threading.Thread(target=forward, args=(client_sock, remote_sock), daemon=True).start()
    threading.Thread(target=forward, args=(remote_sock, client_sock), daemon=True).start()

def start_proxy(real_host):
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind((LISTEN_HOST, LISTEN_PORT))
    server.listen(5)
    server.settimeout(1.0)
    print(f"[+] Proxy listening on {LISTEN_HOST}:{LISTEN_PORT} -> {real_host}:{REAL_PORT}")

    try:
        while True:
            try:
                client_sock, addr = server.accept()
                print(f"[+] Connection from {addr}")
                threading.Thread(target=handle_client, args=(client_sock, real_host), daemon=True).start()
            except socket.timeout:
                continue
    except KeyboardInterrupt:
        print("\n[!] Shutting down proxy.")
    finally:
        server.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="NATS TCP Proxy")
    parser.add_argument("real_host", help="IP address of the real NATS server")
    args = parser.parse_args()

    start_proxy(args.real_host)