import os
import time
import asyncio
import websockets
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

TARGET_URL = os.getenv("TARGET_URL", "https://omegleweb.io/")
TARGET_WS_URL = os.getenv("TARGET_WS_URL", "wss://omegleweb.io:8443/socket.io/?EIO=4&transport=websocket")
LOCAL_HOST = os.getenv("LOCAL_HOST", "0.0.0.0")
LOCAL_PORT = int(os.getenv("LOCAL_PORT", "8765"))
HEADLESS = os.getenv("HEADLESS", "false").lower() in ("true", "1", "yes")
CF_WAIT_TIME = int(os.getenv("CF_WAIT_TIME", "60"))
SELENIUM_URL = os.getenv("SELENIUM_URL", "http://localhost:4444")

credentials = {
    "user_agent": None,
    "cookies": None
}

def stealth_js(driver):
    driver.execute_script("""
        Object.defineProperty(navigator, 'webdriver', {get: () => undefined});
        Object.defineProperty(navigator, 'plugins', {get: () => [1, 2, 3, 4, 5]});
        Object.defineProperty(navigator, 'languages', {get: () => ['en-US', 'en']});
        window.chrome = {runtime: {}};
        const originalQuery = window.navigator.permissions.query;
        window.navigator.permissions.query = (parameters) => (
            parameters.name === 'notifications' ?
                Promise.resolve({state: Notification.permission}) :
                originalQuery(parameters)
        );
        Object.defineProperty(navigator, 'permissions', {get: () => window.navigator.permissions});
    """)

def extract_credentials():
    print("\n" + "="*50)
    print("[*] PHASE 1: SELENIUM EXTRACTION")
    print("="*50)
    print("[*] Launching Chrome browser...")

    options = Options()
    if HEADLESS:
        options.add_argument("--headless=new")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-dev-shm-usage")
    options.add_argument("--disable-gpu")
    options.add_argument("--disable-software-rasterizer")
    options.add_argument("--disable-extensions")
    options.add_argument("--disable-blink-features=AutomationControlled")
    options.add_argument("--disable-setuid-sandbox")
    options.add_argument("--disable-background-networking")
    options.add_argument("--disable-default-apps")
    options.add_argument("--disable-sync")
    options.add_argument("--disable-translate")
    options.add_argument("--metrics-recording-only")
    options.add_argument("--mute-audio")
    options.add_argument("--no-first-run")
    options.add_argument("--safebrowsing-disable-auto-update")
    options.add_argument("--ignore-certificate-errors")
    options.add_argument("--ignore-ssl-errors")
    options.add_argument("--user-data-dir=/tmp/chrome-data")
    options.add_argument("--disable-features=IsolateOrigins,site-per-process")

    driver = webdriver.Remote(
        command_executor=SELENIUM_URL + "/wd/hub",
        options=options
    )

    stealth_js(driver)

    try:
        print("[*] Navigating to OmegleWeb homepage...")
        driver.get(TARGET_URL)

        wait = WebDriverWait(driver, CF_WAIT_TIME)
        wait.until(EC.presence_of_element_located((By.ID, "logo")))

        print("[*] Time's up! Extracting the goods...")
        user_agent = driver.execute_script("return navigator.userAgent;")
        selenium_cookies = driver.get_cookies()
        cookie_string = "; ".join([f"{c['name']}={c['value']}" for c in selenium_cookies])

        credentials["user_agent"] = user_agent
        credentials["cookies"] = cookie_string

        print("[+] Extraction successful!")
        return True
    except Exception as e:
        print(f"[!] Extraction failed: {e}")
        return False
    finally:
        print("[*] Closing browser...")
        driver.quit()

async def start_bridge():
    async def bridge_handler(local_client):
        print(f"\n[{datetime.now().strftime('%H:%M:%S')}] [+] Local client connected to the bridge!")

        while True:
            headers = {
                "User-Agent": credentials["user_agent"],
                "Cookie": credentials["cookies"]
            }

            try:
                print(f"[*] Attempting tunnel to OmegleWeb...")
                async with websockets.connect(TARGET_WS_URL, additional_headers=headers) as target_server:
                    print("[+] Tunnel established! Relaying messages...")

                    async def forward_local_to_target():
                        async for message in local_client:
                            print(f"[{datetime.now().strftime('%H:%M:%S.%f')[:-3]}] [Local -> Omegle]: {message}")
                            await target_server.send(message)

                    async def forward_target_to_local():
                        async for message in target_server:
                            print(f"[{datetime.now().strftime('%H:%M:%S.%f')[:-3]}] [Omegle -> Local]: {message}")
                            await local_client.send(message)

                    t1 = asyncio.create_task(forward_local_to_target())
                    t2 = asyncio.create_task(forward_target_to_local())

                    done, pending = await asyncio.wait(
                        [t1, t2],
                        return_when=asyncio.FIRST_COMPLETED
                    )

                    for task in pending:
                        task.cancel()

                    if t1 in done and t1.exception() is None:
                        print("[-] Local client disconnected.")
                        break

                    print("[-] OmegleWeb connection closed.")
                    break

            except websockets.exceptions.InvalidStatusCode as e:
                if e.status_code == 403:
                    print(f"[!] HTTP 403 Forbidden: Cookies likely expired. Refreshing...")
                    loop = asyncio.get_running_loop()
                    success = await loop.run_in_executor(None, extract_credentials)
                    if not success:
                        print("[!] Failed to refresh cookies. Retrying in 5s...")
                        await asyncio.sleep(5)
                    continue
                else:
                    print(f"[!] Bridge error (Status {e.status_code}): {e}")
                    break
            except websockets.exceptions.ConnectionClosed:
                print("[-] OmegleWeb disconnected.")
                break
            except Exception as e:
                print(f"[!] Bridge error: {type(e).__name__}: {e}")
                break

    print(f"[*] Starting local WebSocket proxy on ws://{LOCAL_HOST}:{LOCAL_PORT}")
    async with websockets.serve(bridge_handler, LOCAL_HOST, LOCAL_PORT):
        await asyncio.Future()

def main():
    if not extract_credentials():
        print("[!] Initial extraction failed. Exiting.")
        return

    print("\n" + "="*50)
    print("[*] PHASE 2: LAUNCH THE WEBSOCKET PROXY")
    print("="*50)

    try:
        asyncio.run(start_bridge())
    except KeyboardInterrupt:
        print("\n[*] Shutting down...")

if __name__ == '__main__':
    main()