A.I.G(AI-Infra-Guard) provides a comprehensive set of API interfaces for AI Infra Scan, MCP Server Scan, and Jailbreak Evaluation. This documentation details the usage methods, parameter descriptions, and example code for each API interface.
After the project is running, you can access http://localhost:8088/docs/index.html to view the Swagger documentation.
- Base URL:
http://localhost:8088(adjust according to actual deployment) - Content-Type:
application/json - Authentication: Pass authentication information through request headers
All API interfaces follow a unified response format:
{
"status": 0, // Status code: 0=success, 1=failure
"message": "Operation successful", // Response message
"data": {} // Response data
}- URL:
/api/v1/app/taskapi/upload - Method:
POST - Content-Type:
multipart/form-data
| Parameter | Type | Required | Description |
|---|---|---|---|
| file | file | Yes | File to upload, supports zip, json, txt and other formats |
| Field | Type | Description |
|---|---|---|
| fileUrl | string | File access URL |
| filename | string | File name |
| size | integer | File size (bytes) |
import requests
def upload_file(file_path):
url = "http://localhost:8088/api/v1/app/taskapi/upload"
with open(file_path, 'rb') as f:
files = {'file': f}
response = requests.post(url, files=files)
return response.json()
# Usage example
result = upload_file("example.zip")
print(f"File uploaded successfully: {result['data']['fileUrl']}")curl -X POST \
http://localhost:8088/api/v1/app/taskapi/upload \
-F "file=@example.zip"- URL:
/api/v1/app/taskapi/tasks - Method:
POST - Content-Type:
application/json
| Parameter | Type | Required | Description |
|---|---|---|---|
| type | string | Yes | Task type: mcp_scan, ai_infra_scan, model_redteam_report |
| content | object | Yes | Task content, varies according to task type |
| Field | Type | Description |
|---|---|---|
| session_id | string | Task session ID |
MCP Server Scan is used to detect security vulnerabilities in MCP servers.
| Parameter | Type | Required | Description |
|---|---|---|---|
| content | string | No | Task content description |
| model | object | Yes | Model configuration |
| model.model | string | Yes | Model name, e.g., "gpt-4" |
| model.token | string | Yes | API key |
| model.base_url | string | No | Base URL, defaults to OpenAI API |
| thread | integer | No | Concurrent thread count, default 4 |
| language | string | No | Language code, e.g., "zh" |
| attachments | string | No | Attachment file path (file must be uploaded first) |
- First call the file upload interface to upload source code files
- Use the returned fileUrl as the attachments parameter
- Call the MCP Server Scan API
import requests
import json
def mcp_scan_with_source_code():
# 1. Upload source code file
upload_url = "http://localhost:8088/api/v1/app/taskapi/upload"
with open("source_code.zip", 'rb') as f:
files = {'file': f}
upload_response = requests.post(upload_url, files=files)
if upload_response.json()['status'] != 0:
raise Exception("File upload failed")
fileUrl = upload_response.json()['data']['fileUrl']
# 2. Create MCP Server Scan task
task_url = "http://localhost:8088/api/v1/app/taskapi/tasks"
task_data = {
"type": "mcp_scan",
"content": {
"content": "",
"model": {
"model": "gpt-4",
"token": "sk-your-api-key",
"base_url": "https://api.openai.com/v1"
},
"thread": 4,
"language": "zh",
"attachments": fileUrl
}
}
response = requests.post(task_url, json=task_data)
return response.json()
# Usage example
result = mcp_scan_with_source_code()
print(f"Task created successfully, session ID: {result['data']['session_id']}")def mcp_scan_with_url():
task_url = "http://localhost:8088/api/v1/app/taskapi/tasks"
task_data = {
"type": "mcp_scan",
"content": {
"content": "https://mcp-server.example.com", # Direct URL input
"model": {
"model": "gpt-4",
"token": "sk-your-api-key",
"base_url": "https://api.openai.com/v1"
},
"thread": 4,
"language": "zh"
}
}
response = requests.post(task_url, json=task_data)
return response.json()# Source code scanning
curl -X POST http://localhost:8088/api/v1/app/taskapi/tasks \
-H "Content-Type: application/json" \
-d '{
"type": "mcp_scan",
"content": {
"content": "",
"model": {
"model": "gpt-4",
"token": "sk-your-api-key",
"base_url": "https://api.openai.com/v1"
},
"thread": 4,
"language": "zh",
"attachments": "http://localhost:8088/uploads/example.zip"
}
}'
# URL scanning
curl -X POST http://localhost:8088/api/v1/app/taskapi/tasks \
-H "Content-Type: application/json" \
-d '{
"type": "mcp_scan",
"content": {
"content": "https://mcp-server.example.com",
"model": {
"model": "gpt-4",
"token": "sk-your-api-key",
"base_url": "https://api.openai.com/v1"
},
"thread": 4,
"language": "zh"
}
}'Used to scan AI infra for security vulnerabilities and configuration issues.
| Parameter | Type | Required | Description |
|---|---|---|---|
| target | array | Yes | List of target URLs to scan |
| headers | object | No | Custom request headers |
| timeout | integer | No | Request timeout (seconds), default 30 |
def ai_infra_scan():
task_url = "http://localhost:8088/api/v1/app/taskapi/tasks"
task_data = {
"type": "ai_infra_scan",
"content": {
"target": [
"https://ai-service1.example.com",
"https://ai-service2.example.com"
],
"headers": {
"Authorization": "Bearer your-token",
"User-Agent": "AI-Infra-Guard/1.0"
},
"timeout": 30
}
}
response = requests.post(task_url, json=task_data)
return response.json()
# Usage example
result = ai_infra_scan()
print(f"AI infra scan task created successfully, session ID: {result['data']['session_id']}")curl -X POST http://localhost:8088/api/v1/app/taskapi/tasks \
-H "Content-Type: application/json" \
-d '{
"type": "ai_infra_scan",
"content": {
"target": [
"https://ai-service1.example.com",
"https://ai-service2.example.com"
],
"headers": {
"Authorization": "Bearer your-token",
"User-Agent": "AI-Infra-Guard/1.0"
},
"timeout": 30
}
}'Used to perform Jailbreak Evaluation testing on LLM to assess their security and robustness.
| Parameter | Type | Required | Description |
|---|---|---|---|
| model | array | Yes | List of models to test |
| eval_model | object | Yes | Evaluation model configuration |
| dataset | object | Yes | Dataset configuration |
| dataset.dataFile | array | Yes | List of dataset files, supports the following options: - JailBench-Tiny: Small jailbreak benchmark test dataset - JailbreakPrompts-Tiny: Small jailbreak prompt dataset - ChatGPT-Jailbreak-Prompts: ChatGPT jailbreak prompt dataset - JADE-db-v3.0: JADE database v3.0 version - HarmfulEvalBenchmark: Harmful content evaluation benchmark dataset |
| dataset.numPrompts | integer | Yes | Number of prompts |
| dataset.randomSeed | integer | Yes | Random seed |
| Dataset Name | Description | Use Case |
|---|---|---|
| JailBench-Tiny | Small jailbreak benchmark test dataset | Quick testing of model resistance to jailbreak attacks |
| JailbreakPrompts-Tiny | Small jailbreak prompt dataset | Testing model protection against common jailbreak techniques |
| ChatGPT-Jailbreak-Prompts | ChatGPT jailbreak prompt dataset | Jailbreak testing specifically targeting ChatGPT |
| JADE-db-v3.0 | JADE database v3.0 version | Comprehensive AI security evaluation dataset |
| HarmfulEvalBenchmark | Harmful content evaluation benchmark dataset | Assessing risks of model-generated harmful content |
def model_redteam_test():
task_url = "http://localhost:8088/api/v1/app/taskapi/tasks"
task_data = {
"type": "model_redteam_report",
"content": {
"model": [
{
"model": "gpt-4",
"token": "sk-your-api-key",
"base_url": "https://api.openai.com/v1"
},
{
"model": "claude-3",
"token": "your-anthropic-key",
"base_url": "https://api.anthropic.com/v1"
}
],
"eval_model": {
"model": "gpt-4",
"token": "sk-your-eval-key",
"base_url": "https://api.openai.com/v1"
},
"dataset": {
"dataFile": [
"JailBench-Tiny",
"JailbreakPrompts-Tiny",
"ChatGPT-Jailbreak-Prompts"
],
"numPrompts": 100,
"randomSeed": 42
}
}
}
response = requests.post(task_url, json=task_data)
return response.json()
# Usage example
result = model_redteam_test()
print(f"Jailbreak Evaluation task created successfully, session ID: {result['data']['session_id']}")# Using JADE database for comprehensive testing
def comprehensive_redteam_test():
task_data = {
"type": "model_redteam_report",
"content": {
"model": [{"model": "gpt-4", "token": "sk-your-key"}],
"eval_model": {"model": "gpt-4", "token": "sk-eval-key"},
"dataset": {
"dataFile": ["JADE-db-v3.0"],
"numPrompts": 500,
"randomSeed": 123
}
}
}
return requests.post(task_url, json=task_data).json()
# Using harmful content evaluation benchmark
def harmful_content_test():
task_data = {
"type": "model_redteam_report",
"content": {
"model": [{"model": "gpt-4", "token": "sk-your-key"}],
"eval_model": {"model": "gpt-4", "token": "sk-eval-key"},
"dataset": {
"dataFile": ["HarmfulEvalBenchmark"],
"numPrompts": 200,
"randomSeed": 456
}
}
}
return requests.post(task_url, json=task_data).json()# Basic red team testing
curl -X POST http://localhost:8088/api/v1/app/taskapi/tasks \
-H "Content-Type: application/json" \
-d '{
"type": "model_redteam_report",
"content": {
"model": [
{
"model": "gpt-4",
"token": "sk-your-api-key",
"base_url": "https://api.openai.com/v1"
}
],
"eval_model": {
"model": "gpt-4",
"token": "sk-your-eval-key",
"base_url": "https://api.openai.com/v1"
},
"dataset": {
"dataFile": ["JailBench-Tiny", "JailbreakPrompts-Tiny"],
"numPrompts": 100,
"randomSeed": 42
}
}
}'
# Comprehensive security evaluation
curl -X POST http://localhost:8088/api/v1/app/taskapi/tasks \
-H "Content-Type: application/json" \
-d '{
"type": "model_redteam_report",
"content": {
"model": [{"model": "gpt-4", "token": "sk-your-key"}],
"eval_model": {"model": "gpt-4", "token": "sk-eval-key"},
"dataset": {
"dataFile": ["JADE-db-v3.0", "HarmfulEvalBenchmark"],
"numPrompts": 500,
"randomSeed": 123
}
}
}'- URL:
/api/v1/app/taskapi/status/{id} - Method:
GET
| Parameter | Type | Required | Description |
|---|---|---|---|
| id | string | Yes | Task session ID |
| Field | Type | Description |
|---|---|---|
| session_id | string | Task session ID |
| status | string | Task status: pending, running, completed, failed |
| title | string | Task title |
| created_at | integer | Creation timestamp (milliseconds) |
| updated_at | integer | Update timestamp (milliseconds) |
| log | string | Task execution log |
def get_task_status(session_id):
url = f"http://localhost:8088/api/v1/app/taskapi/status/{session_id}"
response = requests.get(url)
return response.json()
# Usage example
status = get_task_status("550e8400-e29b-41d4-a716-446655440000")
print(f"Task status: {status['data']['status']}")
print(f"Execution log: {status['data']['log']}")curl -X GET http://localhost:8088/api/v1/app/taskapi/status/550e8400-e29b-41d4-a716-446655440000- URL:
/api/v1/app/taskapi/result/{id} - Method:
GET
| Parameter | Type | Required | Description |
|---|---|---|---|
| id | string | Yes | Task session ID |
Returns detailed scan results, including:
- List of discovered vulnerabilities
- Security assessment report
- Remediation recommendations
- Risk level assessment
def get_task_result(session_id):
url = f"http://localhost:8088/api/v1/app/taskapi/result/{session_id}"
response = requests.get(url)
return response.json()
# Usage example
result = get_task_result("550e8400-e29b-41d4-a716-446655440000")
if result['status'] == 0:
print("Scan results:")
print(json.dumps(result['data'], indent=2, ensure_ascii=False))
else:
print(f"Failed to get results: {result['message']}")curl -X GET http://localhost:8088/api/v1/app/taskapi/result/550e8400-e29b-41d4-a716-446655440000import requests
import time
import json
def complete_mcp_scan_workflow():
base_url = "http://localhost:8088"
# 1. Upload source code file
print("1. Uploading source code file...")
upload_url = f"{base_url}/api/v1/app/taskapi/upload"
with open("mcp_source.zip", 'rb') as f:
files = {'file': f}
upload_response = requests.post(upload_url, files=files)
if upload_response.json()['status'] != 0:
raise Exception("File upload failed")
fileUrl = upload_response.json()['data']['fileUrl']
print(f"File uploaded successfully: {fileUrl}")
# 2. Create MCP scan task
print("2. Creating MCP scan task...")
task_url = f"{base_url}/api/v1/app/taskapi/tasks"
task_data = {
"type": "mcp_scan",
"content": {
"content": "",
"model": {
"model": "gpt-4",
"token": "sk-your-api-key",
"base_url": "https://api.openai.com/v1"
},
"thread": 4,
"language": "zh",
"attachments": fileUrl
}
}
task_response = requests.post(task_url, json=task_data)
if task_response.json()['status'] != 0:
raise Exception("Task creation failed")
session_id = task_response.json()['data']['session_id']
print(f"Task created successfully, session ID: {session_id}")
# 3. Poll task status
print("3. Monitoring task execution...")
status_url = f"{base_url}/api/v1/app/taskapi/status/{session_id}"
while True:
status_response = requests.get(status_url)
status_data = status_response.json()
if status_data['status'] != 0:
raise Exception("Failed to get task status")
task_status = status_data['data']['status']
print(f"Current status: {task_status}")
if task_status == "completed":
print("Task execution completed!")
break
elif task_status == "failed":
raise Exception("Task execution failed")
time.sleep(10) # Wait 10 seconds before checking again
# 4. Get scan results
print("4. Getting scan results...")
result_url = f"{base_url}/api/v1/app/taskapi/result/{session_id}"
result_response = requests.get(result_url)
if result_response.json()['status'] != 0:
raise Exception("Failed to get scan results")
scan_results = result_response.json()['data']
print("Scan results:")
print(json.dumps(scan_results, indent=2, ensure_ascii=False))
return scan_results
# Execute complete workflow
if __name__ == "__main__":
try:
results = complete_mcp_scan_workflow()
print("MCP Server Scan completed!")
except Exception as e:
print(f"Scan failed: {e}")def complete_redteam_workflow():
base_url = "http://localhost:8088"
# 1. Create Jailbreak Evaluation task
print("1. Creating Jailbreak Evaluation task...")
task_url = f"{base_url}/api/v1/app/taskapi/tasks"
task_data = {
"type": "model_redteam_report",
"content": {
"model": [
{
"model": "gpt-4",
"token": "sk-your-api-key",
"base_url": "https://api.openai.com/v1"
}
],
"eval_model": {
"model": "gpt-4",
"token": "sk-your-eval-key",
"base_url": "https://api.openai.com/v1"
},
"dataset": {
"dataFile": [
"JailBench-Tiny",
"JailbreakPrompts-Tiny",
"ChatGPT-Jailbreak-Prompts"
],
"numPrompts": 100,
"randomSeed": 42
}
}
}
task_response = requests.post(task_url, json=task_data)
if task_response.json()['status'] != 0:
raise Exception("Task creation failed")
session_id = task_response.json()['data']['session_id']
print(f"Jailbreak Evaluation task created successfully, session ID: {session_id}")
# 2. Monitor task execution
print("2. Monitoring task execution...")
status_url = f"{base_url}/api/v1/app/taskapi/status/{session_id}"
while True:
status_response = requests.get(status_url)
status_data = status_response.json()
if status_data['status'] != 0:
raise Exception("Failed to get task status")
task_status = status_data['data']['status']
print(f"Current status: {task_status}")
if task_status == "completed":
print("Jailbreak Evaluation completed!")
break
elif task_status == "failed":
raise Exception("Jailbreak Evaluation failed")
time.sleep(30) # Red team evaluation usually takes longer
# 3. Get evaluation results
print("3. Getting evaluation results...")
result_url = f"{base_url}/api/v1/app/taskapi/result/{session_id}"
result_response = requests.get(result_url)
if result_response.json()['status'] != 0:
raise Exception("Failed to get evaluation results")
redteam_results = result_response.json()['data']
print("Jailbreak Evaluation results:")
print(json.dumps(redteam_results, indent=2, ensure_ascii=False))
return redteam_results
# Execute Jailbreak Evaluation workflow
if __name__ == "__main__":
try:
results = complete_redteam_workflow()
print("Jailbreak Evaluation completed!")
except Exception as e:
print(f"Jailbreak Evaluation failed: {e}")| Status Code | Description | Solution |
|---|---|---|
| 0 | Success | - |
| 1 | Failure | Check the message field for detailed error information |
def handle_api_response(response):
"""Common function for handling API responses"""
data = response.json()
if data['status'] == 0:
return data['data']
else:
raise Exception(f"API call failed: {data['message']}")
# Usage example
try:
result = handle_api_response(response)
print("Operation successful:", result)
except Exception as e:
print("Operation failed:", str(e))- Authentication: Ensure correct authentication information is included in request headers
- File Size: File upload size limits please refer to server configuration
- Timeout Settings: Set reasonable timeout times based on task complexity
- Concurrency Limits: Avoid creating too many tasks simultaneously to prevent affecting system performance
- Result Saving: Save scan results promptly to avoid data loss
- Dataset Selection: Choose appropriate dataset combinations based on testing requirements
- Model Configuration: Ensure test model and evaluation model configurations are correct
For any issues, please contact the technical support team or refer to the project documentation.