#!/usr/bin/env python3

import socket
import sys
import signal
import threading
from scapy.all import sniff, IP, TCP

# Capture Configuration
CAPTURE_INTERFACE = "enp2s0"  # Your promiscuous interface
SOURCE_IP = "172.16.0.108"    # Source IP of the device sending data
TARGET_PORT = 2201            # The port you're capturing
LOCAL_PORT = 3000             # Local port to serve the captured data on

# Global variables
clients = []
clients_lock = threading.Lock()

def signal_handler(sig, frame):
    print('Stopping capture and closing connections...')
    # Close local client connections
    with clients_lock:
        for client in clients:
            try:
                client.close()
            except:
                pass
    sys.exit(0)

def packet_handler(packet):
    """Process each captured packet and forward its payload"""
    if IP in packet and TCP in packet:
        if packet[IP].src == SOURCE_IP and packet[TCP].dport == TARGET_PORT:
            if packet[TCP].payload:
                payload = bytes(packet[TCP].payload)
                if payload:
                    # Forward payload to all connected clients
                    with clients_lock:
                        disconnected = []
                        for client in clients:
                            try:
                                client.sendall(payload)
                            except:
                                disconnected.append(client)

                        # Remove disconnected clients
                        for client in disconnected:
                            clients.remove(client)

def client_handler(client_socket, addr):
    """Handle a connected client"""
    print(f"New client connected: {addr}")
    with clients_lock:
        clients.append(client_socket)

    # Keep connection open until client disconnects
    try:
        while True:
            data = client_socket.recv(1024)
            if not data:
                break
    except:
        pass
    finally:
        with clients_lock:
            if client_socket in clients:
                clients.remove(client_socket)
        client_socket.close()
        print(f"Client disconnected: {addr}")

def start_server():
    """Start a TCP server to serve captured data"""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(('0.0.0.0', LOCAL_PORT))
    server.listen(5)
    print(f"Listening for connections on port {LOCAL_PORT}")

    while True:
        client_sock, address = server.accept()
        client_thread = threading.Thread(target=client_handler, args=(client_sock, address))
        client_thread.daemon = True
        client_thread.start()

if __name__ == "__main__":
    # Set up signal handler for clean shutdown
    signal.signal(signal.SIGINT, signal_handler)

    # Start TCP server in a separate thread
    server_thread = threading.Thread(target=start_server)
    server_thread.daemon = True
    server_thread.start()

    print(f"Starting payload capture from {SOURCE_IP}:{TARGET_PORT}")
    print(f"Sniffing on {CAPTURE_INTERFACE}...")

    # Start packet capture
    sniff(
        iface=CAPTURE_INTERFACE,
        filter=f"tcp and src host {SOURCE_IP} and dst port {TARGET_PORT}",
        prn=packet_handler,
        store=0
    )