feat(add scan-csv api route):

This commit is contained in:
Alexander Myasoedov
2025-02-14 11:40:55 +02:00
parent 0987f05c4d
commit 4ef7473a56
3 changed files with 51 additions and 3 deletions
+4 -2
View File
@@ -5,6 +5,8 @@ from typing import Protocol
class IntegrationProto(Protocol): class IntegrationProto(Protocol):
def __init__( def __init__(
self, prompt_groups: list, tools_inbox: asyncio.Queue, opts: dict = {} self, prompt_groups: list, tools_inbox: asyncio.Queue, opts: dict = {}
): ... ):
...
async def apply(self) -> list: ... async def apply(self) -> list:
...
+26 -1
View File
@@ -1,6 +1,6 @@
from datetime import datetime from datetime import datetime
from fastapi import APIRouter, BackgroundTasks, HTTPException from fastapi import APIRouter, BackgroundTasks, File, HTTPException, Query, UploadFile
from fastapi.responses import StreamingResponse from fastapi.responses import StreamingResponse
from ..core.app import get_stop_event, get_tools_inbox, set_current_run from ..core.app import get_stop_event, get_tools_inbox, set_current_run
@@ -52,3 +52,28 @@ async def scan(scan_parameters: Scan, background_tasks: BackgroundTasks):
async def stop_scan(): async def stop_scan():
get_stop_event().set() get_stop_event().set()
return {"status": "Scan stopped"} return {"status": "Scan stopped"}
@router.post("/scan-csv")
async def scan_csv(
background_tasks: BackgroundTasks,
file: UploadFile = File(...),
llmSpec: UploadFile = File(...),
optimize: bool = Query(False),
maxBudget: int = Query(10_000),
enableMultiStepAttack: bool = Query(False),
):
# TODO: content dataset to fuzzer
content = await file.read()
llm_spec = await llmSpec.read()
scan_parameters = Scan(
llmSpec=llm_spec,
optimize=optimize,
maxBudget=1000,
enableMultiStepAttack=enableMultiStepAttack,
)
return StreamingResponse(
streaming_response_generator(scan_parameters), media_type="application/json"
)
+21
View File
@@ -0,0 +1,21 @@
import agentic_security.test_spec_assets as test_spec_assets
from agentic_security.routes.scan import router
from fastapi.testclient import TestClient
client = TestClient(router)
def test_upload_csv_and_run():
# Create a sample CSV content
csv_content = "id,prompt\nspec1,value1\nspec2,value3"
# Send a POST request to the /upload-csv endpoint
response = client.post(
"/scan-csv?optimize=false&enableMultiStepAttack=false&maxBudget=1000",
files={
"file": ("test.csv", csv_content, "text/csv"),
"llmSpec": ("spec.txt", test_spec_assets.SAMPLE_SPEC, "text/plain"),
},
)
assert response.status_code == 200
assert "Scan completed." in response.text