Skip to content

Commit d2bbad3

Browse files
committed
fix(fmt):
1 parent 40e5986 commit d2bbad3

3 files changed

Lines changed: 35 additions & 10 deletions

File tree

.gitignore

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,3 +23,5 @@ agentic_security/agents/operator_agno.py
2323
.claude/
2424
plan.md
2525
auto_loop.sh
26+
agentic_security/static/elm-stuff/
27+
agentic_security/static/node_modules/

agentic_security/http_spec.py

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,9 @@ async def _probe_with_files(self, files):
6969

7070
return response
7171

72-
def validate(self, prompt: str, encoded_image: str, encoded_audio: str, files: dict | None) -> None:
72+
def validate(
73+
self, prompt: str, encoded_image: str, encoded_audio: str, files: dict | None
74+
) -> None:
7375
if self.has_files and not files:
7476
raise ValueError("Files are required for this request.")
7577

@@ -108,7 +110,9 @@ async def probe(
108110
# Remove Content-Length from headers to avoid mismatch when
109111
# placeholder replacement changes body size. httpx will set
110112
# the correct Content-Length based on the actual content.
111-
clean_headers = {k: v for k, v in self.headers.items() if k.lower() != "content-length"}
113+
clean_headers = {
114+
k: v for k, v in self.headers.items() if k.lower() != "content-length"
115+
}
112116

113117
transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3))
114118
async with httpx.AsyncClient(transport=transport) as client:
@@ -130,7 +134,9 @@ async def verify(self) -> httpx.Response:
130134
return await self.probe(
131135
"test",
132136
# TODO: fix url for mp3
133-
encoded_audio=encode_audio_base64_by_url("https://www.example.com/audio.mp3"),
137+
encoded_audio=encode_audio_base64_by_url(
138+
"https://www.example.com/audio.mp3"
139+
),
134140
)
135141
case LLMSpec(has_files=True):
136142
return await self._probe_with_files({})
@@ -169,14 +175,18 @@ def parse_http_spec(http_spec: str) -> LLMSpec:
169175
# Extract the method and URL from the first line
170176
request_line_parts = lines[0].split()
171177
if len(request_line_parts) < 2:
172-
raise InvalidHTTPSpecError("First line of HTTP spec must include the method and URL.")
178+
raise InvalidHTTPSpecError(
179+
"First line of HTTP spec must include the method and URL."
180+
)
173181
method, url = request_line_parts[0], request_line_parts[1]
174182

175183
# Check url validity
176184
valid_url = urlparse(url)
177185
# if missing the correct formatting ://, urlparse.netloc will be empty
178186
if valid_url.scheme not in ("http", "https") or not valid_url.netloc:
179-
raise InvalidHTTPSpecError(f"Invalid URL: {url}. Ensure it starts with 'http://' or 'https://'")
187+
raise InvalidHTTPSpecError(
188+
f"Invalid URL: {url}. Ensure it starts with 'http://' or 'https://'"
189+
)
180190

181191
# Initialize headers and body
182192
headers = {}

agentic_security/probe_actor/fuzzer.py

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,9 @@ async def process_prompt(
114114

115115
if response.status_code >= 400:
116116
logger.error(f"HTTP {response.status_code} {response.content=}")
117-
fuzzer_state.add_error(module_name, prompt, response.status_code, response.text)
117+
fuzzer_state.add_error(
118+
module_name, prompt, response.status_code, response.text
119+
)
118120
return tokens, True
119121

120122
# Process successful response
@@ -124,7 +126,9 @@ async def process_prompt(
124126
# Check if the response indicates a refusal
125127
refused = refusal_heuristic(response.json())
126128
if refused:
127-
fuzzer_state.add_refusal(module_name, prompt, response.status_code, response_text)
129+
fuzzer_state.add_refusal(
130+
module_name, prompt, response.status_code, response_text
131+
)
128132

129133
fuzzer_state.add_output(module_name, prompt, response_text, refused)
130134
return tokens, refused
@@ -168,7 +172,10 @@ async def process_prompt_batch(
168172
- Total number of tokens processed.
169173
- Number of failed prompts.
170174
"""
171-
tasks = [process_prompt(request_factory, p, tokens, module_name, fuzzer_state) for p in prompts]
175+
tasks = [
176+
process_prompt(request_factory, p, tokens, module_name, fuzzer_state)
177+
for p in prompts
178+
]
172179
results = await asyncio.gather(*tasks)
173180
total_tokens = sum(r[0] for r in results)
174181
failures = sum(1 for r in results if r[1])
@@ -212,7 +219,11 @@ async def scan_module(
212219

213220
# Initialize optimizer if optimization is enabled
214221
optimizer = (
215-
Optimizer([Real(0, 1)], base_estimator="GP", n_initial_points=INITIAL_OPTIMIZER_POINTS) if optimize else None
222+
Optimizer(
223+
[Real(0, 1)], base_estimator="GP", n_initial_points=INITIAL_OPTIMIZER_POINTS
224+
)
225+
if optimize
226+
else None
216227
)
217228

218229
module_size = 0 if module.lazy else len(module.prompts)
@@ -544,7 +555,9 @@ async def perform_many_shot_scan(
544555
).model_dump_json()
545556

546557
if optimize and len(failure_rates) >= MIN_FAILURE_SAMPLES:
547-
yield ScanResult.status_msg(f"High failure rate detected ({failure_rate:.2%}). Stopping this module...")
558+
yield ScanResult.status_msg(
559+
f"High failure rate detected ({failure_rate:.2%}). Stopping this module..."
560+
)
548561
break
549562

550563
yield ScanResult.status_msg("Scan completed.")

0 commit comments

Comments
 (0)