humanizetech commited on
Commit
e671617
·
1 Parent(s): 8726684

feat: Add VAPT agent client for MCP security tests

Browse files
Files changed (2) hide show
  1. .gitignore +207 -0
  2. vapt_mcp_client.py +133 -0
.gitignore ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py.cover
50
+ .hypothesis/
51
+ .pytest_cache/
52
+ cover/
53
+
54
+ # Translations
55
+ *.mo
56
+ *.pot
57
+
58
+ # Django stuff:
59
+ *.log
60
+ local_settings.py
61
+ db.sqlite3
62
+ db.sqlite3-journal
63
+
64
+ # Flask stuff:
65
+ instance/
66
+ .webassets-cache
67
+
68
+ # Scrapy stuff:
69
+ .scrapy
70
+
71
+ # Sphinx documentation
72
+ docs/_build/
73
+
74
+ # PyBuilder
75
+ .pybuilder/
76
+ target/
77
+
78
+ # Jupyter Notebook
79
+ .ipynb_checkpoints
80
+
81
+ # IPython
82
+ profile_default/
83
+ ipython_config.py
84
+
85
+ # pyenv
86
+ # For a library or package, you might want to ignore these files since the code is
87
+ # intended to run in multiple environments; otherwise, check them in:
88
+ # .python-version
89
+
90
+ # pipenv
91
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
93
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
94
+ # install all needed dependencies.
95
+ #Pipfile.lock
96
+
97
+ # UV
98
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
99
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
100
+ # commonly ignored for libraries.
101
+ #uv.lock
102
+
103
+ # poetry
104
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
105
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
106
+ # commonly ignored for libraries.
107
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
108
+ #poetry.lock
109
+ #poetry.toml
110
+
111
+ # pdm
112
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
113
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
114
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
115
+ #pdm.lock
116
+ #pdm.toml
117
+ .pdm-python
118
+ .pdm-build/
119
+
120
+ # pixi
121
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
122
+ #pixi.lock
123
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
124
+ # in the .venv directory. It is recommended not to include this directory in version control.
125
+ .pixi
126
+
127
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
128
+ __pypackages__/
129
+
130
+ # Celery stuff
131
+ celerybeat-schedule
132
+ celerybeat.pid
133
+
134
+ # SageMath parsed files
135
+ *.sage.py
136
+
137
+ # Environments
138
+ .env
139
+ .envrc
140
+ .venv
141
+ env/
142
+ venv/
143
+ ENV/
144
+ env.bak/
145
+ venv.bak/
146
+
147
+ # Spyder project settings
148
+ .spyderproject
149
+ .spyproject
150
+
151
+ # Rope project settings
152
+ .ropeproject
153
+
154
+ # mkdocs documentation
155
+ /site
156
+
157
+ # mypy
158
+ .mypy_cache/
159
+ .dmypy.json
160
+ dmypy.json
161
+
162
+ # Pyre type checker
163
+ .pyre/
164
+
165
+ # pytype static type analyzer
166
+ .pytype/
167
+
168
+ # Cython debug symbols
169
+ cython_debug/
170
+
171
+ # PyCharm
172
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
173
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
174
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
175
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
176
+ #.idea/
177
+
178
+ # Abstra
179
+ # Abstra is an AI-powered process automation framework.
180
+ # Ignore directories containing user credentials, local state, and settings.
181
+ # Learn more at https://abstra.io/docs
182
+ .abstra/
183
+
184
+ # Visual Studio Code
185
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
186
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
187
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
188
+ # you could uncomment the following to ignore the entire vscode folder
189
+ # .vscode/
190
+
191
+ # Ruff stuff:
192
+ .ruff_cache/
193
+
194
+ # PyPI configuration file
195
+ .pypirc
196
+
197
+ # Cursor
198
+ # Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
199
+ # exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
200
+ # refer to https://docs.cursor.com/context/ignore-files
201
+ .cursorignore
202
+ .cursorindexingignore
203
+
204
+ # Marimo
205
+ marimo/_static/
206
+ marimo/_lsp/
207
+ __marimo__/
vapt_mcp_client.py ADDED
@@ -0,0 +1,133 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import asyncio
2
+ import ast
3
+ import os
4
+ from pathlib import Path
5
+ from typing import Dict, Any
6
+
7
+ from dotenv import load_dotenv
8
+ from mcp import ClientSession
9
+ from mcp.client.streamable_http import streamablehttp_client
10
+
11
+ # Load .env automatically
12
+ load_dotenv()
13
+
14
+ MCP_SERVER_URL = "https://mcp-1st-birthday-vapt-agent.hf.space/gradio_api/mcp/"
15
+ TOOL_NAME = "vapt_agent_run_security_test"
16
+
17
+ BASE_DIR = Path(__file__).parent
18
+
19
+
20
+ async def call_vapt_security_test(
21
+ api_endpoint: str,
22
+ http_method: str = "GET",
23
+ api_key: str = "",
24
+ ) -> Dict[str, Any]:
25
+
26
+ report_md_text = None
27
+ report_file_name = None
28
+
29
+ async with streamablehttp_client(url=MCP_SERVER_URL, headers={}) as (
30
+ read_stream,
31
+ write_stream,
32
+ _,
33
+ ):
34
+ async with ClientSession(read_stream, write_stream) as session:
35
+ await session.initialize()
36
+
37
+ tools_resp = await session.list_tools()
38
+ tool_names = [t.name for t in tools_resp.tools]
39
+ if TOOL_NAME not in tool_names:
40
+ raise RuntimeError(
41
+ f"Tool {TOOL_NAME!r} not found on server. Available: {tool_names}"
42
+ )
43
+
44
+ args = {
45
+ "api_endpoint": api_endpoint,
46
+ "http_method": http_method,
47
+ "api_key": api_key,
48
+ }
49
+
50
+ result = await session.call_tool(TOOL_NAME, args)
51
+
52
+ def handle_text_chunk(text: str):
53
+ nonlocal report_md_text, report_file_name
54
+
55
+ parsed = None
56
+ try:
57
+ parsed = ast.literal_eval(text)
58
+ except Exception:
59
+ parsed = None
60
+
61
+ # Expected: ['progress...', 'report_md', file_info_dict]
62
+ if (
63
+ isinstance(parsed, list)
64
+ and len(parsed) >= 3
65
+ and isinstance(parsed[0], str)
66
+ and isinstance(parsed[1], str)
67
+ and isinstance(parsed[2], dict)
68
+ ):
69
+ progress_str, md_str, file_info = parsed[0], parsed[1], parsed[2]
70
+
71
+ print("\n=== Progress ===")
72
+ for line in progress_str.splitlines():
73
+ print(line)
74
+
75
+ print("\n=== Report (first 20 lines) ===")
76
+ md_lines = md_str.splitlines()
77
+ for line in md_lines[:20]:
78
+ print(line)
79
+ if len(md_lines) > 20:
80
+ print("... [truncated, full report saved to .md file]")
81
+
82
+ report_md_text = md_str
83
+ report_file_name = file_info.get("orig_name", "vapt_report.md")
84
+
85
+ else:
86
+ print(text)
87
+
88
+ # Future compatibility: streaming async iterator
89
+ if hasattr(result, "__aiter__"):
90
+ async for event in result:
91
+ for block in event.content:
92
+ if getattr(block, "type", None) == "text":
93
+ handle_text_chunk(block.text)
94
+ else:
95
+ print("\n=== RAW TOOL RESULT METADATA ===")
96
+ print(f"isError: {result.isError}")
97
+
98
+ for block in result.content:
99
+ if getattr(block, "type", None) == "text":
100
+ handle_text_chunk(block.text)
101
+
102
+ # Save markdown
103
+ if report_md_text:
104
+ out_name = report_file_name or "vapt_report.md"
105
+ out_path = BASE_DIR / out_name
106
+ out_path.write_text(report_md_text, encoding="utf-8")
107
+ print(f"\n✅ Markdown report saved to: {out_path.resolve()}")
108
+
109
+ return {
110
+ "is_error": getattr(result, "isError", None),
111
+ "content": getattr(result, "content", None),
112
+ }
113
+
114
+
115
+ async def main():
116
+ api_endpoint = os.getenv("TEST_API_ENDPOINT")
117
+ http_method = os.getenv("TEST_API_METHOD", "GET")
118
+ api_key = os.getenv("TEST_API_KEY", "")
119
+
120
+ print("\nUsing environment settings:")
121
+ print(f" TEST_API_ENDPOINT = {api_endpoint}")
122
+ print(f" TEST_API_METHOD = {http_method}")
123
+ # Do NOT print API key
124
+
125
+ await call_vapt_security_test(
126
+ api_endpoint=api_endpoint,
127
+ http_method=http_method,
128
+ api_key=api_key,
129
+ )
130
+
131
+
132
+ if __name__ == "__main__":
133
+ asyncio.run(main())