Enhance test executor and pipeline service functionality

Av7danger · Av7danger · commit 2be88b423480 · 2026-03-08T10:57:51.000+05:30
- Introduced ExecutionStats dataclass to track test execution metrics.
- Refactored execute_proactive_tests to return execution statistics alongside results.
- Added detailed execution logging for proactive tests.
- Implemented new test cases for identity context handling and classification of findings.
- Enhanced pipeline service to enforce request budgets and manage test case branching.
diff --git a/src/secnodeapi/test_executor.py b/src/secnodeapi/test_executor.py
@@ -2,33 +2,43 @@
 Test Executor engine for SecNode API Pentester.
 Asynchronously fires HTTP requests based on generated test cases.
 """
-import time
+
 import asyncio
+import time
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
 import httpx
-from typing import List, Dict, Optional
 import structlog
 
 from .vulnerability_models import TestCase, TestResult
 
 logger = structlog.get_logger(__name__)
 
 
+@dataclass(frozen=True)
+class ExecutionStats:
+    attempted: int
+    successful_requests: int
+    failed_requests: int
+
+
 async def _execute_single_test(
     client: httpx.AsyncClient,
     test_case: TestCase,
     base_url: str,
     auth_headers: Dict[str, str],
-    semaphore: asyncio.Semaphore
+    semaphore: asyncio.Semaphore,
 ) -> Optional[TestResult]:
     """Execute a single test case with retry logic and concurrency control."""
     max_retries = 3
     method = test_case.method.upper()
     url = f"{base_url.rstrip('/')}/{test_case.endpoint.lstrip('/')}"
-    
+
     headers = {**test_case.headers, **auth_headers}
     params = test_case.params
     json_body = test_case.body if isinstance(test_case.body, dict) else None
-    
+
     async with semaphore:
         for attempt in range(max_retries):
             start_time = time.time()
@@ -39,32 +49,34 @@ async def _execute_single_test(
                     headers=headers,
                     params=params,
                     json=json_body,
-                    timeout=10.0
+                    timeout=10.0,
                 )
-                
+
                 duration_ms = (time.time() - start_time) * 1000
-                
+
                 return TestResult(
                     test_case=test_case,
                     status_code=response.status_code,
                     response_body=response.text[:1000] if response.text else "EMPTY",
                     response_headers=dict(response.headers),
                     request_url=str(response.request.url),
                     request_headers=dict(response.request.headers),
-                    request_body=response.request.content.decode('utf-8')[:500] if response.request.content else None,
-                    response_time_ms=duration_ms
+                    request_body=response.request.content.decode("utf-8")[:500]
+                    if response.request.content
+                    else None,
+                    response_time_ms=duration_ms,
                 )
             except (httpx.RequestError, httpx.TimeoutException) as e:
                 logger.warning(
                     "Request failed, retrying",
                     test_id=test_case.id,
                     attempt=attempt + 1,
-                    error=str(e)
+                    error=str(e),
                 )
                 if attempt == max_retries - 1:
                     logger.error("Max retries reached", test_id=test_case.id)
                     return None
-                await asyncio.sleep(2 ** attempt)  # Exponential backoff
+                await asyncio.sleep(2**attempt)  # Exponential backoff
     return None
 
 
@@ -75,29 +87,57 @@ async def execute_proactive_tests(
     auth_headers: Optional[Dict[str, str]] = None,
     proxy: Optional[str] = None,
     verify_ssl: bool = True,
+    max_requests: Optional[int] = None,
 ) -> List[TestResult]:
     """Execute generated test cases concurrently."""
-    logger.info("Executing proactive tests", total=len(test_cases), concurrency=concurrency)
-    
-    results = []
+    results, _ = await execute_proactive_tests_detailed(
+        test_cases=test_cases,
+        base_url=base_url,
+        concurrency=concurrency,
+        auth_headers=auth_headers,
+        proxy=proxy,
+        verify_ssl=verify_ssl,
+        max_requests=max_requests,
+    )
+    return results
+
+
+async def execute_proactive_tests_detailed(
+    test_cases: List[TestCase],
+    base_url: str,
+    concurrency: int = 5,
+    auth_headers: Optional[Dict[str, str]] = None,
+    proxy: Optional[str] = None,
+    verify_ssl: bool = True,
+    max_requests: Optional[int] = None,
+) -> Tuple[List[TestResult], ExecutionStats]:
+    """Execute generated tests and return both results and execution stats."""
+    selected_cases = test_cases[:max_requests] if max_requests is not None else test_cases
+    logger.info("Executing proactive tests", total=len(selected_cases), concurrency=concurrency)
+
+    results: List[TestResult] = []
     auth_headers = auth_headers or {}
     semaphore = asyncio.Semaphore(concurrency)
-    
-    # We use a single AsyncClient for connection pooling optimizations
+
+    # We use a single AsyncClient for connection pooling optimizations.
     async with httpx.AsyncClient(verify=verify_ssl, proxy=proxy) as client:
         tasks = [
             _execute_single_test(client, tc, base_url, auth_headers, semaphore)
-            for tc in test_cases
+            for tc in selected_cases
         ]
-        
-        # gather and filter out any None from failures
+
         outcomes = await asyncio.gather(*tasks, return_exceptions=True)
-        
-        for idx, outcome in enumerate(outcomes):
+
+        for outcome in outcomes:
             if isinstance(outcome, Exception):
                 logger.error("Uncaught exception in test executor worker", exc_info=outcome)
             elif outcome is not None:
                 results.append(outcome)
 
     logger.info("Finished executing tests", successful_requests=len(results))
-    return results
+    stats = ExecutionStats(
+        attempted=len(selected_cases),
+        successful_requests=len(results),
+        failed_requests=max(0, len(selected_cases) - len(results)),
+    )
+    return results, stats
diff --git a/tests/test_ai_modules.py b/tests/test_ai_modules.py
@@ -2,7 +2,7 @@
 
 from secnodeapi.ai.generate import generate_test_cases
 from secnodeapi.ai.understand import understand_api_with_ai
-from secnodeapi.ai.validate import validate_findings_with_ai
+from secnodeapi.ai.validate import classify_findings, validate_findings_with_ai
 from secnodeapi.vulnerability_models import (
     APIUnderstanding,
     SchemaStructure,
@@ -86,3 +86,69 @@ async def fake_call_llm(*args, **kwargs):
     findings = await validate_findings_with_ai([result])
     assert len(findings) == 1
     assert findings[0].test_case_id == "T-1"
+
+
+@pytest.mark.asyncio
+async def test_classify_findings_deterministic_validator_no_llm(monkeypatch) -> None:
+    test_case = TestCase(
+        id="T-DET-1",
+        name="BOLA object access",
+        description="idor",
+        owasp_category="API1: Broken Object Level Authorization",
+        endpoint="/users/2",
+        method="GET",
+    )
+    result = TestResult(
+        test_case=test_case,
+        status_code=200,
+        response_body='{"id":2}',
+        response_headers={},
+        request_url="https://api.example.com/users/2",
+        request_headers={},
+        request_body=None,
+        response_time_ms=8.0,
+    )
+
+    async def fail_if_called(*args, **kwargs):
+        raise AssertionError("LLM should not be called for deterministic finding")
+
+    monkeypatch.setattr("secnodeapi.ai.validate.call_llm", fail_if_called)
+    confirmed, suspected = await classify_findings([result])
+    assert len(confirmed) == 1
+    assert confirmed[0].validation_source == "deterministic"
+    assert suspected == []
+
+
+@pytest.mark.asyncio
+async def test_classify_findings_low_confidence_goes_to_suspected(monkeypatch) -> None:
+    test_case = TestCase(
+        id="T-SUS-1",
+        name="custom logic probe",
+        description="custom category",
+        owasp_category="BIZ-LOGIC",
+        endpoint="/orders",
+        method="GET",
+    )
+    result = TestResult(
+        test_case=test_case,
+        status_code=200,
+        response_body="ok",
+        response_headers={},
+        request_url="https://api.example.com/orders",
+        request_headers={},
+        request_body=None,
+        response_time_ms=10.0,
+    )
+
+    async def fake_call_llm(*args, **kwargs):
+        return (
+            '{"analysis":"weak signal","is_vulnerable":true,"cvss_score":5.0,'
+            '"cvss_vector":"CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:L/I:L/A:N",'
+            '"description":"possible issue","remediation":"review","confidence":0.62}'
+        )
+
+    monkeypatch.setattr("secnodeapi.ai.validate.call_llm", fake_call_llm)
+    confirmed, suspected = await classify_findings([result])
+    assert confirmed == []
+    assert len(suspected) == 1
+    assert suspected[0].validation_source == "ai-suspected"
diff --git a/tests/test_cli_and_config.py b/tests/test_cli_and_config.py
@@ -34,6 +34,24 @@ def test_parse_auth_file(tmp_path: Path) -> None:
     assert headers == {"X-Key": "abc"}
 
 
+def test_parse_identities_file(tmp_path: Path) -> None:
+    identities_file = tmp_path / "identities.json"
+    identities_file.write_text(
+        json.dumps(
+            {
+                "identities": [
+                    {"name": "anon", "headers": {}},
+                    {"name": "user", "headers": {"Authorization": "Bearer x"}},
+                ]
+            }
+        ),
+        encoding="utf-8",
+    )
+    identities = cli.parse_identities(str(identities_file))
+    assert len(identities) == 2
+    assert identities[1].name == "user"
+
+
 def test_require_provider_key_raises(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("OPENAI_API_KEY", raising=False)
     monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
@@ -59,6 +77,7 @@ def test_parse_args_with_dry_run_output(monkeypatch: pytest.MonkeyPatch) -> None
     assert args.target == "https://api.example.com/swagger.json"
     assert args.dry_run is True
     assert args.dry_run_output == "out.json"
+    assert args.mode == "agent"
 
 
 def test_parse_args_requires_dry_run_for_output(monkeypatch: pytest.MonkeyPatch) -> None:
diff --git a/tests/test_pipeline_service.py b/tests/test_pipeline_service.py
