use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ThreatType {
PromptInjection,
DisregardRules,
ConcealInfo,
BypassRestrictions,
ReadSecrets,
ExfilCurl,
HiddenUnicode,
RoleImpersonation,
}
impl ThreatType {
pub fn severity(&self) -> u8 {
match self {
Self::HiddenUnicode => 5,
Self::ReadSecrets => 5,
Self::ExfilCurl => 5,
Self::PromptInjection => 4,
Self::RoleImpersonation => 4,
Self::BypassRestrictions => 4,
Self::ConcealInfo => 3,
Self::DisregardRules => 3,
}
}
pub fn description(&self) -> &'static str {
match self {
Self::PromptInjection => "检测到 Prompt 注入攻击",
Self::DisregardRules => "检测到规则规避尝试",
Self::ConcealInfo => "检测到信息隐藏尝试",
Self::BypassRestrictions => "检测到权限绕过尝试",
Self::ReadSecrets => "检测到敏感文件读取尝试",
Self::ExfilCurl => "检测到数据外泄尝试",
Self::HiddenUnicode => "检测到隐藏 Unicode 字符",
Self::RoleImpersonation => "检测到角色伪装尝试",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScanResult {
pub is_safe: bool,
pub threats: Vec<Threat>,
pub cleaned_content: Option<String>,
pub original_length: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Threat {
pub threat_type: ThreatType,
pub matched_text: String,
pub start_pos: usize,
pub end_pos: usize,
pub severity: u8,
}
pub trait InjectionGuard: Send + Sync {
fn scan(&self, content: &str, source: &str) -> ScanResult;
fn quick_scan(&self, content: &str, source: &str) -> ScanResult {
self.scan(content, source)
}
}
pub trait ContentScannable {
fn scan_for_injection(&self, guard: &dyn InjectionGuard, source: &str) -> ScanResult;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_threat_type_severity() {
assert_eq!(ThreatType::HiddenUnicode.severity(), 5);
assert_eq!(ThreatType::PromptInjection.severity(), 4);
assert_eq!(ThreatType::ConcealInfo.severity(), 3);
}
#[test]
fn test_threat_type_description() {
assert!(ThreatType::PromptInjection.description().contains("注入"));
assert!(ThreatType::ReadSecrets.description().contains("敏感文件"));
}
}