From 71787c6ec9331e29ea73ea2becea156e5b371513 Mon Sep 17 00:00:00 2001 From: nemanjaASE <93867316+nemanjaASE@users.noreply.github.com> Date: Sun, 9 Mar 2025 13:50:18 +0100 Subject: [PATCH 1/4] Add type annotations to functions and methods for improved clarity and maintainabiliy --- agentic_security/core/app.py | 19 +++++++++++++------ agentic_security/misc/banner.py | 17 +++++++++-------- agentic_security/report_chart.py | 6 ++++-- agentic_security/routes/scan.py | 11 ++++++----- 4 files changed, 32 insertions(+), 21 deletions(-) diff --git a/agentic_security/core/app.py b/agentic_security/core/app.py index 3938ba8..f97f380 100644 --- a/agentic_security/core/app.py +++ b/agentic_security/core/app.py @@ -4,11 +4,18 @@ from asyncio import Event, Queue from fastapi import FastAPI from fastapi.responses import ORJSONResponse +from agentic_security.http_spec import LLMSpec +from typing import Any, Dict + tools_inbox: Queue = Queue() stop_event: Event = Event() current_run: str = {"spec": "", "id": ""} -_secrets = {} +_secrets: dict[str, str] = {} +current_run: Dict[str, int | LLMSpec] = { + "spec": "", + "id": "" +} def create_app() -> FastAPI: """Create and configure the FastAPI application.""" @@ -26,29 +33,29 @@ def get_stop_event() -> Event: return stop_event -def get_current_run() -> str: +def get_current_run() -> Dict[str, int | LLMSpec]: """Get the current run id.""" return current_run -def set_current_run(spec): +def set_current_run(spec : LLMSpec) -> Dict[str, int | LLMSpec]: """Set the current run id.""" current_run["id"] = hash(id(spec)) current_run["spec"] = spec return current_run -def get_secrets(): +def get_secrets() -> dict[str, str]: return _secrets -def set_secrets(secrets): +def set_secrets(secrets : dict[str, str]) -> dict[str, str]: _secrets.update(secrets) expand_secrets(_secrets) return _secrets -def expand_secrets(secrets): +def expand_secrets(secrets : dict[str, str]) -> None: for key in secrets: val = secrets[key] if val.startswith("$"): diff --git a/agentic_security/misc/banner.py b/agentic_security/misc/banner.py index f758c2f..308d42a 100644 --- a/agentic_security/misc/banner.py +++ b/agentic_security/misc/banner.py @@ -1,5 +1,6 @@ from pyfiglet import Figlet, FontNotFound from termcolor import colored +from typing import Optional try: from importlib.metadata import version @@ -8,14 +9,14 @@ except ImportError: def generate_banner( - title="Agentic Security", - font="slant", - version="v2.1.0", - tagline="Proactive Threat Detection & Automated Security Protocols", - author="Developed by: [Security Team]", - website="Website: https://github.com/msoedov/agentic_security", - warning="", -): + title: str = "Agentic Security", + font: str = "slant", + version: str = "v2.1.0", + tagline: str = "Proactive Threat Detection & Automated Security Protocols", + author: str = "Developed by: [Security Team]", + website: str = "Website: https://github.com/msoedov/agentic_security", + warning: Optional[str] = "", # Using Optional for warning since it might be None +) -> str: """Generate a visually enhanced banner with dynamic width and borders.""" # Define the text elements diff --git a/agentic_security/report_chart.py b/agentic_security/report_chart.py index c76e576..4f13ae9 100644 --- a/agentic_security/report_chart.py +++ b/agentic_security/report_chart.py @@ -1,5 +1,6 @@ import io import string +from typing import List import matplotlib.pyplot as plt import numpy as np @@ -7,8 +8,9 @@ import pandas as pd from matplotlib.cm import ScalarMappable from matplotlib.colors import LinearSegmentedColormap, Normalize +from .primitives import Table -def plot_security_report(table): +def plot_security_report(table: Table) -> io.BytesIO: # Data preprocessing data = pd.DataFrame(table) @@ -141,7 +143,7 @@ def plot_security_report(table): return buf -def generate_identifiers(data): +def generate_identifiers(data : pd.DataFrame) -> List[str]: data_length = len(data) alphabet = string.ascii_uppercase num_letters = len(alphabet) diff --git a/agentic_security/routes/scan.py b/agentic_security/routes/scan.py index f5e2977..b4400f1 100644 --- a/agentic_security/routes/scan.py +++ b/agentic_security/routes/scan.py @@ -1,4 +1,5 @@ from datetime import datetime +from typing import Any, Generator from fastapi import ( APIRouter, @@ -24,7 +25,7 @@ router = APIRouter() @router.post("/verify") async def verify( info: LLMInfo, secrets: InMemorySecrets = Depends(get_in_memory_secrets) -): +) -> dict[str, int | str | float]: spec = LLMSpec.from_string(info.spec) try: r = await spec.verify() @@ -42,7 +43,7 @@ async def verify( ) -def streaming_response_generator(scan_parameters: Scan): +def streaming_response_generator(scan_parameters: Scan) -> Generator[str, Any, None]: request_factory = LLMSpec.from_string(scan_parameters.llmSpec) set_current_run(request_factory) @@ -63,7 +64,7 @@ async def scan( scan_parameters: Scan, background_tasks: BackgroundTasks, secrets: InMemorySecrets = Depends(get_in_memory_secrets), -): +) -> StreamingResponse: scan_parameters.with_secrets(secrets) return StreamingResponse( streaming_response_generator(scan_parameters), media_type="application/json" @@ -71,7 +72,7 @@ async def scan( @router.post("/stop") -async def stop_scan(): +async def stop_scan() -> dict[str, str]: get_stop_event().set() return {"status": "Scan stopped"} @@ -85,7 +86,7 @@ async def scan_csv( maxBudget: int = Query(10_000), enableMultiStepAttack: bool = Query(False), secrets: InMemorySecrets = Depends(get_in_memory_secrets), -): +) -> StreamingResponse: # TODO: content dataset to fuzzer content = await file.read() # noqa llm_spec = await llmSpec.read() From 6f3c522d5979437b846739b0054f108da36feece Mon Sep 17 00:00:00 2001 From: Praveenk8051 Date: Sun, 9 Mar 2025 14:01:50 +0100 Subject: [PATCH 2/4] chore: update poetry.lock and pyproject.toml to include pytest-xdist --- poetry.lock | 38 ++++++++++++++++++++++++++++++++++++-- pyproject.toml | 1 + 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/poetry.lock b/poetry.lock index 3b0e30e..389d518 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -786,6 +786,20 @@ files = [ {file = "distlib-0.3.8.tar.gz", hash = "sha256:1530ea13e350031b6312d8580ddb6b27a104275a31106523b8f123787f494f64"}, ] +[[package]] +name = "execnet" +version = "2.1.1" +description = "execnet: rapid multi-Python deployment" +optional = false +python-versions = ">=3.8" +files = [ + {file = "execnet-2.1.1-py3-none-any.whl", hash = "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc"}, + {file = "execnet-2.1.1.tar.gz", hash = "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3"}, +] + +[package.extras] +testing = ["hatch", "pre-commit", "pytest", "tox"] + [[package]] name = "executing" version = "2.2.0" @@ -3244,6 +3258,26 @@ pytest = ">=6.2.5" [package.extras] dev = ["pre-commit", "pytest-asyncio", "tox"] +[[package]] +name = "pytest-xdist" +version = "3.6.1" +description = "pytest xdist plugin for distributed testing, most importantly across multiple CPUs" +optional = false +python-versions = ">=3.8" +files = [ + {file = "pytest_xdist-3.6.1-py3-none-any.whl", hash = "sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7"}, + {file = "pytest_xdist-3.6.1.tar.gz", hash = "sha256:ead156a4db231eec769737f57668ef58a2084a34b2e55c4a8fa20d861107300d"}, +] + +[package.dependencies] +execnet = ">=2.1" +pytest = ">=7.0.0" + +[package.extras] +psutil = ["psutil (>=3.0)"] +setproctitle = ["setproctitle"] +testing = ["filelock"] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -4447,4 +4481,4 @@ propcache = ">=0.2.0" [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "28a2b74bfafa9f93d14d2f8d1fcaffa340db212acce6469d6714d342203ad77f" +content-hash = "35e03dba41d30cf6129a4a4f3107eca560f779205b21d3ffb2871eeffc5d5a64" diff --git a/pyproject.toml b/pyproject.toml index 9ca085c..631ad15 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ pyfiglet = "^1.0.2" termcolor = "^2.4.0" # garak = { version = "*", optional = true } +pytest-xdist = "3.6.1" [tool.poetry.group.dev.dependencies] From 8bd76b9f055090dcf8e1a49f7e15e764d46757d9 Mon Sep 17 00:00:00 2001 From: Praveenk8051 Date: Sun, 9 Mar 2025 14:15:44 +0100 Subject: [PATCH 3/4] chore: update pytest configuration to enable parallel test execution --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 631ad15..14de7ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,7 +87,7 @@ build-backend = "poetry.core.masonry.api" [tool.pytest.ini_options] -addopts = "--durations=5 -m 'not slow'" +addopts = "--durations=5 -m 'not slow' -n auto" asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" markers = "slow: marks tests as slow" From 5a4b5e11b2302118218e4b8619464d835aa499f7 Mon Sep 17 00:00:00 2001 From: nemanjaASE <93867316+nemanjaASE@users.noreply.github.com> Date: Sun, 9 Mar 2025 15:21:50 +0100 Subject: [PATCH 4/4] Fix: Remove unused imports with pycln --- agentic_security/core/app.py | 15 ++++++--------- agentic_security/misc/banner.py | 4 ++-- agentic_security/report_chart.py | 4 ++-- agentic_security/routes/scan.py | 3 ++- 4 files changed, 12 insertions(+), 14 deletions(-) diff --git a/agentic_security/core/app.py b/agentic_security/core/app.py index f97f380..3cb306e 100644 --- a/agentic_security/core/app.py +++ b/agentic_security/core/app.py @@ -5,17 +5,14 @@ from fastapi import FastAPI from fastapi.responses import ORJSONResponse from agentic_security.http_spec import LLMSpec -from typing import Any, Dict tools_inbox: Queue = Queue() stop_event: Event = Event() current_run: str = {"spec": "", "id": ""} _secrets: dict[str, str] = {} -current_run: Dict[str, int | LLMSpec] = { - "spec": "", - "id": "" -} +current_run: dict[str, int | LLMSpec] = {"spec": "", "id": ""} + def create_app() -> FastAPI: """Create and configure the FastAPI application.""" @@ -33,12 +30,12 @@ def get_stop_event() -> Event: return stop_event -def get_current_run() -> Dict[str, int | LLMSpec]: +def get_current_run() -> dict[str, int | LLMSpec]: """Get the current run id.""" return current_run -def set_current_run(spec : LLMSpec) -> Dict[str, int | LLMSpec]: +def set_current_run(spec: LLMSpec) -> dict[str, int | LLMSpec]: """Set the current run id.""" current_run["id"] = hash(id(spec)) current_run["spec"] = spec @@ -49,13 +46,13 @@ def get_secrets() -> dict[str, str]: return _secrets -def set_secrets(secrets : dict[str, str]) -> dict[str, str]: +def set_secrets(secrets: dict[str, str]) -> dict[str, str]: _secrets.update(secrets) expand_secrets(_secrets) return _secrets -def expand_secrets(secrets : dict[str, str]) -> None: +def expand_secrets(secrets: dict[str, str]) -> None: for key in secrets: val = secrets[key] if val.startswith("$"): diff --git a/agentic_security/misc/banner.py b/agentic_security/misc/banner.py index 308d42a..45edba5 100644 --- a/agentic_security/misc/banner.py +++ b/agentic_security/misc/banner.py @@ -1,6 +1,6 @@ + from pyfiglet import Figlet, FontNotFound from termcolor import colored -from typing import Optional try: from importlib.metadata import version @@ -15,7 +15,7 @@ def generate_banner( tagline: str = "Proactive Threat Detection & Automated Security Protocols", author: str = "Developed by: [Security Team]", website: str = "Website: https://github.com/msoedov/agentic_security", - warning: Optional[str] = "", # Using Optional for warning since it might be None + warning: str | None = "", # Using Optional for warning since it might be None ) -> str: """Generate a visually enhanced banner with dynamic width and borders.""" # Define the text elements diff --git a/agentic_security/report_chart.py b/agentic_security/report_chart.py index 4f13ae9..45387cb 100644 --- a/agentic_security/report_chart.py +++ b/agentic_security/report_chart.py @@ -1,6 +1,5 @@ import io import string -from typing import List import matplotlib.pyplot as plt import numpy as np @@ -10,6 +9,7 @@ from matplotlib.colors import LinearSegmentedColormap, Normalize from .primitives import Table + def plot_security_report(table: Table) -> io.BytesIO: # Data preprocessing data = pd.DataFrame(table) @@ -143,7 +143,7 @@ def plot_security_report(table: Table) -> io.BytesIO: return buf -def generate_identifiers(data : pd.DataFrame) -> List[str]: +def generate_identifiers(data: pd.DataFrame) -> list[str]: data_length = len(data) alphabet = string.ascii_uppercase num_letters = len(alphabet) diff --git a/agentic_security/routes/scan.py b/agentic_security/routes/scan.py index b4400f1..8313702 100644 --- a/agentic_security/routes/scan.py +++ b/agentic_security/routes/scan.py @@ -1,5 +1,6 @@ +from collections.abc import Generator from datetime import datetime -from typing import Any, Generator +from typing import Any from fastapi import ( APIRouter,