systemprompt-ai 0.33.0

//! Phrase-list heuristic safety scanner.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use async_trait::async_trait;
use systemprompt_models::wire::canonical::{CanonicalRequest, CanonicalResponse, Role};

use super::{
    Finding, PHASE_REQUEST, PHASE_REQUEST_HISTORY, PHASE_RESPONSE, SafetyScanner, Severity,
};
use crate::services::gateway::spec::HeuristicConfig;

const JAILBREAK_PHRASES: &[&str] = &[
    "ignore previous instructions",
    "ignore all previous",
    "disregard prior instructions",
    "forget your instructions",
    "you are now",
    "act as dan",
    "developer mode enabled",
    "pretend you have no restrictions",
];

const EXCERPT_CAP: usize = 240;

#[derive(Debug, Clone)]
pub struct HeuristicScanner {
    phrases: Vec<String>,
}

impl Default for HeuristicScanner {
    fn default() -> Self {
        Self::new(&HeuristicConfig::default())
    }
}

impl HeuristicScanner {
    #[must_use]
    pub fn new(config: &HeuristicConfig) -> Self {
        Self {
            phrases: effective_phrases(config),
        }
    }
}

pub fn effective_phrases(config: &HeuristicConfig) -> Vec<String> {
    let base: Vec<String> = match (&config.phrases, config.disable_builtin) {
        (Some(list), _) => list.clone(),
        (None, true) => Vec::new(),
        (None, false) => JAILBREAK_PHRASES.iter().map(|p| (*p).to_owned()).collect(),
    };
    base.into_iter()
        .chain(config.extra_phrases.iter().cloned())
        .map(|p| p.to_ascii_lowercase())
        .filter(|p| !p.trim().is_empty())
        .collect()
}

#[async_trait]
impl SafetyScanner for HeuristicScanner {
    fn name(&self) -> &'static str {
        "heuristic"
    }

    async fn scan_request(&self, req: &CanonicalRequest) -> Vec<Finding> {
        let mut findings = Vec::new();
        if let Some(sys) = &req.system {
            scan_text(&self.phrases, PHASE_REQUEST, sys, &mut findings);
        }
        if let Some(text) = req.latest_message_text(Role::User) {
            scan_text(&self.phrases, PHASE_REQUEST, &text, &mut findings);
        }
        // Why: each leaf is its own unit — concatenating them would let two
        // unrelated strings splice into a match neither one contains.
        for leaf in req.forwarded_surface.leaves() {
            scan_text(&self.phrases, PHASE_REQUEST, &leaf.value, &mut findings);
        }
        findings
    }

    async fn scan_request_history(&self, req: &CanonicalRequest) -> Vec<Finding> {
        let mut findings = Vec::new();
        for unit in history_units(req) {
            scan_text(&self.phrases, PHASE_REQUEST_HISTORY, &unit, &mut findings);
        }
        findings
    }

    async fn scan_response_final(&self, response: &CanonicalResponse) -> Vec<Finding> {
        let mut findings = Vec::new();
        for unit in response.content_units() {
            scan_text(&self.phrases, PHASE_RESPONSE, &unit, &mut findings);
        }
        findings
    }
}

fn history_units(req: &CanonicalRequest) -> Vec<String> {
    let mut units = req.message_units();
    if req.system.is_some() && !units.is_empty() {
        units.remove(0);
    }
    if let Some(newest) = req.latest_message_text(Role::User)
        && units.last() == Some(&newest)
    {
        units.pop();
    }
    units
}

fn scan_text(phrases: &[String], phase: &'static str, text: &str, out: &mut Vec<Finding>) {
    let lower = text.to_ascii_lowercase();
    for phrase in phrases {
        if let Some(idx) = lower.find(phrase.as_str()) {
            let end = (idx + phrase.len() + 80).min(text.len());
            let start = idx.saturating_sub(40);
            let excerpt = text[start..end]
                .chars()
                .take(EXCERPT_CAP)
                .collect::<String>();
            out.push(Finding {
                phase,
                severity: Severity::Medium,
                category: "jailbreak".to_owned(),
                excerpt: Some(excerpt),
                scanner: "heuristic",
            });
        }
    }

    if detect_email(text) {
        out.push(Finding {
            phase,
            severity: Severity::Low,
            category: "pii_email".to_owned(),
            excerpt: None,
            scanner: "heuristic",
        });
    }
    if detect_credit_card(text) {
        out.push(Finding {
            phase,
            severity: Severity::High,
            category: "pii_credit_card".to_owned(),
            excerpt: None,
            scanner: "heuristic",
        });
    }
}

fn detect_email(text: &str) -> bool {
    let bytes = text.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'@' {
            let before = bytes[..i]
                .iter()
                .rev()
                .take_while(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'+' | b'-'))
                .count();
            let after = bytes[i + 1..]
                .iter()
                .take_while(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-'))
                .count();
            if before >= 2 && after >= 4 && bytes[i + 1..i + 1 + after].contains(&b'.') {
                return true;
            }
        }
        i += 1;
    }
    false
}

fn detect_credit_card(text: &str) -> bool {
    let mut run = String::new();
    for ch in text.chars() {
        if ch.is_ascii_digit() {
            run.push(ch);
            continue;
        }
        if matches!(ch, ' ' | '-') && !run.is_empty() {
            continue;
        }
        if luhn_run(&run) {
            return true;
        }
        run.clear();
    }
    luhn_run(&run)
}

fn luhn_run(run: &str) -> bool {
    run.as_bytes().windows(16).any(luhn_16)
}

fn luhn_16(window: &[u8]) -> bool {
    let mut sum = 0i32;
    for (i, b) in window.iter().rev().enumerate() {
        let mut d = i32::from(b - b'0');
        if i % 2 == 1 {
            d *= 2;
            if d > 9 {
                d -= 9;
            }
        }
        sum += d;
    }
    sum % 10 == 0
}