atharvawarade9807's picture
Update Version_1/main.py
d70ed25 verified
Raw
History Blame Contribute Delete
6.1 kB
import asyncio
import sys
import time
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from urllib.parse import urlparse # Clean URL parsing utility
# --- Keep all your existing structural imports perfectly intact ---
from models.diagnostic_matrix import DiagnosticMatrix, LayerResults
from layers.layer_1_lexical import analyze_lexical_structure
from layers.layer_2_network import analyze_network_footprint
from layers.layer_3_crypto import analyze_crypto_handshake
from layers.layer_4_sandbox import analyze_content_sandbox
from layers.layer_5_context import analyze_global_context
from fastapi.middleware.cors import CORSMiddleware
# 1. Initialize the Master Web Engine
app = FastAPI(
title="Chimera V1 Master Security Server",
description="Unified API hosting Chimera V1 and future engines"
)
app.add_middleware(
CORSMiddleware,
allow_origins=["http://0.0.0.0:8000",
"http://127.0.0.1:8000"], # Update this to your frontend URL in production
allow_credentials=False, # Must be False if mode is 'omit' in JS
allow_methods=["*"],
allow_headers=["*"],
)
# 2. Define what incoming web data must look like
class URLPayload(BaseModel):
url: str
# --- Updated Core Processing Engine with Userinfo Interception ---
async def process_payload(url: str) -> DiagnosticMatrix:
start_time = time.perf_counter()
circuit_breaks = []
# --- SANITIZATION GATE & EXPLOIT DETECTION ---
try:
parsed_url = urlparse(url)
# Check if the userinfo credential parameter is present (e.g., 'youtube' in youtube@evil-site.com)
if parsed_url.username:
exploit_detected = True
# Reconstruct the true underlying network destination host target for downstream validation
target_url = f"{parsed_url.scheme}://{parsed_url.hostname}"
if parsed_url.path:
target_url += parsed_url.path
if parsed_url.query:
target_url += f"?{parsed_url.query}"
else:
exploit_detected = False
target_url = url
except Exception:
exploit_detected = False
target_url = url
async def safe_layer_4(target):
try:
return await asyncio.wait_for(analyze_content_sandbox(target), timeout=0.150)
except asyncio.TimeoutError:
circuit_breaks.append("Layer_4_Sandbox_Timeout")
return None
# All analysis layers now process the TRUE, extracted network host destination target
results = await asyncio.gather(
analyze_lexical_structure(target_url),
analyze_network_footprint(target_url),
analyze_crypto_handshake(target_url),
safe_layer_4(target_url),
analyze_global_context(target_url),
return_exceptions=True
)
safe_results = []
for res in results:
if isinstance(res, Exception):
# Fallback change: Under high stress/failure, assign an ambiguous 0.5
# penalty state instead of a completely clean 0.0 pass score
safe_results.append(0.5)
else:
safe_results.append(res)
l1, l2, l3, l4, l5 = safe_results
# CRITICAL SECURITY PATCH APPLICATION:
# If a userinfo redirection exploit attempt was caught at the gate,
# force maximum penalty points to override deceptive layer responses.
if exploit_detected:
l1 = max(l1 or 0.0, 1.0) # Max out lexical abnormality risk score
l5 = max(l5 or 0.0, 1.0) # Max out suspicious context profile score
circuit_breaks.append("Userinfo_Obfuscation_Intercepted")
valid_scores = [s for s in [l1, l2, l3, l4, l5] if s is not None]
if valid_scores:
base_avg = sum(valid_scores) / len(valid_scores)
highest_risk = max(valid_scores)
# BLENDED SCORING: Combine the average with the absolute highest threat
# to prevent pure averages from diluting a single critical vulnerability.
composite = (base_avg * 0.5) + (highest_risk * 0.5)
else:
composite = 0
# Determine the final system classification parameters
if composite < 0.3:
verdict = "PASS"
elif composite <= 0.7:
verdict = "ESCALATE_TO_V5"
else:
verdict = "QUARANTINE"
execution_time = (time.perf_counter() - start_time) * 1000
return DiagnosticMatrix(
target_url=url, # Returns original string layout for security logging tracking
composite_score=round(composite, 3),
verdict=verdict,
layer_scores=LayerResults(
layer_1_lexical=round(l1, 2) if l1 is not None else 0.0,
layer_2_network=round(l2, 2) if l2 is not None else 0.0,
layer_3_crypto=round(l3, 2) if l3 is not None else 0.0,
layer_4_sandbox=round(l4, 2) if l4 is not None else None,
layer_5_context=round(l5, 2) if l5 is not None else 0.0
),
execution_time_ms=round(execution_time, 2),
circuit_breaks=circuit_breaks
)
# 3. Create the Web Endpoint for Chimera V1 Scanner
@app.post("/predict/chimera-v1")
async def scan_url_endpoint(payload: URLPayload):
url_to_test = payload.url.strip()
if not url_to_test:
raise HTTPException(status_code=400, detail="URL payload cannot be empty")
if not url_to_test.startswith(('http://', 'https://')):
url_to_test = 'http://' + url_to_test
try:
# Run your powerful parallel engine
matrix = await process_payload(url_to_test)
# Return the matrix directly as JSON data to the web client
return matrix
except Exception as e:
raise HTTPException(status_code=500, detail=f"Internal scan failure: {str(e)}")
# Windows event loop fallback adjustments
if sys.platform == 'win32':
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())