Update blackgat/app/main.py

blackgat/app/main.py

@@ -1,29 +1,57 @@
-from fastapi import FastAPI, Request
-from app.ai_modules import heatseeker_score, scribe_generate, killchain_ai, recon_gpt, exploit_suggestion
+# main.py
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import Dict
+from ai_models import (
+    scribe_generate, killchain_ai,
+    heatseeker_score, recon_gpt,
+    exploit_suggestion
+)
 
-app = FastAPI()
+app = FastAPI(
+    title="BlackGat AI",
+    description="Cloud-based AI for bug bounty automation",
+    version="1.0.0"
+)
 
-@app.post("/score")
-async def score(request: Request):
-    data = await request.json()
-    return heatseeker_score(data)
+# Models
+class VulnReport(BaseModel):
+    type: str
+    url: str
+    payload: str
+    impact: str
+
+class Finding(BaseModel):
+    data: str
+
+class ExploitRequest(BaseModel):
+    data: str
+
+class HeatInput(BaseModel):
+    url: str
+    params: Dict[str, str]
+    status_code: int
 
-@app.post("/report")
-async def report(request: Request):
-    data = await request.json()
-    return scribe_generate(data)
+class ReconPrompt(BaseModel):
+    prompt: str
 
-@app.post("/chain")
-async def chain(request: Request):
-    data = await request.json()
-    return killchain_ai(data)
+# API Routes
+@app.post("/scribe")
+def scribe(data: VulnReport):
+    return scribe_generate(data.dict())
+
+@app.post("/killchain")
+def killchain(data: Finding):
+    return killchain_ai(data.data)
+
+@app.post("/score")
+def score(data: HeatInput):
+    return heatseeker_score(data.dict())
 
 @app.post("/recon")
-async def recon(request: Request):
-    data = await request.json()
-    return recon_gpt(data["prompt"])
+def recon(data: ReconPrompt):
+    return recon_gpt(data.prompt)
 
 @app.post("/exploit")
-async def exploit(request: Request):
-    data = await request.json()
-    return exploit_suggestion(data)
+def exploit(data: ExploitRequest):
+    return exploit_suggestion(data.data)