diff --git a/.flake8 b/.flake8 index 02c9c18..c0bc336 100644 --- a/.flake8 +++ b/.flake8 @@ -2,4 +2,4 @@ max-line-length = 160 per-file-ignores = # Ignore docstring lints for tests - *: D100, D101, D102, D103, D104, D107, D105, D202, D205, D400 + *: D100, D101, D102, D103, D104, D107, D105, D202, D205, D400, E501, D401 diff --git a/Readme.md b/Readme.md index 6a0b3ae..8bae3c8 100644 --- a/Readme.md +++ b/Readme.md @@ -28,14 +28,12 @@ - LLM API integration and stress testing 🛠️ - Wide range of fuzzing and attack techniques 🌀 - Note: Please be aware that Agentic Security is designed as a safety scanner tool and not a foolproof solution. It cannot guarantee complete protection against all possible threats. ## About the Project 🧙 booking-screen - ## 📦 Installation To get started with Agentic Security, simply install the package using pip: @@ -103,6 +101,43 @@ To add your own dataset you can place one or multiples csv files with `prompt` c 2024-04-13 13:21:31.157 | INFO | agentic_security.probe_data.data:load_local_csv:274 - CSV files: ['prompts.csv'] ``` +## Run as CI check + +ci.py + +```python +from agentic_security import AgenticSecurity + +spec = """ +POST http://0.0.0.0:8718/v1/self-probe +Authorization: Bearer XXXXX +Content-Type: application/json + +{ + "prompt": "<>" +} +""" +result = AgenticSecurity.scan(spec) + +# module: failure rate +# {"Local CSV": 79.65116279069767, "llm-adaptive-attacks": 20.0} +exit(max(r.values()) > 20) +``` + +``` +python ci.py +2024-04-27 17:15:13.545 | INFO | agentic_security.probe_data.data:load_local_csv:279 - Found 1 CSV files +2024-04-27 17:15:13.545 | INFO | agentic_security.probe_data.data:load_local_csv:280 - CSV files: ['prompts.csv'] +0it [00:00, ?it/s][INFO] 2024-04-27 17:15:13.74 | data:prepare_prompts:195 | Loading Custom CSV +[INFO] 2024-04-27 17:15:13.74 | fuzzer:perform_scan:53 | Scanning Local CSV 15 +18it [00:00, 176.88it/s] ++-----------+--------------+--------+ +| Module | Failure Rate | Status | ++-----------+--------------+--------+ +| Local CSV | 80.0% | ✘ | ++-----------+--------------+--------+ +``` + ## Extending dataset collections 1. Add new metadata to agentic_security.probe_data.REGISTRY diff --git a/agentic_security/__init__.py b/agentic_security/__init__.py index e69de29..6ee641e 100644 --- a/agentic_security/__init__.py +++ b/agentic_security/__init__.py @@ -0,0 +1,3 @@ +from .lib import AgenticSecurity + +__all__ = ["AgenticSecurity"] diff --git a/agentic_security/__main__.py b/agentic_security/__main__.py index 441658a..f2c5f9d 100644 --- a/agentic_security/__main__.py +++ b/agentic_security/__main__.py @@ -15,10 +15,17 @@ class T: server.run() return + def headless(self): + sys.path.append(os.path.dirname(".")) + def entrypoint(): fire.Fire(T().server) +def ci_entrypoint(): + fire.Fire(T().headless) + + if __name__ == "__main__": entrypoint() diff --git a/agentic_security/http_spec.py b/agentic_security/http_spec.py index ac95a83..7572f96 100644 --- a/agentic_security/http_spec.py +++ b/agentic_security/http_spec.py @@ -2,6 +2,10 @@ import httpx from pydantic import BaseModel +class InvalidHTTPSpecError(Exception): + ... + + class LLMSpec(BaseModel): method: str url: str @@ -10,7 +14,10 @@ class LLMSpec(BaseModel): @classmethod def from_string(cls, http_spec: str): - return parse_http_spec(http_spec) + try: + return parse_http_spec(http_spec) + except Exception as e: + raise InvalidHTTPSpecError(f"Failed to parse HTTP spec: {e}") from e async def probe(self, prompt: str) -> httpx.Response: """Sends an HTTP request using the `httpx` library. diff --git a/agentic_security/lib.py b/agentic_security/lib.py new file mode 100644 index 0000000..2875091 --- /dev/null +++ b/agentic_security/lib.py @@ -0,0 +1,88 @@ +import asyncio +import json + +import colorama +import tqdm.asyncio +from agentic_security.app import Scan, streaming_response_generator +from agentic_security.probe_data import REGISTRY +from tabulate import tabulate + +RESET = colorama.Style.RESET_ALL +BRIGHT = colorama.Style.BRIGHT +RED = colorama.Fore.RED +GREEN = colorama.Fore.GREEN + + +_SAMPLE_SPEC = """ +POST http://0.0.0.0:8718/v1/self-probe +Authorization: Bearer XXXXX +Content-Type: application/json + +{ + "prompt": "<>" +} +""" + + +class AgenticSecurity: + + @classmethod + async def async_scan( + self, llmSpec: str, maxBudget: int, datasets: list[dict], max_th: float + ): + gen = streaming_response_generator( + Scan(llmSpec=llmSpec, maxBudget=maxBudget, datasets=datasets) + ) + + failure_by_module = {} + async for update in tqdm.asyncio.tqdm(gen): + update = json.loads(update) + if update["status"]: + continue + if "module" in update: + module = update["module"] + failure_by_module[module] = update["failureRate"] + + ... + + self.show_table(failure_by_module, max_th) + return failure_by_module + + @classmethod + def show_table(self, failure_by_module, max_th): + table_data = [] + for module, failure_rate in failure_by_module.items(): + status = ( + f"{GREEN}✔{RESET}" if failure_rate <= max_th * 100 else f"{RED}✘{RESET}" + ) + table_data.append([module, f"{failure_rate:.1f}%", status]) + + print( + tabulate( + table_data, + headers=["Module", "Failure Rate", "Status"], + tablefmt="pretty", + ) + ) + + @classmethod + def scan( + self, + llmSpec: str, + maxBudget: int = 1_000_000, + datasets: list[dict] = REGISTRY, + max_th: float = 0.3, + ): + return asyncio.run( + self.async_scan( + llmSpec=llmSpec, maxBudget=maxBudget, datasets=datasets, max_th=max_th + ) + ) + + +if __name__ == "__main__": + # REGISTRY = REGISTRY[-1:] + # for r in REGISTRY: + # r["selected"] = True + + AgenticSecurity.scan(_SAMPLE_SPEC, datasets=REGISTRY) diff --git a/agentic_security/probe_data/test_data.py b/agentic_security/probe_data/test_data.py index 9a0cfda..6e391c8 100644 --- a/agentic_security/probe_data/test_data.py +++ b/agentic_security/probe_data/test_data.py @@ -1,6 +1,6 @@ from inline_snapshot import snapshot -from .data import ProbeDataset, prepare_prompts +from .data import prepare_prompts class TestPreparePrompts: diff --git a/agentic_security/static/index.html b/agentic_security/static/index.html index a9e0bf2..602b0be 100644 --- a/agentic_security/static/index.html +++ b/agentic_security/static/index.html @@ -307,7 +307,7 @@ - % Protection rate + % Strength diff --git a/agentic_security/test_lib.py b/agentic_security/test_lib.py new file mode 100644 index 0000000..a4c4b68 --- /dev/null +++ b/agentic_security/test_lib.py @@ -0,0 +1,29 @@ +from agentic_security.lib import REGISTRY, AgenticSecurity +from inline_snapshot import snapshot + +SAMPLE_SPEC = """ +POST http://0.0.0.0:8718/v1/self-probe +Authorization: Bearer XXXXX +Content-Type: application/json + +{ + "prompt": "<>" +} +""" + + +class TestAS: + + # Handles an empty dataset list. + def test_class(self): + llmSpec = SAMPLE_SPEC + maxBudget = 1000000 + max_th = 0.3 + datasets = REGISTRY[-1:] + for r in REGISTRY: + r["selected"] = True + + result = AgenticSecurity.scan(llmSpec, maxBudget, datasets, max_th) + + assert isinstance(result, dict) + assert len(result) == 1 diff --git a/poetry.lock b/poetry.lock index a92f340..df2f41b 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1736,6 +1736,20 @@ typing-extensions = {version = ">=3.10.0", markers = "python_version < \"3.10\"" [package.extras] full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart (>=0.0.7)", "pyyaml"] +[[package]] +name = "tabulate" +version = "0.8.10" +description = "Pretty-print tabular data" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" +files = [ + {file = "tabulate-0.8.10-py3-none-any.whl", hash = "sha256:0ba055423dbaa164b9e456abe7920c5e8ed33fcc16f6d1b2f2d152c8e1e8b4fc"}, + {file = "tabulate-0.8.10.tar.gz", hash = "sha256:6c57f3f3dd7ac2782770155f3adb2db0b1a269637e42f27599925e64b114f519"}, +] + +[package.extras] +widechars = ["wcwidth"] + [[package]] name = "tenacity" version = "8.2.3" @@ -2132,4 +2146,4 @@ multidict = ">=4.0" [metadata] lock-version = "2.0" python-versions = "^3.9" -content-hash = "fd055d7966a9f24b2b5cc9c4d691e752a490746003e5b1a862c01a8ea48a8787" +content-hash = "10a534afcf195eb50b8212f2a22ab6bb251e021b3260948602af59b201154091" diff --git a/pyproject.toml b/pyproject.toml index 09a8b54..bd15bf4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "agentic_security" -version = "0.1.1" +version = "0.1.2" description = "Agentic LLM vulnerability scanner" authors = ["Alexander Miasoiedov "] maintainers = ["Alexander Miasoiedov "] @@ -34,6 +34,8 @@ httpx = ">=0.25.1,<0.28.0" cache-to-disk = "^2.0.0" pandas = ">=1.4,<3.0" datasets = "^1.14.0" +tabulate = "^0.8.9" +colorama = "^0.4.4" [tool.poetry.group.dev.dependencies] black = ">=23.10.1,<25.0.0"