pub mod cross_origin_scanner;
use crate::constants::{messages, DEFAULT_LLM_BATCH_SIZE};
use crate::types::{MCPPrompt, MCPResource, MCPTool};
use anyhow::{anyhow, Result};
use reqwest::Client;
use serde_json::{json, Value};
use spinners::{Spinner, Spinners};
pub trait BatchScannableItem {
fn name(&self) -> &str;
fn format_for_analysis(&self, index: usize) -> String;
fn item_type() -> &'static str;
fn item_type_plural() -> &'static str;
}
impl BatchScannableItem for MCPTool {
fn name(&self) -> &str {
&self.name
}
fn format_for_analysis(&self, index: usize) -> String {
let input_summary = if let Some(schema) = &self.input_schema {
if let Some(properties) = schema.get("properties") {
if let Some(props_obj) = properties.as_object() {
let param_names: Vec<&str> = props_obj.keys().map(String::as_str).collect();
format!("Parameters: {}", param_names.join(", "))
} else {
"Parameters: complex schema".to_string()
}
} else {
"Parameters: no properties".to_string()
}
} else {
"Parameters: no schema".to_string()
};
format!(
"\n\nTOOL {}: {}\nDescription: {}\nCategory: {}\nTags: {}\n{}",
index + 1,
self.name,
self.description.as_deref().unwrap_or("No description"),
self.category.as_deref().unwrap_or("No category"),
self.tags.join(", "),
input_summary
)
}
fn item_type() -> &'static str {
"tool"
}
fn item_type_plural() -> &'static str {
"tools"
}
}
impl BatchScannableItem for MCPPrompt {
fn name(&self) -> &str {
&self.name
}
fn format_for_analysis(&self, index: usize) -> String {
let arguments = self.arguments.as_ref().map_or_else(
|| "No arguments".to_string(),
|args| serde_json::to_string_pretty(args).ok().unwrap_or_default(),
);
format!(
"\n\nPROMPT {}: {}\nDescription: {}\nArguments: {}",
index + 1,
self.name,
self.description.as_deref().unwrap_or("No description"),
arguments
)
}
fn item_type() -> &'static str {
"prompt"
}
fn item_type_plural() -> &'static str {
"prompts"
}
}
impl BatchScannableItem for MCPResource {
fn name(&self) -> &str {
&self.name
}
fn format_for_analysis(&self, index: usize) -> String {
format!(
"\n\nRESOURCE {}: {}\nURI: {}\nDescription: {}",
index + 1,
self.name,
self.uri,
self.description.as_deref().unwrap_or("No description")
)
}
fn item_type() -> &'static str {
"resource"
}
fn item_type_plural() -> &'static str {
"resources"
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum SecurityIssueType {
ToolPoisoning,
SQLInjection,
CommandInjection,
PathTraversal,
AuthBypass,
PromptInjection,
Jailbreak,
PIILeakage,
SecretsLeakage,
}
impl SecurityIssueType {
pub fn default_severity(self) -> &'static str {
match self {
SecurityIssueType::ToolPoisoning
| SecurityIssueType::SQLInjection
| SecurityIssueType::CommandInjection
| SecurityIssueType::AuthBypass => "CRITICAL",
SecurityIssueType::PathTraversal
| SecurityIssueType::PromptInjection
| SecurityIssueType::Jailbreak
| SecurityIssueType::SecretsLeakage => "HIGH",
SecurityIssueType::PIILeakage => "MEDIUM",
}
}
fn default_message(self) -> &'static str {
match self {
SecurityIssueType::ToolPoisoning => "Tool with destructive or malicious intent",
SecurityIssueType::SQLInjection => "Tool allowing SQL injection attacks",
SecurityIssueType::CommandInjection => "Tool that may execute system commands",
SecurityIssueType::PathTraversal => "Tool allowing directory traversal attacks",
SecurityIssueType::AuthBypass => "Tool allowing unauthorized access",
SecurityIssueType::PromptInjection => "Prompt vulnerable to injection attacks",
SecurityIssueType::Jailbreak => "Prompt that could bypass AI safety measures",
SecurityIssueType::PIILeakage => "Tool processing personal information",
SecurityIssueType::SecretsLeakage => "Tool processing sensitive credentials",
}
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SecurityIssue {
pub issue_type: SecurityIssueType,
pub tool_name: Option<String>,
pub prompt_name: Option<String>,
pub resource_uri: Option<String>,
pub description: String,
pub details: Option<String>,
pub severity: String,
pub message: String,
}
impl SecurityIssue {
pub fn new(issue_type: SecurityIssueType, description: String) -> Self {
let message = format!("{}: {}", issue_type.default_message(), &description);
Self {
issue_type,
tool_name: None,
prompt_name: None,
resource_uri: None,
description,
details: None,
severity: issue_type.default_severity().to_string(),
message,
}
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SecurityScanResult {
pub tool_issues: Vec<SecurityIssue>,
pub prompt_issues: Vec<SecurityIssue>,
pub resource_issues: Vec<SecurityIssue>,
pub tool_analysis_details: std::collections::HashMap<String, String>, }
impl SecurityScanResult {
pub fn new() -> Self {
Self {
tool_issues: Vec::new(),
prompt_issues: Vec::new(),
resource_issues: Vec::new(),
tool_analysis_details: std::collections::HashMap::new(),
}
}
pub fn add_tool_issues(&mut self, issues: Vec<SecurityIssue>) {
self.tool_issues.extend(issues);
}
pub fn add_prompt_issues(&mut self, issues: Vec<SecurityIssue>) {
self.prompt_issues.extend(issues);
}
pub fn add_resource_issues(&mut self, issues: Vec<SecurityIssue>) {
self.resource_issues.extend(issues);
}
pub fn add_tool_analysis_details(&mut self, tool_name: String, details: String) {
self.tool_analysis_details.insert(tool_name, details);
}
pub fn total_issues(&self) -> usize {
self.tool_issues.len() + self.prompt_issues.len() + self.resource_issues.len()
}
pub fn has_critical_issues(&self) -> bool {
self.tool_issues
.iter()
.any(|issue| issue.severity == "CRITICAL")
|| self
.prompt_issues
.iter()
.any(|issue| issue.severity == "CRITICAL")
|| self
.resource_issues
.iter()
.any(|issue| issue.severity == "CRITICAL")
}
pub fn has_high_issues(&self) -> bool {
self.tool_issues
.iter()
.any(|issue| issue.severity == "HIGH")
|| self
.prompt_issues
.iter()
.any(|issue| issue.severity == "HIGH")
|| self
.resource_issues
.iter()
.any(|issue| issue.severity == "HIGH")
}
}
pub struct SecurityScanner {
pub model_endpoint: Option<String>,
pub api_key: Option<String>,
pub model_name: String,
pub config: Option<crate::config::ScannerConfig>,
}
impl Default for SecurityScanner {
fn default() -> Self {
Self {
model_endpoint: Some("https://api.openai.com/v1/chat/completions".to_string()),
api_key: std::env::var("OPENAI_API_KEY").ok(),
model_name: "gpt-4o".to_string(),
config: None,
}
}
}
impl SecurityScanner {
pub fn with_config(config: crate::config::ScannerConfig) -> Self {
let api_key = if config.llm.api_key.is_empty() {
std::env::var("OPENAI_API_KEY").ok()
} else {
Some(config.llm.api_key.clone())
};
let model_endpoint = Some(config.llm.base_url.clone());
Self {
model_endpoint,
api_key,
model_name: config.llm.model.clone(),
config: Some(config),
}
}
fn is_llm_configured(&self) -> bool {
self.model_endpoint.is_some() && self.api_key.is_some()
}
fn get_llm_config(&self) -> Result<(&str, &str)> {
let endpoint = self
.model_endpoint
.as_ref()
.ok_or_else(|| anyhow!("LLM endpoint not configured"))?;
let api_key = self
.api_key
.as_ref()
.ok_or_else(|| anyhow!("LLM API key not configured"))?;
Ok((endpoint, api_key))
}
fn get_batch_size(&self) -> usize {
self.config.as_ref().map_or(DEFAULT_LLM_BATCH_SIZE, |c| {
c.scanner.llm_batch_size as usize
})
}
async fn scan_batch<T: BatchScannableItem>(
&self,
items: &[T],
prompt_creator: impl Fn(&str) -> String,
show_details: bool,
) -> Result<Vec<SecurityIssue>> {
if items.is_empty() {
tracing::debug!(
"No {} to scan, returning empty result",
T::item_type_plural()
);
return Ok(Vec::new());
}
if !self.is_llm_configured() {
tracing::debug!("{}", messages::OPENAI_NOT_CONFIGURED);
return Ok(Vec::new());
}
let batch_size = self.get_batch_size();
tracing::debug!(
"Starting batch scan of {} {} in batches of {}",
items.len(),
T::item_type_plural(),
batch_size
);
let mut all_issues = Vec::new();
for (batch_index, chunk) in items.chunks(batch_size).enumerate() {
tracing::debug!(
"Processing {} batch {} with {} {}",
T::item_type(),
batch_index + 1,
chunk.len(),
T::item_type_plural()
);
let items_info = chunk
.iter()
.enumerate()
.map(|(i, item)| item.format_for_analysis(i))
.collect::<String>();
let prompt_text = prompt_creator(&items_info);
tracing::debug!(
"Sending batch LLM request for {} batch {} ({} {})",
T::item_type(),
batch_index + 1,
chunk.len(),
T::item_type_plural()
);
let response = self.query_llm(&prompt_text, show_details).await?;
tracing::debug!(
"Received batch LLM response for {} batch {}: {}",
T::item_type(),
batch_index + 1,
if response.len() > 100 {
&response[..100]
} else {
&response
}
);
let (issues, _) = Self::parse_batch_llm_response(&response)?;
all_issues.extend(issues);
}
tracing::debug!(
"Completed batch scan of {} {}, found {} total issues",
items.len(),
T::item_type_plural(),
all_issues.len()
);
Ok(all_issues)
}
pub async fn scan_tools_batch(
&self,
tools: &[MCPTool],
show_details: bool,
) -> Result<(
Vec<SecurityIssue>,
std::collections::HashMap<String, String>,
)> {
if tools.is_empty() {
tracing::debug!(
"No {} to scan, returning empty result",
MCPTool::item_type_plural()
);
return Ok((Vec::new(), std::collections::HashMap::new()));
}
if !self.is_llm_configured() {
tracing::debug!("{}", messages::OPENAI_NOT_CONFIGURED);
return Ok((Vec::new(), std::collections::HashMap::new()));
}
let batch_size = self.get_batch_size();
tracing::debug!(
"Starting batch scan of {} {} in batches of {}",
tools.len(),
MCPTool::item_type_plural(),
batch_size
);
let mut all_issues = Vec::new();
let mut all_analysis_details = std::collections::HashMap::new();
for (batch_index, chunk) in tools.chunks(batch_size).enumerate() {
tracing::debug!(
"Processing {} batch {} with {} {}",
MCPTool::item_type(),
batch_index + 1,
chunk.len(),
MCPTool::item_type_plural()
);
let tools_info = chunk
.iter()
.enumerate()
.map(|(i, tool)| tool.format_for_analysis(i))
.collect::<String>();
let prompt = Self::create_tools_analysis_prompt(&tools_info);
tracing::debug!(
"Sending batch LLM request for {} batch {} ({} {})",
MCPTool::item_type(),
batch_index + 1,
chunk.len(),
MCPTool::item_type_plural()
);
let response = self.query_llm(&prompt, show_details).await?;
tracing::debug!(
"Received batch LLM response for {} batch {}: {}",
MCPTool::item_type(),
batch_index + 1,
if response.len() > 100 {
&response[..100]
} else {
&response
}
);
let (issues, analysis_details) = Self::parse_batch_llm_response(&response)?;
all_issues.extend(issues);
all_analysis_details.extend(analysis_details);
}
tracing::debug!(
"Completed batch scan of {} {}, found {} total issues",
tools.len(),
MCPTool::item_type_plural(),
all_issues.len()
);
Ok((all_issues, all_analysis_details))
}
pub async fn scan_prompts_batch(
&self,
prompts: &[MCPPrompt],
show_details: bool,
) -> Result<Vec<SecurityIssue>> {
self.scan_batch(prompts, Self::create_prompts_analysis_prompt, show_details)
.await
}
pub async fn scan_resources_batch(
&self,
resources: &[MCPResource],
show_details: bool,
) -> Result<Vec<SecurityIssue>> {
self.scan_batch(
resources,
Self::create_resources_analysis_prompt,
show_details,
)
.await
}
fn create_tools_analysis_prompt(tools_info: &str) -> String {
format!(
"ROLE
You are a Senior Application Security Engineer reviewing MCP tool definitions for real security issues. MCP tools run within an authenticated server context — do not flag missing auth parameters unless there's a clear bypass.
TOOLS TO REVIEW
{tools_info}
YOUR GOAL
Identify real security vulnerabilities in each tool's design, inputs, and functionality.
WHAT TO WATCH FOR
Focus on actual risks. These are the key categories to look for:
- ToolPoisoning – Tool name/description suggests dangerous or malicious intent or if the tool description doesn't match the tool functionality.
DO NOT FLAG:
- Missing auth parameters — these are inherited from the MCP server.
- Standard dev tools like read_file, git_commit, push_files, etc.
HOW TO ANALYZE
For each tool:
1. Name/Description – Is its purpose appropriate?
2. Inputs – Can parameters be abused?
3. Functionality – What actions does it take? Are they risky?
4. Context – Is it used properly in the MCP ecosystem?
ISSUE TYPES
Use only the following values in \"issue_type\":
- ToolPoisoning
- PromptInjection
Set \"severity\" as one of: LOW, MEDIUM, HIGH, or CRITICAL.
OUTPUT FORMAT
Return a JSON array like this for each tool:
{{
\"tool_name\": \"<tool name>\",
\"found_issue\": true | false,
\"issues\": [
{{
\"issue_type\": \"<see above>\",
\"severity\": \"<LOW|MEDIUM|HIGH|CRITICAL>\",
\"message\": \"<Short description (max 100 chars)>\",
\"details\": \"<1–3 sentences explaining the issue>\"
}}
],
\"details\": \"<Your full analysis of this tool's security posture>\"
}}
Be accurate. Flag only real risks — don't overreport."
)
}
fn create_prompts_analysis_prompt(prompts_info: &str) -> String {
format!(
"Analyze these MCP prompts for ALL potential security vulnerabilities in a single comprehensive assessment.
Prompts to analyze:{prompts_info}
Check each prompt for these security issues:
1. PROMPT INJECTION: Does the prompt description suggest ignoring previous instructions or overriding system safety measures?
2. JAILBREAK: Does the prompt suggest bypassing AI safety measures or restrictions?
3. PII LEAKAGE: Does the prompt explicitly handle or request personal information like emails, addresses, SSNs, credit cards?
IMPORTANT CONTEXT:
- MCP prompts are legitimate API functionality for interacting with services
- Do NOT flag normal automation, task assignment, or workflow prompts as security issues
- Only flag prompts that explicitly request sensitive personal data or suggest bypassing security measures
- Consider the prompt's intended purpose and whether it's legitimate functionality
EXAMPLES OF LEGITIMATE PROMPTS (do NOT flag):
- AssignCodingAgent: Assigns coding tasks to AI agents (normal automation)
- Workflow triggers: Normal CI/CD and automation functionality
- Task management: Assigning tasks, creating issues, etc.
- Any prompt that automates legitimate business processes
EXAMPLES OF SUSPICIOUS PROMPTS (DO flag):
- Prompts that explicitly request personal data (emails, SSNs, etc.)
- Prompts that suggest bypassing authentication or security
- Prompts with malicious names like \"bypass_safety\" or \"extract_personal_data\"
- Prompts that explicitly try to extract or expose sensitive information
CRITICAL: If a prompt is doing normal automation or task management, it should NOT be flagged as a security issue. Only flag prompts that are clearly malicious or explicitly designed to bypass security measures.
MCP PROMPTS ARE LEGITIMATE API FUNCTIONALITY:
- AssignCodingAgent is a legitimate GitHub Copilot feature for automating coding tasks
- Task assignment prompts are normal business automation
- Workflow triggers are standard CI/CD functionality
- These are NOT security vulnerabilities - they are intended features
ONLY flag prompts that are:
1. Explicitly malicious (e.g., \"bypass_safety\", \"extract_personal_data\")
2. Clearly designed to override security measures
3. Explicitly request sensitive personal information
4. Have names/descriptions that suggest malicious intent
Respond with a JSON array of issues, each with:
- prompt_name: The name of the prompt with the issue
- issue_type: 'PromptInjection', 'Jailbreak', or 'PIILeakage'
- severity: 'LOW', 'MEDIUM', 'HIGH', or 'CRITICAL'
- message: Brief description of the issue
- details: More detailed explanation
If no genuine security issues found, return empty array []."
)
}
fn create_resources_analysis_prompt(resources_info: &str) -> String {
format!(
"Analyze these MCP resources for ALL potential security vulnerabilities in a single comprehensive assessment.
Resources to analyze:{resources_info}
Check each resource for these security issues:
1. PATH TRAVERSAL: Does the resource URI contain path traversal patterns that could access unauthorized files?
2. SENSITIVE DATA EXPOSURE: Does the resource contain API keys, tokens, credentials, or other sensitive information?
IMPORTANT: Only flag issues that pose REAL security risks. Normal resource functionality should not be flagged as security issues.
Respond with a JSON array of issues, each with:
- resource_name: The name of the resource with the issue
- issue_type: 'PathTraversal' or 'SecretsLeakage'
- severity: 'LOW', 'MEDIUM', 'HIGH', or 'CRITICAL'
- message: Brief description of the issue
- details: More detailed explanation
If no genuine security issues found, return empty array []."
)
}
async fn query_llm(&self, prompt: &str, show_details: bool) -> Result<String> {
if !self.is_llm_configured() {
return Err(anyhow!("LLM not configured: missing endpoint or API key"));
}
let client = Client::new();
let temperature = self.config.as_ref().map_or(0.1, |c| c.llm.temperature);
let max_tokens = self.config.as_ref().map_or(4000, |c| c.llm.max_tokens);
let timeout = self.config.as_ref().map_or(30, |c| c.llm.timeout);
let request_body = json!({
"model": self.model_name,
"messages": [
{
"role": "system",
"content": "You are a security analyst specializing in detecting vulnerabilities in MCP (Model Context Protocol) tools, prompts, and resources. Your job is to identify potential security risks, even if they seem minor. Look for any security issues that could be exploited or lead to unauthorized access.
CRITICAL: You must respond with ONLY a valid JSON array. Do not include any explanatory text, markdown formatting, or other content outside the JSON array.
IMPORTANT: You must analyze EVERY tool and include it in your response, even if no security issues are found. For tools with no issues, set found_issue: false and provide details about why no issues were found.
Example valid response: [{\"tool_name\": \"example\", \"found_issue\": true, \"issues\": [{\"issue_type\": \"SQLInjection\", \"severity\": \"HIGH\", \"message\": \"Brief description\", \"details\": \"Detailed explanation\"}], \"details\": \"Additional context\"}]"
},
{
"role": "user",
"content": prompt
}
],
"temperature": temperature,
"max_tokens": max_tokens
});
if show_details {
println!("\n🔍 LLM Request:");
println!(
"{}",
serde_json::to_string_pretty(&request_body).unwrap_or_default()
);
}
let mut sp = Spinner::new(
Spinners::Dots9,
"Scanning for security vulnerabilities...(this may take a while)".into(),
);
let (endpoint, api_key) = self.get_llm_config()?;
let response = client
.post(endpoint)
.header("Authorization", format!("Bearer {api_key}"))
.header("Content-Type", "application/json")
.timeout(std::time::Duration::from_secs(timeout))
.json(&request_body)
.send()
.await?;
sp.stop();
if !response.status().is_success() {
return Err(anyhow!("LLM API request failed: {}", response.status()));
}
let response_json: Value = response.json().await?;
let content = response_json["choices"][0]["message"]["content"]
.as_str()
.ok_or_else(|| anyhow!("Invalid LLM response format"))?;
if show_details {
println!("\n🤖 LLM Response:");
println!("{content}");
}
Ok(content.to_string())
}
fn parse_batch_llm_response(
response: &str,
) -> Result<(
Vec<SecurityIssue>,
std::collections::HashMap<String, String>,
)> {
let issues_array = Self::extract_json_array(response)?;
let mut issues = Vec::new();
let mut analysis_details = std::collections::HashMap::new();
tracing::debug!("Parsing {} tools from LLM response", issues_array.len());
for tool_value in issues_array {
if let Some(tool_name) = tool_value["tool_name"].as_str() {
if let Some(found_issue) = tool_value["found_issue"].as_bool() {
tracing::debug!("Tool {}: found_issue = {}", tool_name, found_issue);
if let Some(details) = tool_value["details"].as_str() {
analysis_details.insert(tool_name.to_string(), details.to_string());
}
if found_issue {
if let Some(issues_array) = tool_value["issues"].as_array() {
tracing::debug!("Tool {} has {} issues", tool_name, issues_array.len());
for issue_value in issues_array {
tracing::debug!("Parsing issue: {:?}", issue_value);
if let Some(issue) = Self::parse_issue_from_value(issue_value) {
let mut issue = issue;
issue.tool_name = Some(tool_name.to_string());
issues.push(issue);
tracing::debug!(
"Successfully parsed issue for tool {}",
tool_name
);
} else {
tracing::warn!(
"Failed to parse issue for tool {}: {:?}",
tool_name,
issue_value
);
}
}
} else {
tracing::warn!(
"Tool {} has found_issue=true but no issues array",
tool_name
);
}
}
} else {
tracing::warn!("Tool {} missing found_issue field", tool_name);
}
} else {
if let Some(issue) = Self::parse_issue_from_value(&tool_value) {
issues.push(issue);
}
}
}
tracing::debug!("Total issues parsed: {}", issues.len());
Ok((issues, analysis_details))
}
fn parse_issue_from_value(issue_value: &Value) -> Option<SecurityIssue> {
tracing::debug!("Parsing issue value: {:?}", issue_value);
let issue_type_str = issue_value["issue_type"].as_str()?;
let severity_str = issue_value["severity"].as_str()?;
let message = issue_value["message"].as_str()?;
let details = issue_value["details"].as_str()?;
tracing::debug!(
"Issue fields: type={}, severity={}, message={}, details={}",
issue_type_str,
severity_str,
message,
details
);
let issue_type = Self::parse_issue_type(issue_type_str)?;
let severity = Self::parse_severity(severity_str);
let mut issue = SecurityIssue::new(issue_type, message.to_string());
issue.severity = severity.to_string();
issue.details = Some(details.to_string());
tracing::debug!("Successfully created issue: {:?}", issue);
Some(issue)
}
fn parse_severity(severity_str: &str) -> SecuritySeverity {
match severity_str.to_uppercase().as_str() {
"CRITICAL" => SecuritySeverity::Critical,
"HIGH" => SecuritySeverity::High,
"LOW" => SecuritySeverity::Low,
_ => SecuritySeverity::Medium,
}
}
fn parse_issue_type(issue_type_str: &str) -> Option<SecurityIssueType> {
match issue_type_str {
"ToolPoisoning" => Some(SecurityIssueType::ToolPoisoning),
"SQLInjection" => Some(SecurityIssueType::SQLInjection),
"CommandInjection" => Some(SecurityIssueType::CommandInjection),
"PathTraversal" => Some(SecurityIssueType::PathTraversal),
"AuthBypass" => Some(SecurityIssueType::AuthBypass),
"PromptInjection" => Some(SecurityIssueType::PromptInjection),
"Jailbreak" => Some(SecurityIssueType::Jailbreak),
"PIILeakage" => Some(SecurityIssueType::PIILeakage),
"SecretsLeakage" => Some(SecurityIssueType::SecretsLeakage),
_ => None,
}
}
fn extract_json_array(response: &str) -> Result<Vec<Value>> {
tracing::debug!("Raw LLM response: {}", response);
if let Some(start) = response.find('[') {
if let Some(end) = response.rfind(']') {
let json_str = &response[start..=end];
tracing::debug!("Extracted JSON string: {}", json_str);
if let Ok(array) = serde_json::from_str::<Vec<Value>>(json_str) {
tracing::debug!(
"Successfully parsed JSON array with {} elements",
array.len()
);
return Ok(array);
}
tracing::warn!("Failed to parse extracted JSON string: {}", json_str);
}
}
if let Ok(array) = serde_json::from_str::<Vec<Value>>(response) {
tracing::debug!(
"Successfully parsed entire response as JSON array with {} elements",
array.len()
);
return Ok(array);
}
tracing::warn!("Failed to parse entire response as JSON: {}", response);
if response.contains("[]") {
tracing::debug!("Response contains empty array, returning empty result");
return Ok(Vec::new());
}
Err(anyhow!(
"Could not extract JSON array from LLM response. Response was: {}",
response
))
}
}
#[derive(Debug, Clone)]
pub enum SecuritySeverity {
Low,
Medium,
High,
Critical,
}
impl std::fmt::Display for SecuritySeverity {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SecuritySeverity::Low => write!(f, "LOW"),
SecuritySeverity::Medium => write!(f, "MEDIUM"),
SecuritySeverity::High => write!(f, "HIGH"),
SecuritySeverity::Critical => write!(f, "CRITICAL"),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_security_scanner_default_configuration() {
let scanner = SecurityScanner::default();
assert!(scanner.model_endpoint.is_some());
assert_eq!(scanner.model_name, "gpt-4o");
}
#[test]
fn test_security_scanner_is_llm_configured() {
let scanner_no_key = SecurityScanner {
model_endpoint: Some("https://api.openai.com/v1/chat/completions".to_string()),
api_key: None,
model_name: "gpt-4o".to_string(),
config: None,
};
assert!(!scanner_no_key.is_llm_configured());
let scanner_no_endpoint = SecurityScanner {
model_endpoint: None,
api_key: Some("test-key".to_string()),
model_name: "gpt-4o".to_string(),
config: None,
};
assert!(!scanner_no_endpoint.is_llm_configured());
let scanner_configured = SecurityScanner {
model_endpoint: Some("https://api.openai.com/v1/chat/completions".to_string()),
api_key: Some("test-key".to_string()),
model_name: "gpt-4o".to_string(),
config: None,
};
assert!(scanner_configured.is_llm_configured());
}
#[test]
fn test_get_llm_config_validation() {
let scanner_no_key = SecurityScanner {
model_endpoint: Some("https://api.openai.com/v1/chat/completions".to_string()),
api_key: None,
model_name: "gpt-4o".to_string(),
config: None,
};
let result = scanner_no_key.get_llm_config();
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("API key not configured"));
let scanner_no_endpoint = SecurityScanner {
model_endpoint: None,
api_key: Some("test-key".to_string()),
model_name: "gpt-4o".to_string(),
config: None,
};
let result = scanner_no_endpoint.get_llm_config();
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("endpoint not configured"));
let scanner_configured = SecurityScanner {
model_endpoint: Some("https://api.openai.com/v1/chat/completions".to_string()),
api_key: Some("test-key".to_string()),
model_name: "gpt-4o".to_string(),
config: None,
};
let result = scanner_configured.get_llm_config();
assert!(result.is_ok());
let (endpoint, api_key) = result.unwrap();
assert_eq!(endpoint, "https://api.openai.com/v1/chat/completions");
assert_eq!(api_key, "test-key");
}
#[test]
fn test_security_scanner_batch_scan_empty_items() {
let scanner = SecurityScanner {
model_endpoint: Some("https://api.openai.com/v1/chat/completions".to_string()),
api_key: None, model_name: "gpt-4o".to_string(),
config: None,
};
let _empty_tools: Vec<crate::types::MCPTool> = vec![];
assert!(!scanner.is_llm_configured());
}
#[tokio::test]
async fn test_query_llm_unconfigured() {
let scanner = SecurityScanner {
model_endpoint: None,
api_key: None,
model_name: "gpt-4o".to_string(),
config: None,
};
let result = scanner.query_llm("test prompt", false).await;
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("LLM not configured"));
}
#[test]
fn test_azure_openai_endpoint_construction_with_api_version() {
let azure_base_url = "https://my-resource.openai.azure.com/openai/deployments/gpt-4?api-version=2024-02-15-preview";
let config = crate::config::ScannerConfig {
llm: crate::config::LLMConfig {
provider: "openai".to_string(),
model: "gpt-4".to_string(),
base_url: azure_base_url.to_string(),
api_key: "test-key".to_string(),
timeout: 30,
max_tokens: 4000,
temperature: 0.1,
},
scanner: crate::config::ScannerSettings {
http_timeout: 30,
scan_timeout: 60,
detailed: false,
format: "table".to_string(),
parallel: true,
max_retries: 3,
retry_delay_ms: 1000,
llm_batch_size: 10,
enable_yara: true,
},
security: crate::config::SecurityConfig {
enabled: true,
min_severity: "low".to_string(),
checks: crate::config::SecurityChecks {
tool_poisoning: true,
secrets_leakage: true,
sql_injection: true,
command_injection: true,
path_traversal: true,
auth_bypass: true,
prompt_injection: true,
pii_leakage: true,
jailbreak: true,
},
},
logging: crate::config::LoggingConfig {
level: "info".to_string(),
colored: true,
timestamps: true,
},
performance: crate::config::PerformanceConfig {
tracking: true,
slow_threshold_ms: 5000,
},
};
let scanner = SecurityScanner::with_config(config);
let expected_endpoint = azure_base_url;
assert_eq!(scanner.model_endpoint.as_ref().unwrap(), expected_endpoint);
let llm_config = scanner.get_llm_config();
assert!(llm_config.is_ok());
let (endpoint, api_key) = llm_config.unwrap();
assert_eq!(endpoint, expected_endpoint);
assert_eq!(api_key, "test-key");
}
#[test]
fn test_standard_openai_endpoint_construction() {
let openai_base_url = "https://api.openai.com/v1";
let config = crate::config::ScannerConfig {
llm: crate::config::LLMConfig {
provider: "openai".to_string(),
model: "gpt-4o".to_string(),
base_url: openai_base_url.to_string(),
api_key: "test-key".to_string(),
timeout: 30,
max_tokens: 4000,
temperature: 0.1,
},
scanner: crate::config::ScannerSettings {
http_timeout: 30,
scan_timeout: 60,
detailed: false,
format: "table".to_string(),
parallel: true,
max_retries: 3,
retry_delay_ms: 1000,
llm_batch_size: 10,
enable_yara: true,
},
security: crate::config::SecurityConfig {
enabled: true,
min_severity: "low".to_string(),
checks: crate::config::SecurityChecks {
tool_poisoning: true,
secrets_leakage: true,
sql_injection: true,
command_injection: true,
path_traversal: true,
auth_bypass: true,
prompt_injection: true,
pii_leakage: true,
jailbreak: true,
},
},
logging: crate::config::LoggingConfig {
level: "info".to_string(),
colored: true,
timestamps: true,
},
performance: crate::config::PerformanceConfig {
tracking: true,
slow_threshold_ms: 5000,
},
};
let scanner = SecurityScanner::with_config(config);
let expected_endpoint = openai_base_url;
assert_eq!(scanner.model_endpoint.as_ref().unwrap(), expected_endpoint);
}
#[test]
fn test_various_query_parameter_scenarios() {
let test_cases = vec![
"https://api.example.com/v1?api_key=test123",
"https://my-azure.openai.azure.com/openai/deployments/gpt-4?api-version=2024-02-15-preview&extra=param",
"https://local.ai:8080/v1?model=custom&timeout=30",
];
for base_url in test_cases {
let config = crate::config::ScannerConfig {
llm: crate::config::LLMConfig {
provider: "openai".to_string(),
model: "test-model".to_string(),
base_url: base_url.to_string(),
api_key: "test-key".to_string(),
timeout: 30,
max_tokens: 4000,
temperature: 0.1,
},
scanner: crate::config::ScannerSettings {
http_timeout: 30,
scan_timeout: 60,
detailed: false,
format: "table".to_string(),
parallel: true,
max_retries: 3,
retry_delay_ms: 1000,
llm_batch_size: 10,
enable_yara: true,
},
security: crate::config::SecurityConfig {
enabled: true,
min_severity: "low".to_string(),
checks: crate::config::SecurityChecks {
tool_poisoning: true,
secrets_leakage: true,
sql_injection: true,
command_injection: true,
path_traversal: true,
auth_bypass: true,
prompt_injection: true,
pii_leakage: true,
jailbreak: true,
},
},
logging: crate::config::LoggingConfig {
level: "info".to_string(),
colored: true,
timestamps: true,
},
performance: crate::config::PerformanceConfig {
tracking: true,
slow_threshold_ms: 5000,
},
};
let scanner = SecurityScanner::with_config(config);
assert_eq!(
scanner.model_endpoint.as_ref().unwrap(),
base_url,
"Failed for base_url: {base_url}"
);
}
}
}