main.py

"""
TellusDigital API Response Comparator - Main Application Module

This module implements the FastAPI application for comparing responses from two different APIs.
It provides authentication, WebSocket support for real-time updates, and comparison functionality.
"""

import html
import json
import secrets
from datetime import datetime
from typing import Any

from fastapi import FastAPI, Request, Form, HTTPException, status, WebSocket, Depends
from fastapi.responses import HTMLResponse, RedirectResponse, Response
from fastapi.security import HTTPBasic
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from starlette.middleware.sessions import SessionMiddleware
from starlette.responses import JSONResponse

from utils.logger import structlog
from utils.utility import (
    USER_CREDENTIALS,
    parse_json_input,
    fetch_api_responses,
    compare_responses,
    json_line_diff,
    validate_urls,
    validate_json_inputs
)
from utils.middleware import SmartSchemeMiddleware

log = structlog.get_logger()


class CustomJSONEncoder(json.JSONEncoder):
    """Custom JSON encoder that handles DeepDiff types."""

    def default(self, obj: Any) -> Any:
        """Convert DeepDiff types to JSON-serializable format.
        
        Args:
            obj: The object to serialize
            
        Returns:
            JSON serializable format of the object
        """
        try:
            values = getattr(obj, '_values', None)
            if values is not None:
                return list(values)
            return json.JSONEncoder.default(self, obj)
        except TypeError:
            return str(obj)


# Generate a random secret key
SECRET_KEY = secrets.token_urlsafe(32)

# Session settings
SESSION_TIMEOUT = 20 * 60  # 20 minutes in seconds

# Create FastAPI app
app = FastAPI(
    title="API Response Comparator",
    description="TellusDigital Internal API Comparison Tool",
    version="1.0.0",
    terms_of_service="Internal use only",
    contact={
        "name": "Kumud Raj",
        "organization": "TellusDigital",
        "email": "kumud.raj01@telusdigital.com"
    },
    license_info={
        "name": "Proprietary - TellusDigital",
        "url": "https://www.telusdigital.com/",
    }
)

# custom middleware for URL check/convert
app.add_middleware(SmartSchemeMiddleware)

# Custom middleware for authentication before session middleware
@app.middleware("http")
async def auth_middleware(request: Request, call_next):
    """Authentication middleware to handle redirects consistently."""
    # Allow access to login, static files and root without authentication
    if request.url.path in ["/login", "/"] or request.url.path.startswith("/static/"):
        return await call_next(request)

    # Check authentication for all other routes
    if not hasattr(request, "session") or "username" not in request.session:
        return RedirectResponse(
            url="/login",
            status_code=status.HTTP_302_FOUND
        )

    # Proceed with authenticated request
    response = await call_next(request)

    # Convert 401 responses to 302 redirects
    if response.status_code == status.HTTP_401_UNAUTHORIZED:
        return RedirectResponse(
            url="/login",
            status_code=status.HTTP_302_FOUND
        )

    return response


# Add session middleware after auth middleware
app.add_middleware(
    SessionMiddleware,
    secret_key=SECRET_KEY,
    session_cookie="api_comparator_session",
    max_age=SESSION_TIMEOUT
)

# Mount static files and templates
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")
security = HTTPBasic()

# WebSocket connections store
active_connections: list[WebSocket] = []


# Update get_current_user to return the username directly
async def get_current_user(request: Request) -> str:
    """Validate current user session and return username."""
    if not hasattr(request, "session") or "username" not in request.session:
        return RedirectResponse(
            url="/login",
            status_code=status.HTTP_302_FOUND
        )
    return request.session["username"]


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket) -> None:
    """Handle WebSocket connections for real-time progress updates.
    
    Args:
        websocket: The WebSocket connection
    """
    await websocket.accept()
    active_connections.append(websocket)
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        active_connections.remove(websocket)
    except Exception as e:
        log.error("WebSocket error", error=str(e))
        if websocket in active_connections:
            active_connections.remove(websocket)


async def broadcast_progress(message: str) -> None:
    """Broadcast progress updates to all connected WebSocket clients.
    
    Args:
        message: The status message to broadcast
    """
    connections = active_connections.copy()
    for connection in connections:
        try:
            await connection.send_json({"type": "processing", "status": message})
        except Exception as e:
            log.error("Broadcast error", error=str(e))
            if connection in active_connections:
                active_connections.remove(connection)


@app.get("/", response_class=RedirectResponse)
async def root() -> RedirectResponse:
    """Redirect root path to login page."""
    return RedirectResponse("/login")


@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request) -> Response:
    """Render the login page."""
    return templates.TemplateResponse("login.html", {"request": request})


@app.post("/login")
async def login(
        request: Request,
        username: str = Form(...),
        password: str = Form(...)
) -> Response:
    """Handle user login."""
    if username in USER_CREDENTIALS and USER_CREDENTIALS[username] == password:
        request.session["username"] = username
        request.session["last_activity"] = datetime.now().timestamp()
        return RedirectResponse(url="/compare", status_code=status.HTTP_302_FOUND)

    return templates.TemplateResponse(
        "login.html",
        {"request": request, "error": "Invalid credentials"},
        status_code=status.HTTP_401_UNAUTHORIZED
    )


@app.get("/compare", response_class=HTMLResponse)
async def compare_page(request: Request) -> Response:
    """Render the API comparison page."""
    if not hasattr(request, "session") or "username" not in request.session:
        return RedirectResponse(url="/login", status_code=status.HTTP_302_FOUND)
    return templates.TemplateResponse("compare.html", {
        "request": request,
        "username": request.session["username"],
        "methods": ["GET", "POST", "PUT", "DELETE", "PATCH"]
    })


@app.post("/compare")
async def compare_apis(
        request: Request,
        api1_url: str = Form(...),
        api1_method: str = Form(...),
        api1_payload: str = Form(...),
        api1_headers: str = Form(...),
        api2_url: str = Form(...),
        api2_method: str = Form(...),
        api2_payload: str = Form(...),
        api2_headers: str = Form(...),
        view_mode: str = Form(default="line"),
        username: str = Depends(get_current_user)
) -> Response:
    """Compare responses from two APIs and render the results page."""
    try:
        await broadcast_progress("starting")

        # Validate inputs
        validate_urls(api1_url, api2_url)
        validate_json_inputs(
            (api1_payload, "API 1 Payload"),
            (api2_payload, "API 2 Payload"),
            (api1_headers, "API 1 Headers"),
            (api2_headers, "API 2 Headers")
        )

        await broadcast_progress("validating")

        # Parse inputs
        payload1 = parse_json_input(api1_payload, "API 1 Payload")
        payload2 = parse_json_input(api2_payload, "API 2 Payload")
        headers1 = parse_json_input(api1_headers, "API 1 Headers")
        headers2 = parse_json_input(api2_headers, "API 2 Headers")

        await broadcast_progress("processing")

        # Fetch API responses
        responses = await fetch_api_responses(
            api1_url, api1_method, payload1, headers1,
            api2_url, api2_method, payload2, headers2
        )
        (response1, execution_time_api1), (response2, execution_time_api2) = responses

        await broadcast_progress("comparing")

        # Get JSON data from responses
        json1 = response1['data']
        json2 = response2['data']

        # Generate diffs
        diff_tree = compare_responses(json1, json2, view='tree')
        diff_text = compare_responses(json1, json2, view='text')
        diff_html = json_line_diff(json1, json2)

        # Prepare results
        api1_result = html.escape(json.dumps(json1, indent=4))
        api2_result = html.escape(json.dumps(json2, indent=4))

        diff_result = (
            html.escape(diff_tree) if isinstance(diff_tree, str) else
            html.escape(
                json.dumps(diff_tree, indent=4, cls=CustomJSONEncoder))
        ) if view_mode == "tree" else diff_html

        return templates.TemplateResponse(
            "results.html",
            {
                "request": request,
                "username": username,
                "execution_time_api1": f"{execution_time_api1:.2f}",
                "execution_time_api2": f"{execution_time_api2:.2f}",
                "status_code1": response1['status'],
                "status_code2": response2['status'],
                "api1_result": api1_result,
                "api2_result": api2_result,
                "diff_result": diff_result,
                "view_mode": view_mode,
                "diff_text": html.escape(str(diff_text)),
                "api1_url": api1_url,
                "api2_url": api2_url,
                "api1_method": api1_method,
                "api2_method": api2_method
            }
        )
    except Exception as e:
        return templates.TemplateResponse(
            "compare.html",
            {
                "request": request,
                "username": username,
                "error": str(e),
                "methods": ["GET", "POST", "PUT", "DELETE", "PATCH"]
            }
        )


@app.post("/extend-session")
async def extend_session(
    request: Request,
    username: str = Depends(get_current_user)
) -> JSONResponse:
    """Extend the user session."""
    request.session["last_activity"] = datetime.now().timestamp()
    return JSONResponse({"status": "success"})


@app.get("/logout")
async def logout(request: Request) -> RedirectResponse:
    """Log out the user and clear the session."""
    if not hasattr(request, "session") or "username" not in request.session:
        return RedirectResponse(url="/login", status_code=status.HTTP_302_FOUND)
    request.session.clear()
    return RedirectResponse(url="/login", status_code=status.HTTP_302_FOUND)


@app.post("/download/")
async def download_result(
        request: Request,
        content: str = Form(...),
        filename: str = Form(...),
        username: str = Form(...)
) -> Response:
    """Handle file download requests."""
    try:
        # Verify authentication
        if username != request.session.get("username"):
            log.error("Authentication failed for download", username=username)
            return RedirectResponse(
                url="/login",
                status_code=status.HTTP_302_FOUND
            )

        content_type = "application/json" if filename.endswith(".json") else "text/plain"

        # Handle JSON content
        if content_type == "application/json":
            try:
                # Try to parse and re-format JSON content
                json_data = json.loads(html.unescape(content))  # Unescape HTML entities
                formatted_content = json.dumps(json_data, indent=2)
            except json.JSONDecodeError as e:
                log.error("Invalid JSON content in download", error=str(e))
                return templates.TemplateResponse(
                    "error.html",
                    {
                        "request": request,
                        "error": "Invalid JSON content",
                        "detail": f"The content could not be parsed as JSON: {str(e)}"
                    },
                    status_code=400
                )
        else:
            # For text content, use as is but unescape HTML entities
            formatted_content = html.unescape(content)

        headers = {
            "Content-Disposition": f'attachment; filename="{filename}"',
            "Content-Type": content_type
        }

        return Response(
            content=formatted_content,
            headers=headers
        )

    except HTTPException as he:
        raise he
    except Exception as e:
        log.error("Download error", error=str(e))
        return templates.TemplateResponse(
            "error.html",
            {
                "request": request,
                "error": "Something went wrong",
                "detail": "Failed to process your download request. Please try again."
            },
            status_code=500
        )


@app.exception_handler(401)
async def unauthorized_handler(request: Request, exc: HTTPException) -> Response:
    """Handle unauthorized access errors with redirect."""
    return RedirectResponse(
        url="/login",
        status_code=status.HTTP_302_FOUND
    )


@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception) -> Response:
    """Handle general exceptions."""
    # For authentication errors, redirect to login
    if isinstance(exc, HTTPException) and exc.status_code in [401, 403]:
        return RedirectResponse(
            url="/login",
            status_code=status.HTTP_302_FOUND
        )

    return templates.TemplateResponse(
        "error.html",
        {
            "request": request,
            "error": "An unexpected error occurred. Please try again.",
            "detail": str(exc) if app.debug else None
        },
        status_code=500
    )