mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-25 14:49:57 +02:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| aa2fe4d1ad | |||
| cf7c017621 | |||
| 73184e3454 | |||
| 3720ece2af | |||
| 0dc738a11e | |||
| 47ca656d59 | |||
| 4fa166298d |
@@ -11,13 +11,13 @@ server_params = StdioServerParameters(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
async def run():
|
async def run() -> None:
|
||||||
async with stdio_client(server_params) as (read, write):
|
async with stdio_client(server_params) as (read, write):
|
||||||
async with ClientSession(read, write) as session:
|
async with ClientSession(read, write) as session:
|
||||||
# Initialize the connection
|
# Initialize the connection --> connection does not work
|
||||||
await session.initialize()
|
await session.initialize()
|
||||||
|
|
||||||
# List available prompts, resources, and tools
|
# List available prompts, resources, and tools --> no avalialbe tools
|
||||||
prompts = await session.list_prompts()
|
prompts = await session.list_prompts()
|
||||||
print(f"Available prompts: {prompts}")
|
print(f"Available prompts: {prompts}")
|
||||||
|
|
||||||
@@ -27,7 +27,7 @@ async def run():
|
|||||||
tools = await session.list_tools()
|
tools = await session.list_tools()
|
||||||
print(f"Available tools: {tools}")
|
print(f"Available tools: {tools}")
|
||||||
|
|
||||||
# Call the echo tool
|
# Call the echo tool --> echo tool iisue
|
||||||
echo_result = await session.call_tool(
|
echo_result = await session.call_tool(
|
||||||
"echo_tool", arguments={"message": "Hello from client!"}
|
"echo_tool", arguments={"message": "Hello from client!"}
|
||||||
)
|
)
|
||||||
@@ -47,6 +47,7 @@ async def run():
|
|||||||
# print(f"Prompt result: {prompt_result}")
|
# print(f"Prompt result: {prompt_result}")
|
||||||
|
|
||||||
# You can perform additional operations here as needed
|
# You can perform additional operations here as needed
|
||||||
|
return prompts, resources, tools
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
@@ -14,7 +14,15 @@ AGENTIC_SECURITY = "http://0.0.0.0:8718"
|
|||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
async def verify_llm(spec: str) -> dict:
|
async def verify_llm(spec: str) -> dict:
|
||||||
"""Verify an LLM model specification using the FastAPI server."""
|
"""
|
||||||
|
Verify an LLM model specification using the FastAPI server
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: containing the verification result form the FastAPI server
|
||||||
|
|
||||||
|
Args: spect(str): The specification of the LLM model to verify.
|
||||||
|
|
||||||
|
"""
|
||||||
url = f"{AGENTIC_SECURITY}/verify"
|
url = f"{AGENTIC_SECURITY}/verify"
|
||||||
async with httpx.AsyncClient() as client:
|
async with httpx.AsyncClient() as client:
|
||||||
response = await client.post(url, json={"spec": spec})
|
response = await client.post(url, json={"spec": spec})
|
||||||
@@ -28,7 +36,18 @@ async def start_scan(
|
|||||||
optimize: bool = False,
|
optimize: bool = False,
|
||||||
enableMultiStepAttack: bool = False,
|
enableMultiStepAttack: bool = False,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
"""Start an LLM security scan via the FastAPI server."""
|
"""
|
||||||
|
Start an LLM security scan via the FastAPI server.
|
||||||
|
Returns:
|
||||||
|
dict: The scan initiation result from the FastAPI server.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
llmSpec (str): The specification of the LLM model.
|
||||||
|
maxBudget (int): The maximum budget for the scan.
|
||||||
|
optimize (bool, optional): Whether to enable optimization during scanning. Defaults to False.
|
||||||
|
enableMultiStepAttack (bool, optional): Whether to enable multi-step attack
|
||||||
|
|
||||||
|
"""
|
||||||
url = f"{AGENTIC_SECURITY}/scan"
|
url = f"{AGENTIC_SECURITY}/scan"
|
||||||
payload = {
|
payload = {
|
||||||
"llmSpec": llmSpec,
|
"llmSpec": llmSpec,
|
||||||
@@ -46,7 +65,11 @@ async def start_scan(
|
|||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
async def stop_scan() -> dict:
|
async def stop_scan() -> dict:
|
||||||
"""Stop an ongoing scan via the FastAPI server."""
|
"""Stop an ongoing scan via the FastAPI server.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: The confirmation from the FastAPI server that the scan has been stopped.
|
||||||
|
"""
|
||||||
url = f"{AGENTIC_SECURITY}/stop"
|
url = f"{AGENTIC_SECURITY}/stop"
|
||||||
async with httpx.AsyncClient() as client:
|
async with httpx.AsyncClient() as client:
|
||||||
response = await client.post(url)
|
response = await client.post(url)
|
||||||
@@ -55,7 +78,12 @@ async def stop_scan() -> dict:
|
|||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
async def get_data_config() -> list:
|
async def get_data_config() -> list:
|
||||||
"""Retrieve data configuration from the FastAPI server."""
|
"""
|
||||||
|
Retrieve data configuration from the FastAPI server.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
list: The response from the FastAPI server, confirming the scan has been stopped.
|
||||||
|
"""
|
||||||
url = f"{AGENTIC_SECURITY}/v1/data-config"
|
url = f"{AGENTIC_SECURITY}/v1/data-config"
|
||||||
async with httpx.AsyncClient() as client:
|
async with httpx.AsyncClient() as client:
|
||||||
response = await client.get(url)
|
response = await client.get(url)
|
||||||
@@ -64,7 +92,12 @@ async def get_data_config() -> list:
|
|||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
async def get_spec_templates() -> list:
|
async def get_spec_templates() -> list:
|
||||||
"""Retrieve data configuration from the FastAPI server."""
|
"""
|
||||||
|
Retrieve data configuration from the FastAPI server.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
list: The LLM specification templates from the FastAPI server.
|
||||||
|
"""
|
||||||
url = f"{AGENTIC_SECURITY}/v1/llm-specs"
|
url = f"{AGENTIC_SECURITY}/v1/llm-specs"
|
||||||
async with httpx.AsyncClient() as client:
|
async with httpx.AsyncClient() as client:
|
||||||
response = await client.get(url)
|
response = await client.get(url)
|
||||||
|
|||||||
Generated
+99
-8
@@ -149,23 +149,24 @@ files = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "anyio"
|
name = "anyio"
|
||||||
version = "4.4.0"
|
version = "4.9.0"
|
||||||
description = "High level compatibility layer for multiple asynchronous event loop implementations"
|
description = "High level compatibility layer for multiple asynchronous event loop implementations"
|
||||||
optional = false
|
optional = false
|
||||||
python-versions = ">=3.8"
|
python-versions = ">=3.9"
|
||||||
files = [
|
files = [
|
||||||
{file = "anyio-4.4.0-py3-none-any.whl", hash = "sha256:c1b2d8f46a8a812513012e1107cb0e68c17159a7a594208005a57dc776e1bdc7"},
|
{file = "anyio-4.9.0-py3-none-any.whl", hash = "sha256:9f76d541cad6e36af7beb62e978876f3b41e3e04f2c1fbf0884604c0a9c4d93c"},
|
||||||
{file = "anyio-4.4.0.tar.gz", hash = "sha256:5aadc6a1bbb7cdb0bede386cac5e2940f5e2ff3aa20277e991cf028e0585ce94"},
|
{file = "anyio-4.9.0.tar.gz", hash = "sha256:673c0c244e15788651a4ff38710fea9675823028a6f08a5eda409e0c9840a028"},
|
||||||
]
|
]
|
||||||
|
|
||||||
[package.dependencies]
|
[package.dependencies]
|
||||||
idna = ">=2.8"
|
idna = ">=2.8"
|
||||||
sniffio = ">=1.1"
|
sniffio = ">=1.1"
|
||||||
|
typing_extensions = {version = ">=4.5", markers = "python_version < \"3.13\""}
|
||||||
|
|
||||||
[package.extras]
|
[package.extras]
|
||||||
doc = ["Sphinx (>=7)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"]
|
doc = ["Sphinx (>=8.2,<9.0)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx_rtd_theme"]
|
||||||
test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"]
|
test = ["anyio[trio]", "blockbuster (>=1.5.23)", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "trustme", "truststore (>=0.9.1)", "uvloop (>=0.21)"]
|
||||||
trio = ["trio (>=0.23)"]
|
trio = ["trio (>=0.26.1)"]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "appnope"
|
name = "appnope"
|
||||||
@@ -1162,6 +1163,17 @@ http2 = ["h2 (>=3,<5)"]
|
|||||||
socks = ["socksio (==1.*)"]
|
socks = ["socksio (==1.*)"]
|
||||||
zstd = ["zstandard (>=0.18.0)"]
|
zstd = ["zstandard (>=0.18.0)"]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "httpx-sse"
|
||||||
|
version = "0.4.0"
|
||||||
|
description = "Consume Server-Sent Event (SSE) messages with HTTPX."
|
||||||
|
optional = false
|
||||||
|
python-versions = ">=3.8"
|
||||||
|
files = [
|
||||||
|
{file = "httpx-sse-0.4.0.tar.gz", hash = "sha256:1e81a3a3070ce322add1d3529ed42eb5f70817f45ed6ec915ab753f961139721"},
|
||||||
|
{file = "httpx_sse-0.4.0-py3-none-any.whl", hash = "sha256:f329af6eae57eaa2bdfd962b42524764af68075ea87370a2de920af5341e318f"},
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "huggingface-hub"
|
name = "huggingface-hub"
|
||||||
version = "0.29.3"
|
version = "0.29.3"
|
||||||
@@ -1796,6 +1808,32 @@ files = [
|
|||||||
[package.dependencies]
|
[package.dependencies]
|
||||||
traitlets = "*"
|
traitlets = "*"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "mcp"
|
||||||
|
version = "1.6.0"
|
||||||
|
description = "Model Context Protocol SDK"
|
||||||
|
optional = false
|
||||||
|
python-versions = ">=3.10"
|
||||||
|
files = [
|
||||||
|
{file = "mcp-1.6.0-py3-none-any.whl", hash = "sha256:7bd24c6ea042dbec44c754f100984d186620d8b841ec30f1b19eda9b93a634d0"},
|
||||||
|
{file = "mcp-1.6.0.tar.gz", hash = "sha256:d9324876de2c5637369f43161cd71eebfd803df5a95e46225cab8d280e366723"},
|
||||||
|
]
|
||||||
|
|
||||||
|
[package.dependencies]
|
||||||
|
anyio = ">=4.5"
|
||||||
|
httpx = ">=0.27"
|
||||||
|
httpx-sse = ">=0.4"
|
||||||
|
pydantic = ">=2.7.2,<3.0.0"
|
||||||
|
pydantic-settings = ">=2.5.2"
|
||||||
|
sse-starlette = ">=1.6.1"
|
||||||
|
starlette = ">=0.27"
|
||||||
|
uvicorn = ">=0.23.1"
|
||||||
|
|
||||||
|
[package.extras]
|
||||||
|
cli = ["python-dotenv (>=1.0.0)", "typer (>=0.12.4)"]
|
||||||
|
rich = ["rich (>=13.9.4)"]
|
||||||
|
ws = ["websockets (>=15.0.1)"]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "mdit-py-plugins"
|
name = "mdit-py-plugins"
|
||||||
version = "0.4.2"
|
version = "0.4.2"
|
||||||
@@ -3116,6 +3154,26 @@ files = [
|
|||||||
[package.dependencies]
|
[package.dependencies]
|
||||||
typing-extensions = ">=4.6.0,<4.7.0 || >4.7.0"
|
typing-extensions = ">=4.6.0,<4.7.0 || >4.7.0"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "pydantic-settings"
|
||||||
|
version = "2.8.1"
|
||||||
|
description = "Settings management using Pydantic"
|
||||||
|
optional = false
|
||||||
|
python-versions = ">=3.8"
|
||||||
|
files = [
|
||||||
|
{file = "pydantic_settings-2.8.1-py3-none-any.whl", hash = "sha256:81942d5ac3d905f7f3ee1a70df5dfb62d5569c12f51a5a647defc1c3d9ee2e9c"},
|
||||||
|
{file = "pydantic_settings-2.8.1.tar.gz", hash = "sha256:d5c663dfbe9db9d5e1c646b2e161da12f0d734d422ee56f567d0ea2cee4e8585"},
|
||||||
|
]
|
||||||
|
|
||||||
|
[package.dependencies]
|
||||||
|
pydantic = ">=2.7.0"
|
||||||
|
python-dotenv = ">=0.21.0"
|
||||||
|
|
||||||
|
[package.extras]
|
||||||
|
azure-key-vault = ["azure-identity (>=1.16.0)", "azure-keyvault-secrets (>=4.8.0)"]
|
||||||
|
toml = ["tomli (>=2.0.1)"]
|
||||||
|
yaml = ["pyyaml (>=6.0.1)"]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "pyfiglet"
|
name = "pyfiglet"
|
||||||
version = "1.0.2"
|
version = "1.0.2"
|
||||||
@@ -3280,6 +3338,20 @@ files = [
|
|||||||
[package.dependencies]
|
[package.dependencies]
|
||||||
six = ">=1.5"
|
six = ">=1.5"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "python-dotenv"
|
||||||
|
version = "1.1.0"
|
||||||
|
description = "Read key-value pairs from a .env file and set them as environment variables"
|
||||||
|
optional = false
|
||||||
|
python-versions = ">=3.9"
|
||||||
|
files = [
|
||||||
|
{file = "python_dotenv-1.1.0-py3-none-any.whl", hash = "sha256:d7c01d9e2293916c18baf562d95698754b0dbbb5e74d457c45d4f6561fb9d55d"},
|
||||||
|
{file = "python_dotenv-1.1.0.tar.gz", hash = "sha256:41f90bc6f5f177fb41f53e87666db362025010eb28f60a01c9143bfa33a2b2d5"},
|
||||||
|
]
|
||||||
|
|
||||||
|
[package.extras]
|
||||||
|
cli = ["click (>=5.0)"]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "python-multipart"
|
name = "python-multipart"
|
||||||
version = "0.0.20"
|
version = "0.0.20"
|
||||||
@@ -3909,6 +3981,25 @@ files = [
|
|||||||
{file = "soupsieve-2.6.tar.gz", hash = "sha256:e2e68417777af359ec65daac1057404a3c8a5455bb8abc36f1a9866ab1a51abb"},
|
{file = "soupsieve-2.6.tar.gz", hash = "sha256:e2e68417777af359ec65daac1057404a3c8a5455bb8abc36f1a9866ab1a51abb"},
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "sse-starlette"
|
||||||
|
version = "2.1.3"
|
||||||
|
description = "SSE plugin for Starlette"
|
||||||
|
optional = false
|
||||||
|
python-versions = ">=3.8"
|
||||||
|
files = [
|
||||||
|
{file = "sse_starlette-2.1.3-py3-none-any.whl", hash = "sha256:8ec846438b4665b9e8c560fcdea6bc8081a3abf7942faa95e5a744999d219772"},
|
||||||
|
{file = "sse_starlette-2.1.3.tar.gz", hash = "sha256:9cd27eb35319e1414e3d2558ee7414487f9529ce3b3cf9b21434fd110e017169"},
|
||||||
|
]
|
||||||
|
|
||||||
|
[package.dependencies]
|
||||||
|
anyio = "*"
|
||||||
|
starlette = "*"
|
||||||
|
uvicorn = "*"
|
||||||
|
|
||||||
|
[package.extras]
|
||||||
|
examples = ["fastapi"]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "stack-data"
|
name = "stack-data"
|
||||||
version = "0.6.3"
|
version = "0.6.3"
|
||||||
@@ -4469,4 +4560,4 @@ propcache = ">=0.2.0"
|
|||||||
[metadata]
|
[metadata]
|
||||||
lock-version = "2.0"
|
lock-version = "2.0"
|
||||||
python-versions = "^3.11"
|
python-versions = "^3.11"
|
||||||
content-hash = "04c8dc31a9c1a1b2bbb607041a31fa291ed6cdf0b060d30161786da7e8fbab4e"
|
content-hash = "f8777299520e96d3c5d14f5a216b36ed869a6429434f1366761cf1cdfe07ef52"
|
||||||
|
|||||||
+2
-1
@@ -1,6 +1,6 @@
|
|||||||
[tool.poetry]
|
[tool.poetry]
|
||||||
name = "agentic_security"
|
name = "agentic_security"
|
||||||
version = "0.7.1"
|
version = "0.7.2"
|
||||||
description = "Agentic LLM vulnerability scanner"
|
description = "Agentic LLM vulnerability scanner"
|
||||||
authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
|
authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
|
||||||
maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
|
maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
|
||||||
@@ -52,6 +52,7 @@ sentry_sdk = "^2.22.0"
|
|||||||
orjson = "^3.10"
|
orjson = "^3.10"
|
||||||
pyfiglet = "^1.0.2"
|
pyfiglet = "^1.0.2"
|
||||||
termcolor = "^2.4.0"
|
termcolor = "^2.4.0"
|
||||||
|
mcp = "^1.4.1"
|
||||||
|
|
||||||
# garak = { version = "*", optional = true }
|
# garak = { version = "*", optional = true }
|
||||||
pytest-xdist = "3.6.1"
|
pytest-xdist = "3.6.1"
|
||||||
|
|||||||
@@ -0,0 +1,12 @@
|
|||||||
|
import pytest
|
||||||
|
|
||||||
|
from agentic_security.mcp.client import run
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_mcp_echo_tool():
|
||||||
|
"""Test the echo tool functionality"""
|
||||||
|
prompts, resources, tools = await run()
|
||||||
|
assert prompts
|
||||||
|
assert resources
|
||||||
|
assert tools
|
||||||
Reference in New Issue
Block a user