Spaces:
Sleeping
Sleeping
File size: 2,740 Bytes
c6e6f10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
import os
import socket
from urllib.parse import urlparse
from models.gnn import PhishingGNN_Model
from pipeline.graph_engine import TopologicalGraphEngine
from config import *
import uvicorn
app = FastAPI(title="Defender V5 Sovereign Threat Analysis API")
# Define node dimension configurations
in_channels_dict = {
'ip': 16,
'domain': 32,
'asn': 8,
'cert': 16
}
# Load model directly into state memory matching training dimension maps
model = PhishingGNN_Model(
metadata=GRAPH_METADATA,
in_channels_dict=in_channels_dict,
hidden_channels=HIDDEN_CHANNELS,
num_heads=NUM_HEADS,
num_layers=NUM_LAYERS,
dropout_rate=0.0
)
if os.path.exists(MODEL_SAVE_PATH):
model.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location='cpu', weights_only=True))
model.eval()
# Payload model expects a URL
class URLPayload(BaseModel):
url: str
@app.post("/analyze")
async def analyze_url(payload: URLPayload):
try:
# 1. Parse the URL to extract the domain
parsed_url = urlparse(payload.url)
domain = parsed_url.netloc or parsed_url.path.split('/')[0]
if ':' in domain:
domain = domain.split(':')[0]
if not domain:
raise HTTPException(status_code=400, detail="Invalid URL format.")
# 2. Resolve the domain to an IP address
try:
ip = socket.gethostbyname(domain)
except socket.gaierror:
raise HTTPException(status_code=400, detail=f"DNS Resolution failed for domain: {domain}")
# 3. Build telemetry dictionary
telemetry_log = {
"ip": ip,
"domain": domain,
"asn": None
}
# 4. Pass through Graph Engine
engine = TopologicalGraphEngine()
x_dict, edge_index_dict = engine.extract_and_build([telemetry_log])
with torch.no_grad():
raw_scores = model(x_dict, edge_index_dict)
threat_probability = float(raw_scores.max().item())
return {
"input_url": payload.url,
"resolved_domain": domain,
"resolved_ip": ip,
"structural_anomaly_score": round(threat_probability, 5),
"remediation_verdict": "ISOLATE_ROUTING" if threat_probability > 0.70 else "ALLOW"
}
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=f"Graph Engine Exception: {str(e)}")
if __name__ == "__main__":
uvicorn.run(app, host="127.0.0.1", port=8005) |