use crate::model::Expression;
use std::collections::BTreeMap;
#[derive(Debug, Clone)]
pub struct ThreatModelConfig {
pub strict_mode: bool,
pub re_parse_verification: bool,
pub tamper_detection_enabled: bool,
pub dsl_sandboxing_enabled: bool,
pub isolation_verification: bool,
}
impl Default for ThreatModelConfig {
fn default() -> Self {
Self {
strict_mode: true,
re_parse_verification: true,
tamper_detection_enabled: true,
dsl_sandboxing_enabled: true,
isolation_verification: true,
}
}
}
pub type ThreatResult<T> = Result<T, ThreatModelError>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ThreatModelError {
ReParseVerificationFailed(String),
TamperDetected(String),
SandboxEscapeDetected(String),
MutationUncertaintyDetected(String),
IsolationViolationDetected(String),
Custom(String),
}
impl std::fmt::Display for ThreatModelError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ReParseVerificationFailed(msg) => {
write!(f, "re-parse verification failed: {}", msg)
}
Self::TamperDetected(msg) => write!(f, "macro tampering detected: {}", msg),
Self::SandboxEscapeDetected(msg) => write!(f, "DSL sandbox escape: {}", msg),
Self::MutationUncertaintyDetected(msg) => {
write!(f, "mutation uncertainty in strict mode: {}", msg)
}
Self::IsolationViolationDetected(msg) => {
write!(f, "simulation isolation violation: {}", msg)
}
Self::Custom(msg) => write!(f, "{}", msg),
}
}
}
pub struct InjectionVerifier;
impl InjectionVerifier {
pub fn verify_coverage(generated_code: &str, expected_checks: &[String]) -> ThreatResult<()> {
for check in expected_checks {
if !generated_code.contains(&format!("// Invariant: {}", check)) {
return Err(ThreatModelError::ReParseVerificationFailed(format!(
"invariant check not found in generated code: {}",
check
)));
}
}
Ok(())
}
pub fn verify_scope_containment(generated_code: &str) -> ThreatResult<()> {
let dangerous_patterns = [
"unsafe",
"extern",
"use std::process",
"std::fs",
"std::net",
];
for pattern in &dangerous_patterns {
if generated_code.contains(pattern) {
return Err(ThreatModelError::ReParseVerificationFailed(format!(
"dangerous pattern found in generated code: {}",
pattern
)));
}
}
Ok(())
}
}
pub struct TamperDetector;
impl TamperDetector {
pub fn compute_hash(checks: &[String]) -> String {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
let mut sorted_checks = checks.to_vec();
sorted_checks.sort();
for check in sorted_checks {
check.hash(&mut hasher);
}
format!("{:016x}", hasher.finish())
}
pub fn verify_tampering(generated_code: &str, expected_checks: &[String]) -> ThreatResult<()> {
let expected_hash = Self::compute_hash(expected_checks);
let hash_pattern = format!("SENTRI_HASH: {}", expected_hash);
if !generated_code.contains(&hash_pattern) {
return Err(ThreatModelError::TamperDetected(
"hash mismatch: generated code does not contain expected SENTRI_HASH".to_string(),
));
}
Ok(())
}
}
pub struct DSLSandbox;
impl DSLSandbox {
pub fn validate_expression(expr: &Expression) -> ThreatResult<()> {
let forbidden_prefixes = ["file_", "io_", "extern_", "unsafe_"];
Self::check_expression_recursive(expr, &forbidden_prefixes)
}
fn check_expression_recursive(
expr: &Expression,
forbidden_prefixes: &[&str],
) -> ThreatResult<()> {
match expr {
Expression::Var(name) => {
for prefix in forbidden_prefixes {
if name.to_lowercase().starts_with(prefix) {
return Err(ThreatModelError::SandboxEscapeDetected(format!(
"forbidden variable name: {}",
name
)));
}
}
Ok(())
}
Expression::LayerVar { layer, var } => {
for prefix in forbidden_prefixes {
if layer.to_lowercase().starts_with(prefix)
|| var.to_lowercase().starts_with(prefix)
{
return Err(ThreatModelError::SandboxEscapeDetected(format!(
"forbidden layer/variable name: {}::{}",
layer, var
)));
}
}
Ok(())
}
Expression::FunctionCall { name, args } => {
let allowed_functions = [
"sum", "len", "min", "max", "abs", "mod", "div", "add", "sub", "mul", "and",
"or", "not",
];
if !allowed_functions.contains(&name.as_str()) {
return Err(ThreatModelError::SandboxEscapeDetected(format!(
"forbidden function call: {}",
name
)));
}
for arg in args {
Self::check_expression_recursive(arg, forbidden_prefixes)?;
}
Ok(())
}
Expression::BinaryOp { left, op: _, right } => {
Self::check_expression_recursive(left, forbidden_prefixes)?;
Self::check_expression_recursive(right, forbidden_prefixes)?;
Ok(())
}
Expression::Logical { left, op: _, right } => {
Self::check_expression_recursive(left, forbidden_prefixes)?;
Self::check_expression_recursive(right, forbidden_prefixes)?;
Ok(())
}
Expression::Not(inner) => {
Self::check_expression_recursive(inner, forbidden_prefixes)?;
Ok(())
}
Expression::Tuple(exprs) => {
for e in exprs {
Self::check_expression_recursive(e, forbidden_prefixes)?;
}
Ok(())
}
Expression::PhaseQualifiedVar { phase, layer, var } => {
for prefix in forbidden_prefixes {
if phase.to_lowercase().starts_with(prefix)
|| layer.to_lowercase().starts_with(prefix)
|| var.to_lowercase().starts_with(prefix)
{
return Err(ThreatModelError::SandboxEscapeDetected(format!(
"forbidden phase/layer/variable name: {}::{}::{}",
phase, layer, var
)));
}
}
Ok(())
}
Expression::PhaseConstraint {
phase: _,
constraint,
} => {
Self::check_expression_recursive(constraint, forbidden_prefixes)
}
Expression::CrossPhaseRelation {
phase1: _,
expr1,
phase2: _,
expr2,
op: _,
} => {
Self::check_expression_recursive(expr1, forbidden_prefixes)?;
Self::check_expression_recursive(expr2, forbidden_prefixes)?;
Ok(())
}
Expression::Boolean(_) | Expression::Int(_) => Ok(()),
}
}
}
pub struct StrictModeAnalyzer {
enabled: bool,
}
impl StrictModeAnalyzer {
pub fn new(enabled: bool) -> Self {
Self { enabled }
}
pub fn verify_mutation_coverage(
&self,
_analyzed_mutations: &[String],
uncertainty_warnings: &[String],
) -> ThreatResult<()> {
if !self.enabled {
return Ok(());
}
if !uncertainty_warnings.is_empty() {
return Err(ThreatModelError::MutationUncertaintyDetected(format!(
"strict mode detected {} uncertain mutations: {}",
uncertainty_warnings.len(),
uncertainty_warnings.join(", ")
)));
}
Ok(())
}
}
pub struct SimulationIsolation;
impl SimulationIsolation {
pub fn verify_isolation(
context_vars: &BTreeMap<String, String>,
allowed_types: &[&str],
) -> ThreatResult<()> {
for (name, type_str) in context_vars {
let is_allowed = allowed_types
.iter()
.any(|&allowed| type_str.contains(allowed));
if !is_allowed {
return Err(ThreatModelError::IsolationViolationDetected(format!(
"variable '{}' has disallowed type '{}' in simulation context",
name, type_str
)));
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_injection_verification() {
let generated_code = r#"
fn transfer(from: &mut Account, to: &mut Account, amount: u64) {
from.balance -= amount;
to.balance += amount;
// Invariant: balance >= 0
// SENTRI_HASH: abcd1234
}
"#;
let checks = vec!["balance >= 0".to_string()];
assert!(InjectionVerifier::verify_coverage(generated_code, &checks).is_ok());
}
#[test]
fn test_injection_verification_missing_check() {
let generated_code = "fn transfer() { /* no checks */ }";
let checks = vec!["balance >= 0".to_string()];
assert!(InjectionVerifier::verify_coverage(generated_code, &checks).is_err());
}
#[test]
fn test_scope_containment() {
let safe_code = "let x = a + b; assert!(x > 0);";
assert!(InjectionVerifier::verify_scope_containment(safe_code).is_ok());
let unsafe_code = "use std::fs; fs::write(\"file.txt\", \"\");";
assert!(InjectionVerifier::verify_scope_containment(unsafe_code).is_err());
}
#[test]
fn test_tamper_hash_deterministic() {
let checks1 = vec!["a".to_string(), "b".to_string()];
let checks2 = vec!["b".to_string(), "a".to_string()];
let hash1 = TamperDetector::compute_hash(&checks1);
let hash2 = TamperDetector::compute_hash(&checks2);
assert_eq!(hash1, hash2);
}
#[test]
fn test_dsl_sandbox_forbidden_variable() {
let expr = Expression::Var("file_handle".to_string());
assert!(DSLSandbox::validate_expression(&expr).is_err());
}
#[test]
fn test_dsl_sandbox_allowed_variable() {
let expr = Expression::Var("balance".to_string());
assert!(DSLSandbox::validate_expression(&expr).is_ok());
}
#[test]
fn test_dsl_sandbox_forbidden_function() {
let expr = Expression::FunctionCall {
name: "system_call".to_string(),
args: vec![],
};
assert!(DSLSandbox::validate_expression(&expr).is_err());
}
#[test]
fn test_dsl_sandbox_allowed_function() {
let expr = Expression::FunctionCall {
name: "sum".to_string(),
args: vec![Expression::Var("balances".to_string())],
};
assert!(DSLSandbox::validate_expression(&expr).is_ok());
}
#[test]
fn test_strict_mode_with_uncertainty() {
let analyzer = StrictModeAnalyzer::new(true);
let mutations = vec!["balance -= amount".to_string()];
let warnings = vec!["mutation from function pointer call (uncertain)".to_string()];
assert!(analyzer
.verify_mutation_coverage(&mutations, &warnings)
.is_err());
}
#[test]
fn test_strict_mode_disabled() {
let analyzer = StrictModeAnalyzer::new(false);
let mutations = vec!["balance -= amount".to_string()];
let warnings = vec!["mutation from function pointer call (uncertain)".to_string()];
assert!(analyzer
.verify_mutation_coverage(&mutations, &warnings)
.is_ok());
}
}