"""mitmproxy addon to capture all HTTP/HTTPS/WebSocket traffic from cricfy"""
import json
from datetime import datetime

class WsCapture:
    def __init__(self):
        self.flows = []

    def request(self, flow):
        url = flow.request.pretty_url
        method = flow.request.method
        headers = dict(flow.request.headers)
        content = flow.request.text if flow.request.text else ""
        print(f"\n>>> [{method}] {url}")
        if content:
            print(f"    Body: {content[:500]}")

    def response(self, flow):
        url = flow.request.pretty_url
        status = flow.response.status_code
        headers = dict(flow.response.headers)
        content = flow.response.text if flow.response.text else ""
        print(f"\n<<< [{status}] {url}")
        if content:
            print(f"    Body: {content[:500]}")

    def websocket_message(self, flow):
        msg = flow.messages[-1]
        direction = ">>>" if msg.from_client else "<<<"
        text = msg.content
        print(f"\n{ direction } [WS] {flow.request.pretty_url}")
        try:
            parsed = json.loads(text)
            print(f"    JSON: {json.dumps(parsed, indent=2)[:1000]}")
        except (json.JSONDecodeError, UnicodeDecodeError):
            print(f"    Raw: {text[:500]}")

    def websocket_end(self, flow):
        print(f"\n=== WS Connection Closed: {flow.request.pretty_url} ===")
        print(f"    Total messages: {len(flow.messages)}")

    def error(self, flow):
        print(f"\n!!! Error: {flow.error}")
        print(f"    URL: {flow.request.pretty_url if flow.request else 'N/A'}")

addons = [WsCapture()]