main.py

#!/usr/bin/env python3
import asyncio
import aiohttp
import uuid
import json
import os
import datetime
import yaml
import sys
from aiohttp import web
from pathlib import Path
import signal
import threading
from flask import Flask, render_template, jsonify

class HTTPProxyLogger:
    def __init__(self, config):
        self.listen_port = config['server']['listen_port']
        self.forward_port = config['server']['forward_port']
        self.host = config['server']['host']
        self.log_dir = Path(config['logging']['log_dir'])
        self.log_dir.mkdir(exist_ok=True)
        self.master_log = self.log_dir / "requests.log"
        
    async def proxy_handler(self, request):
        request_id = str(uuid.uuid4())
        timestamp = datetime.datetime.now().isoformat()
        
        # Extract request data
        request_data = {
            "id": request_id,
            "timestamp": timestamp,
            "method": request.method,
            "path": request.path,
            "query_string": str(request.query_string),
            "headers": dict(request.headers),
            "body": await request.text() if request.can_read_body else None
        }
        
        # Forward the request
        async with aiohttp.ClientSession() as session:
            forward_url = f"http://{self.host}:{self.forward_port}{request.path}"
            if request.query_string:
                forward_url += f"?{request.query_string}"
                
            try:
                async with session.request(
                    method=request.method,
                    url=forward_url,
                    headers={k: v for k, v in request.headers.items() if k.lower() != 'host'},
                    data=await request.read()
                ) as response:
                    response_data = {
                        "status": response.status,
                        "headers": dict(response.headers),
                        "body": await response.text()
                    }
                    
                    # Log to individual file
                    request_file = self.log_dir / f"{request_id}.log"
                    with open(request_file, 'w') as f:
                        json.dump({
                            "request": request_data,
                            "response": response_data
                        }, f, indent=2)
                    
                    # Log to master file with response code
                    with open(self.master_log, 'a') as f:
                        f.write(f"{request_data['method']} {request_data['path']} {request_id} {response_data['status']}\n")
                    
                    # Return response to client
                    return web.Response(
                        status=response.status,
                        headers={k: v for k, v in response.headers.items() if k.lower() not in ['content-encoding', 'transfer-encoding', 'content-length']},
                        body=response_data["body"]
                    )
            except Exception as e:
                error_data = {
                    "error": str(e),
                    "timestamp": timestamp
                }
                
                # Log error
                request_file = self.log_dir / f"{request_id}.log"
                with open(request_file, 'w') as f:
                    json.dump({
                        "request": request_data,
                        "error": error_data
                    }, f, indent=2)
                
                with open(self.master_log, 'a') as f:
                    f.write(f"{request_data['method']} {request_data['path']} {request_id} 502 ERROR\n")
                
                return web.Response(status=502, text=f"Proxy Error: {str(e)}")
    
    async def start(self):
        app = web.Application()
        app.router.add_route('*', '/{path:.*}', self.proxy_handler)
        
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, self.host, self.listen_port)
        
        print(f"HTTP Proxy listening on {self.host}:{self.listen_port}")
        print(f"Forwarding requests to {self.host}:{self.forward_port}")
        print(f"Logging to directory: {self.log_dir}")
        
        await site.start()
        
        # Keep the server running
        while True:
            await asyncio.sleep(3600)

def load_config(config_path):
    try:
        with open(config_path, 'r') as file:
            return yaml.safe_load(file)
    except FileNotFoundError:
        print(f"Error: Configuration file not found at {config_path}")
        sys.exit(1)
    except yaml.YAMLError as e:
        print(f"Error parsing YAML configuration: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"Unexpected error loading configuration: {e}")
        sys.exit(1)

# Flask UI setup
app = Flask(__name__)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/requests')
def get_requests():
    config_path = os.environ.get("PROXY_CONFIG", "config.yaml")
    config = load_config(config_path)
    log_dir = Path(config['logging']['log_dir'])
    
    master_log = log_dir / "requests.log"
    requests_list = []
    
    if master_log.exists():
        with open(master_log, 'r') as f:
            for line in f:
                parts = line.strip().split()
                if len(parts) >= 4:
                    method = parts[0]
                    path = parts[1]
                    request_id = parts[2]
                    status_code = parts[3]
                    
                    has_error = len(parts) > 4 and parts[4] == "ERROR"
                    
                    requests_list.insert(0, {
                        "method": method,
                        "path": path,
                        "id": request_id,
                        "status_code": status_code,
                        "status": "ERROR" if has_error else "OK"
                    })
    
    return jsonify(requests_list)

@app.route('/request/<request_id>')
def get_request_details(request_id):
    config_path = os.environ.get("PROXY_CONFIG", "config.yaml")
    config = load_config(config_path)
    log_dir = Path(config['logging']['log_dir'])
    
    request_file = log_dir / f"{request_id}.log"
    
    if request_file.exists():
        with open(request_file, 'r') as f:
            try:
                data = json.load(f)
                return jsonify(data)
            except json.JSONDecodeError:
                return jsonify({"error": "Invalid JSON in log file"}), 500
    else:
        return jsonify({"error": "Request not found"}), 404

@app.route('/clean', methods=['POST'])
def clean_logs():
    config_path = os.environ.get("PROXY_CONFIG", "config.yaml")
    config = load_config(config_path)
    log_dir = Path(config['logging']['log_dir'])
    
    master_log = log_dir / "requests.log"
    if master_log.exists():
        open(master_log, 'w').close()
    
    for file in log_dir.glob("*.log"):
        if file.name != "requests.log":
            try:
                file.unlink()
            except Exception:
                pass
    
    return jsonify({"success": True})

def start_ui_server(config):
    if not config.get('ui', {}).get('enabled', False):
        return
    
    templates_dir = Path(__file__).parent / "templates"
    templates_dir.mkdir(exist_ok=True)
    
    ui_port = config['ui']['port']
    print(f"Starting UI server on port {ui_port}")
    threading.Thread(target=lambda: app.run(host='0.0.0.0', port=ui_port, debug=False, use_reloader=False)).start()

def main():
    config_path = os.environ.get("PROXY_CONFIG", "config.yaml")
    config = load_config(config_path)
    
    # Start UI server if enabled
    if config.get('ui', {}).get('enabled', False):
        start_ui_server(config)
        print(f"UI server started on port {config['ui']['port']}")
    
    # Start proxy server
    proxy = HTTPProxyLogger(config)
    loop = asyncio.get_event_loop()
    
    # Handle graceful shutdown
    def signal_handler():
        print("\nShutting down...")
        loop.stop()
    
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, signal_handler)
    
    try:
        loop.run_until_complete(proxy.start())
    except KeyboardInterrupt:
        print("\nProxy stopped")
    finally:
        loop.close()

if __name__ == "__main__":
    main()