diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 48484f3..5d0db2e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,26 +1,24 @@ - default_language_version: - python: python3 + python: python3.11 repos: - - repo: https://github.com/asottile/pyupgrade - rev: v2.31.1 + rev: v3.15.0 hooks: - id: pyupgrade - args: [--py39-plus] + args: [--py311-plus] - repo: https://github.com/psf/black - rev: 22.8.0 + rev: 23.11.0 hooks: - id: black - language_version: python3.9 + language_version: python3.11 - repo: https://github.com/pycqa/flake8 - rev: 5.0.4 + rev: 6.1.0 hooks: - id: flake8 - language_version: python3 + language_version: python3.11 additional_dependencies: [flake8-docstrings] - repo: https://github.com/PyCQA/isort @@ -30,7 +28,7 @@ repos: args: [--profile, black] - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.3.0 + rev: v4.5.0 hooks: - id: check-ast exclude: '^(third_party)/' @@ -47,16 +45,15 @@ repos: args: ['--maxkb=100'] - repo: https://github.com/executablebooks/mdformat - rev: 0.7.14 + rev: 0.7.17 hooks: - id: mdformat name: mdformat entry: mdformat . - language_version: python3 - + language_version: python3.11 - repo: https://github.com/hadialqattan/pycln - rev: v2.1.1 # Possible releases: https://github.com/hadialqattan/pycln/releases + rev: v2.4.0 hooks: - id: pycln @@ -65,9 +62,8 @@ repos: hooks: - id: teyit - - repo: https://github.com/python-poetry/poetry - rev: '1.6.0' + rev: '1.7.0' hooks: - id: poetry-check - id: poetry-lock @@ -75,9 +71,8 @@ repos: args: - --check - - repo: https://github.com/codespell-project/codespell - rev: v2.2.5 + rev: v2.2.6 hooks: - id: codespell exclude: '^(third_party/)|(poetry.lock)' diff --git a/agentic_security/app.py b/agentic_security/app.py index 28fd40a..f6388ee 100644 --- a/agentic_security/app.py +++ b/agentic_security/app.py @@ -3,11 +3,11 @@ from .middleware.cors import setup_cors from .middleware.logging import LogNon200ResponsesMiddleware from .routes import ( - static_router, - scan_router, probe_router, proxy_router, report_router, + scan_router, + static_router, ) # Create the FastAPI app diff --git a/agentic_security/core/app.py b/agentic_security/core/app.py index 7a8d55d..57893ca 100644 --- a/agentic_security/core/app.py +++ b/agentic_security/core/app.py @@ -1,4 +1,5 @@ from asyncio import Event, Queue + from fastapi import FastAPI tools_inbox: Queue = Queue() diff --git a/agentic_security/lib.py b/agentic_security/lib.py index 751fe36..46808c8 100644 --- a/agentic_security/lib.py +++ b/agentic_security/lib.py @@ -6,8 +6,8 @@ from tabulate import tabulate from agentic_security.models.schemas import Scan -from agentic_security.routes.scan import streaming_response_generator from agentic_security.probe_data import REGISTRY +from agentic_security.routes.scan import streaming_response_generator RESET = colorama.Style.RESET_ALL BRIGHT = colorama.Style.BRIGHT diff --git a/agentic_security/middleware/logging.py b/agentic_security/middleware/logging.py index 9171a4f..0cf5be8 100644 --- a/agentic_security/middleware/logging.py +++ b/agentic_security/middleware/logging.py @@ -1,6 +1,6 @@ from fastapi import Request -from starlette.middleware.base import BaseHTTPMiddleware from loguru import logger +from starlette.middleware.base import BaseHTTPMiddleware class LogNon200ResponsesMiddleware(BaseHTTPMiddleware): diff --git a/agentic_security/models/schemas.py b/agentic_security/models/schemas.py index 8fa287a..c4c274a 100644 --- a/agentic_security/models/schemas.py +++ b/agentic_security/models/schemas.py @@ -1,4 +1,5 @@ import os + from pydantic import BaseModel, Field diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py index 104270a..15924e9 100644 --- a/agentic_security/probe_actor/fuzzer.py +++ b/agentic_security/probe_actor/fuzzer.py @@ -1,16 +1,17 @@ import asyncio import os import random -from typing import AsyncGenerator +from collections.abc import AsyncGenerator + import httpx import pandas as pd from loguru import logger from skopt import Optimizer from skopt.space import Real +from agentic_security.models.schemas import ScanResult from agentic_security.probe_actor.refusal import refusal_heuristic from agentic_security.probe_data.data import prepare_prompts -from agentic_security.models.schemas import ScanResult IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t" diff --git a/agentic_security/probe_actor/test_refusal.py b/agentic_security/probe_actor/test_refusal.py index 281a202..f2a88e5 100644 --- a/agentic_security/probe_actor/test_refusal.py +++ b/agentic_security/probe_actor/test_refusal.py @@ -2,7 +2,6 @@ class TestCheckRefusal: - # The function correctly identifies a refusal phrase in the response. def test_identify_refusal_phrase(self): response = "I'm sorry, but I cannot provide that information." diff --git a/agentic_security/probe_data/data.py b/agentic_security/probe_data/data.py index f75d63e..8d48d8b 100644 --- a/agentic_security/probe_data/data.py +++ b/agentic_security/probe_data/data.py @@ -152,7 +152,6 @@ def load_dataset_v6(): @cache_to_disk() def load_dataset_v7(): - splits = { "mini_JailBreakV_28K": "JailBreakV_28K/mini_JailBreakV_28K.csv", "JailBreakV_28K": "JailBreakV_28K/JailBreakV_28K.csv", @@ -173,7 +172,6 @@ def load_dataset_v7(): @cache_to_disk() def load_dataset_v8(): - df = pd.read_csv( "hf://datasets/ShawnMenz/jailbreak_sft_rm_ds/jailbreak_sft_rm_ds.csv", names=["jailbreak", "prompt"], @@ -321,7 +319,6 @@ def __iter__(self): def apply(self): for prompt_group in self.prompt_groups: - size = len(prompt_group.prompts) for name, fn in self.fn_library.items(): logger.info(f"Applying {name} to {prompt_group.dataset_name}") diff --git a/agentic_security/probe_data/modules/adaptive_attacks.py b/agentic_security/probe_data/modules/adaptive_attacks.py index 31268fb..7ac0d00 100644 --- a/agentic_security/probe_data/modules/adaptive_attacks.py +++ b/agentic_security/probe_data/modules/adaptive_attacks.py @@ -9,7 +9,6 @@ class Module: def __init__(self, prompt_groups: []): - r = httpx.get(url) content = r.content diff --git a/agentic_security/probe_data/modules/test_adaptive_attacks.py b/agentic_security/probe_data/modules/test_adaptive_attacks.py index 589cae5..10d0630 100644 --- a/agentic_security/probe_data/modules/test_adaptive_attacks.py +++ b/agentic_security/probe_data/modules/test_adaptive_attacks.py @@ -4,7 +4,6 @@ class TestModule: - # Module can be initialized with a list of prompt groups. def test_initialize_with_prompt_groups(self): prompt_groups = [] diff --git a/agentic_security/routes/__init__.py b/agentic_security/routes/__init__.py index d334d65..5997717 100644 --- a/agentic_security/routes/__init__.py +++ b/agentic_security/routes/__init__.py @@ -1,8 +1,8 @@ -from .static import router as static_router -from .scan import router as scan_router from .probe import router as probe_router from .proxy import router as proxy_router from .report import router as report_router +from .scan import router as scan_router +from .static import router as static_router __all__ = [ "static_router", diff --git a/agentic_security/routes/probe.py b/agentic_security/routes/probe.py index b2bbddf..612f03d 100644 --- a/agentic_security/routes/probe.py +++ b/agentic_security/routes/probe.py @@ -1,5 +1,7 @@ import random + from fastapi import APIRouter + from ..models.schemas import Probe from ..probe_actor.refusal import REFUSAL_MARKS from ..probe_data import REGISTRY diff --git a/agentic_security/routes/proxy.py b/agentic_security/routes/proxy.py index 9732ed7..ac75ac5 100644 --- a/agentic_security/routes/proxy.py +++ b/agentic_security/routes/proxy.py @@ -1,9 +1,11 @@ import random from asyncio import Event + from fastapi import APIRouter + +from ..core.app import get_tools_inbox from ..models.schemas import CompletionRequest, Settings from ..probe_actor.refusal import REFUSAL_MARKS -from ..core.app import get_tools_inbox router = APIRouter() diff --git a/agentic_security/routes/report.py b/agentic_security/routes/report.py index 25988aa..f45dbd3 100644 --- a/agentic_security/routes/report.py +++ b/agentic_security/routes/report.py @@ -1,6 +1,8 @@ from pathlib import Path + from fastapi import APIRouter, Response from fastapi.responses import FileResponse, StreamingResponse + from ..models.schemas import Table from ..report_chart import plot_security_report diff --git a/agentic_security/routes/scan.py b/agentic_security/routes/scan.py index cd7062d..46344e6 100644 --- a/agentic_security/routes/scan.py +++ b/agentic_security/routes/scan.py @@ -1,10 +1,12 @@ from datetime import datetime + from fastapi import APIRouter, BackgroundTasks, HTTPException from fastapi.responses import StreamingResponse -from ..models.schemas import LLMInfo, Scan + +from ..core.app import get_stop_event, get_tools_inbox from ..http_spec import LLMSpec +from ..models.schemas import LLMInfo, Scan from ..probe_actor import fuzzer -from ..core.app import get_tools_inbox, get_stop_event router = APIRouter() diff --git a/agentic_security/routes/static.py b/agentic_security/routes/static.py index a97bdf8..3f01eb8 100644 --- a/agentic_security/routes/static.py +++ b/agentic_security/routes/static.py @@ -1,6 +1,8 @@ from pathlib import Path + from fastapi import APIRouter from fastapi.responses import FileResponse + from ..models.schemas import Settings router = APIRouter() diff --git a/agentic_security/test_lib.py b/agentic_security/test_lib.py index c733fe1..da7eda6 100644 --- a/agentic_security/test_lib.py +++ b/agentic_security/test_lib.py @@ -14,7 +14,6 @@ class TestAS: - # Handles an empty dataset list. def test_class(self): llmSpec = SAMPLE_SPEC diff --git a/agentic_security/test_spec.py b/agentic_security/test_spec.py index b916542..46c9687 100644 --- a/agentic_security/test_spec.py +++ b/agentic_security/test_spec.py @@ -2,7 +2,6 @@ class TestParseHttpSpec: - # Should correctly parse a simple HTTP spec with headers and body def test_parse_simple_http_spec(self): http_spec = ( diff --git a/poetry.lock b/poetry.lock index 6f16dc3..6d97aac 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -114,7 +114,6 @@ files = [ [package.dependencies] aiohappyeyeballs = ">=2.3.0" aiosignal = ">=1.1.2" -async-timeout = {version = ">=4.0,<6.0", markers = "python_version < \"3.11\""} attrs = ">=17.3.0" frozenlist = ">=1.1.1" multidict = ">=4.5,<7.0" @@ -160,10 +159,8 @@ files = [ ] [package.dependencies] -exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""} idna = ">=2.8" sniffio = ">=1.1" -typing-extensions = {version = ">=4.1", markers = "python_version < \"3.11\""} [package.extras] doc = ["Sphinx (>=7)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"] @@ -188,17 +185,6 @@ six = ">=1.12.0" astroid = ["astroid (>=1,<2)", "astroid (>=2,<4)"] test = ["astroid (>=1,<2)", "astroid (>=2,<4)", "pytest"] -[[package]] -name = "async-timeout" -version = "4.0.3" -description = "Timeout context manager for asyncio programs" -optional = false -python-versions = ">=3.7" -files = [ - {file = "async-timeout-4.0.3.tar.gz", hash = "sha256:4640d96be84d82d02ed59ea2b7105a0f7b33abe8703703cd0ab0bf87c427522f"}, - {file = "async_timeout-4.0.3-py3-none-any.whl", hash = "sha256:7405140ff1230c310e51dc27b3145b9092d659ce68ff733fb0cefe3ee42be028"}, -] - [[package]] name = "attrs" version = "24.2.0" @@ -255,8 +241,6 @@ mypy-extensions = ">=0.4.3" packaging = ">=22.0" pathspec = ">=0.9.0" platformdirs = ">=2" -tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""} -typing-extensions = {version = ">=4.0.1", markers = "python_version < \"3.11\""} [package.extras] colorama = ["colorama (>=0.4.3)"] @@ -578,20 +562,6 @@ files = [ {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, ] -[[package]] -name = "exceptiongroup" -version = "1.2.2" -description = "Backport of PEP 654 (exception groups)" -optional = false -python-versions = ">=3.7" -files = [ - {file = "exceptiongroup-1.2.2-py3-none-any.whl", hash = "sha256:3111b9d131c238bec2f8f516e123e14ba243563fb135d3fe885990585aa7795b"}, - {file = "exceptiongroup-1.2.2.tar.gz", hash = "sha256:47c2edf7c6738fafb49fd34290706d1a1a2f4d1c6df275526b62cbb4aa5393cc"}, -] - -[package.extras] -test = ["pytest (>=6)"] - [[package]] name = "executing" version = "2.1.0" @@ -1011,8 +981,6 @@ black = ">=23.3.0" click = ">=8.1.4" executing = ">=2.1.0" rich = ">=13.7.1" -toml = {version = ">=0.10.2", markers = "python_version < \"3.11\""} -types-toml = {version = ">=0.10.8.7", markers = "python_version < \"3.11\""} typing-extensions = "*" [[package]] @@ -1504,7 +1472,6 @@ files = [ [package.dependencies] mypy-extensions = ">=1.0.0" -tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""} typing-extensions = ">=4.6.0" [package.extras] @@ -1729,7 +1696,6 @@ files = [ [package.dependencies] numpy = [ {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, - {version = ">=1.22.4", markers = "python_version < \"3.11\""}, {version = ">=1.23.2", markers = "python_version == \"3.11\""}, ] python-dateutil = ">=2.8.2" @@ -2258,11 +2224,9 @@ files = [ [package.dependencies] colorama = {version = "*", markers = "sys_platform == \"win32\""} -exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""} iniconfig = "*" packaging = "*" pluggy = ">=1.5,<2" -tomli = {version = ">=1", markers = "python_version < \"3.11\""} [package.extras] dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"] @@ -2415,11 +2379,6 @@ files = [ {file = "scikit_learn-1.5.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f60021ec1574e56632be2a36b946f8143bf4e5e6af4a06d85281adc22938e0dd"}, {file = "scikit_learn-1.5.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:394397841449853c2290a32050382edaec3da89e35b3e03d6cc966aebc6a8ae6"}, {file = "scikit_learn-1.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:57cc1786cfd6bd118220a92ede80270132aa353647684efa385a74244a41e3b1"}, - {file = "scikit_learn-1.5.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e9a702e2de732bbb20d3bad29ebd77fc05a6b427dc49964300340e4c9328b3f5"}, - {file = "scikit_learn-1.5.2-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:b0768ad641981f5d3a198430a1d31c3e044ed2e8a6f22166b4d546a5116d7908"}, - {file = "scikit_learn-1.5.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:178ddd0a5cb0044464fc1bfc4cca5b1833bfc7bb022d70b05db8530da4bb3dd3"}, - {file = "scikit_learn-1.5.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f7284ade780084d94505632241bf78c44ab3b6f1e8ccab3d2af58e0e950f9c12"}, - {file = "scikit_learn-1.5.2-cp313-cp313-win_amd64.whl", hash = "sha256:b7b0f9a0b1040830d38c39b91b3a44e1b643f4b36e36567b80b7c6bd2202a27f"}, {file = "scikit_learn-1.5.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:757c7d514ddb00ae249832fe87100d9c73c6ea91423802872d9e74970a0e40b9"}, {file = "scikit_learn-1.5.2-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:52788f48b5d8bca5c0736c175fa6bdaab2ef00a8f536cda698db61bd89c551c1"}, {file = "scikit_learn-1.5.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:643964678f4b5fbdc95cbf8aec638acc7aa70f5f79ee2cdad1eec3df4ba6ead8"}, @@ -2610,28 +2569,6 @@ files = [ {file = "threadpoolctl-3.5.0.tar.gz", hash = "sha256:082433502dd922bf738de0d8bcc4fdcbf0979ff44c42bd40f5af8a282f6fa107"}, ] -[[package]] -name = "toml" -version = "0.10.2" -description = "Python Library for Tom's Obvious, Minimal Language" -optional = false -python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*" -files = [ - {file = "toml-0.10.2-py2.py3-none-any.whl", hash = "sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b"}, - {file = "toml-0.10.2.tar.gz", hash = "sha256:b3bda1d108d5dd99f4a20d24d9c348e91c4db7ab1b749200bded2f839ccbe68f"}, -] - -[[package]] -name = "tomli" -version = "2.0.1" -description = "A lil' TOML parser" -optional = false -python-versions = ">=3.7" -files = [ - {file = "tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc"}, - {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, -] - [[package]] name = "tqdm" version = "4.66.5" @@ -2652,17 +2589,6 @@ notebook = ["ipywidgets (>=6)"] slack = ["slack-sdk"] telegram = ["requests"] -[[package]] -name = "types-toml" -version = "0.10.8.20240310" -description = "Typing stubs for toml" -optional = false -python-versions = ">=3.8" -files = [ - {file = "types-toml-0.10.8.20240310.tar.gz", hash = "sha256:3d41501302972436a6b8b239c850b26689657e25281b48ff0ec06345b8830331"}, - {file = "types_toml-0.10.8.20240310-py3-none-any.whl", hash = "sha256:627b47775d25fa29977d9c70dc0cbab3f314f32c8d8d0c012f2ef5de7aaec05d"}, -] - [[package]] name = "typing-extensions" version = "4.12.2" @@ -2716,7 +2642,6 @@ files = [ [package.dependencies] click = ">=7.0" h11 = ">=0.8" -typing-extensions = {version = ">=4.0", markers = "python_version < \"3.11\""} [package.extras] standard = ["colorama (>=0.4)", "httptools (>=0.5.0)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.14.0,!=0.15.0,!=0.15.1)", "watchfiles (>=0.13)", "websockets (>=10.4)"] @@ -2970,5 +2895,5 @@ propcache = ">=0.2.0" [metadata] lock-version = "2.0" -python-versions = "^3.10" -content-hash = "1b1034745163165a2d3c528c8b1e65f7e0c1bfadd4d0bb1d195eb83a48d7474a" +python-versions = "^3.11" +content-hash = "5237f88b9526bbf893e73c1ca2be04a3b598e7aaf4c3d72e7cf7eaac9ff72ecd" diff --git a/pyproject.toml b/pyproject.toml index 6abf0c6..ebc6829 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "agentic_security" -version = "0.3.0" +version = "0.3.1" description = "Agentic LLM vulnerability scanner" authors = ["Alexander Miasoiedov "] maintainers = ["Alexander Miasoiedov "]