#!/usr/bin/env python3
"""
View mitmproxy flows from log file
Usage: python3 view_flows.py [options]

Options:
  --last N      Show last N requests (default: 20)
  --domain X    Filter by domain
  --method X    Filter by method (GET/POST)
  --json        Show only JSON responses
  --full        Show full request/response
"""

import sys
import json
from pathlib import Path
from mitmproxy import io


FLOWS_FILE = "/root/kwork-parser/mitmproxy/flows.mitm"


def view_flows(last_n=20, domain_filter=None, method_filter=None, json_only=False, full=False):
    """View flows from mitmproxy dump."""
    
    if not Path(FLOWS_FILE).exists():
        print(f"❌ Flows file not found: {FLOWS_FILE}")
        print("Make sure mitmproxy is running and capturing traffic.")
        return
    
    flows = []
    
    with open(FLOWS_FILE, 'rb') as f:
        reader = io.FlowReader(f)
        for flow in reader.read():
            from mitmproxy.http import HTTPFlow
            if isinstance(flow, HTTPFlow) and flow.request and flow.response:
                flows.append(flow)
    
    # Filter
    if domain_filter:
        flows = [f for f in flows if domain_filter in f.request.url]
    
    if method_filter:
        flows = [f for f in flows if f.request.method == method_filter]
    
    if json_only:
        flows = [f for f in flows if f.response.headers.get("Content-Type", "").startswith("application/json")]
    
    # Sort by time (newest first)
    flows.sort(key=lambda x: x.request.timestamp_start, reverse=True)
    
    # Limit
    flows = flows[:last_n]
    
    # Display
    print(f"📊 Showing {len(flows)} flows (newest first)\n")
    
    for i, flow in enumerate(flows, 1):
        method = flow.request.method
        url = flow.request.url
        status = flow.response.status_code
        size = len(flow.response.content) if flow.response.content else 0
        content_type = flow.response.headers.get("Content-Type", "unknown").split(";")[0]
        time_ms = int((flow.response.timestamp_end - flow.request.timestamp_start) * 1000) if flow.response.timestamp_end else 0
        
        # Color coding
        status_color = "🟢" if status < 400 else "🟡" if status < 500 else "🔴"
        
        print(f"{i}. {status_color} {method} {url[:80]}")
        print(f"   Status: {status} | Size: {size:,} bytes | Time: {time_ms}ms | Type: {content_type}")
        
        if full:
            print(f"\n   Request Headers:")
            for k, v in flow.request.headers.items():
                print(f"     {k}: {v[:100]}")
            
            if flow.request.method == "POST" and flow.request.content:
                print(f"\n   Request Body:")
                print(f"     {flow.request.get_text()[:500]}")
            
            print(f"\n   Response:")
            print(f"     {flow.response.get_text()[:500]}")
        
        print()


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="View mitmproxy flows")
    parser.add_argument("--last", type=int, default=20, help="Show last N requests")
    parser.add_argument("--domain", type=str, help="Filter by domain")
    parser.add_argument("--method", type=str, help="Filter by method")
    parser.add_argument("--json", action="store_true", help="Show only JSON responses")
    parser.add_argument("--full", action="store_true", help="Show full request/response")
    
    args = parser.parse_args()
    
    view_flows(
        last_n=args.last,
        domain_filter=args.domain,
        method_filter=args.method,
        json_only=args.json,
        full=args.full
    )