use std::collections::HashSet;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ValidationError {
     UnbalancedBraces { open: usize, close: usize },
     UnbalancedParentheses { open: usize, close: usize },
     MissingExternBlock,
     MissingUnsafeBlock,
     InvalidFunctionSignature(String),
     MissingRequiredImport(String),
     DuplicateExternDeclaration(String),
 }
 
 impl std::fmt::Display for ValidationError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             ValidationError::UnbalancedBraces { open, close } => {
                 write!(f, "Unbalanced braces: {} open, {} close", open, close)
             }
             ValidationError::UnbalancedParentheses { open, close } => {
                 write!(f, "Unbalanced parentheses: {} open, {} close", open, close)
             }
             ValidationError::MissingExternBlock => {
                 write!(f, "Missing extern block for Windows API functions")
             }
             ValidationError::MissingUnsafeBlock => {
                 write!(f, "Windows API calls must be in unsafe block")
             }
             ValidationError::InvalidFunctionSignature(sig) => {
                 write!(f, "Invalid function signature: {}", sig)
             }
             ValidationError::MissingRequiredImport(import) => {
                 write!(f, "Missing required import: {}", import)
             }
             ValidationError::DuplicateExternDeclaration(name) => {
                 write!(f, "Duplicate extern declaration: {}", name)
             }
         }
     }
 }
 
 pub struct CodeValidator;
 
 impl CodeValidator {
     pub fn validate_technique_code(code: &str) -> Result<(), Vec<ValidationError>> {
         let mut errors = Vec::new();
 
         if let Err(e) = Self::check_balanced_braces(code) {
             errors.push(e);
         }
 
         if let Err(e) = Self::check_balanced_parentheses(code) {
             errors.push(e);
         }
 
         if let Some(e) = Self::check_unsafe_usage(code) {
             errors.push(e);
         }
 
         if let Some(e) = Self::check_duplicate_externs(code) {
             errors.push(e);
         }
 
         if errors.is_empty() {
             Ok(())
         } else {
             Err(errors)
         }
     }
 
     fn check_balanced_braces(code: &str) -> Result<(), ValidationError> {
         let open = code.matches('{').count();
         let close = code.matches('}').count();
         if open != close {
             Err(ValidationError::UnbalancedBraces { open, close })
         } else {
             Ok(())
         }
     }
 
     fn check_balanced_parentheses(code: &str) -> Result<(), ValidationError> {
         let open = code.matches('(').count();
         let close = code.matches(')').count();
         if open != close {
             Err(ValidationError::UnbalancedParentheses { open, close })
         } else {
             Ok(())
         }
     }
 
     fn check_unsafe_usage(code: &str) -> Option<ValidationError> {
         let has_windows_api_calls = code.contains("VirtualAlloc")
             || code.contains("VirtualProtect")
             || code.contains("CreateProcess")
             || code.contains("WriteProcessMemory")
             || code.contains("QueueUserAPC");
 
         let has_unsafe_block = code.contains("unsafe");
 
         if has_windows_api_calls && !has_unsafe_block {
             Some(ValidationError::MissingUnsafeBlock)
         } else {
             None
         }
     }
 
     fn check_duplicate_externs(code: &str) -> Option<ValidationError> {
         let mut extern_functions = HashSet::new();
         let lines: Vec<&str> = code.lines().collect();
         let mut in_extern_block = false;
 
         for line in lines {
             let trimmed = line.trim();
 
             if trimmed.starts_with("extern") && trimmed.contains('{') {
                 in_extern_block = true;
                 continue;
             }
 
             if in_extern_block && trimmed == "}" {
                 in_extern_block = false;
                 continue;
             }
 
             if in_extern_block && trimmed.starts_with("fn ") {
                 let fn_name = trimmed
                     .trim_start_matches("fn ")
                     .split('(')
                     .next()
                     .unwrap_or("");
                 if !fn_name.is_empty() && !extern_functions.insert(fn_name.to_string()) {
                     return Some(ValidationError::DuplicateExternDeclaration(
                         fn_name.to_string(),
                     ));
                 }
             }
         }
         None
     }
 
     pub fn validate_with_context(
         code: &str,
         context: &str,
     ) -> Result<(), Vec<(String, ValidationError)>> {
         Self::validate_technique_code(code).map_err(|errors| {
             errors
                 .into_iter()
                 .map(|e| (context.to_string(), e))
                 .collect()
         })
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn test_valid_code_passes() {
         let code = r#"
             unsafe {
                 let ptr = VirtualAlloc(ptr::null_mut(), 1024, MEM_COMMIT, PAGE_READWRITE);
             }
         "#;
         assert!(CodeValidator::validate_technique_code(code).is_ok());
     }
 
     #[test]
     fn test_unbalanced_braces_detected() {
         let code = r#"
             unsafe {
                 let ptr = VirtualAlloc(ptr::null_mut(), 1024, MEM_COMMIT, PAGE_READWRITE);
         "#;
         let result = CodeValidator::validate_technique_code(code);
         assert!(result.is_err());
         let errors = result.unwrap_err();
         assert!(matches!(
             errors[0],
             ValidationError::UnbalancedBraces { .. }
         ));
     }
 
     #[test]
     fn test_unbalanced_parentheses_detected() {
         let code = r#"
             unsafe {
                 let ptr = VirtualAlloc(ptr::null_mut(), 1024, MEM_COMMIT, PAGE_READWRITE;
             }
         "#;
         let result = CodeValidator::validate_technique_code(code);
         assert!(result.is_err());
         let errors = result.unwrap_err();
         assert!(matches!(
             errors[0],
             ValidationError::UnbalancedParentheses { .. }
         ));
     }
 
     #[test]
     fn test_missing_unsafe_block_detected() {
         let code = r#"
             {
                 let ptr = VirtualAlloc(ptr::null_mut(), 1024, MEM_COMMIT, PAGE_READWRITE);
             }
         "#;
         let result = CodeValidator::validate_technique_code(code);
         assert!(result.is_err());
         let errors = result.unwrap_err();
         assert!(matches!(errors[0], ValidationError::MissingUnsafeBlock));
     }
 
     #[test]
     fn test_duplicate_extern_detected() {
         let code = r#"
             extern "system" {
                 fn VirtualAlloc(addr: *mut c_void) -> *mut c_void;
                 fn VirtualAlloc(addr: *mut c_void) -> *mut c_void;
             }
         "#;
         let result = CodeValidator::validate_technique_code(code);
         assert!(result.is_err());
         let errors = result.unwrap_err();
         eprintln!("Errors found: {:?}", errors);
         assert!(errors
             .iter()
             .any(|e| matches!(e, ValidationError::DuplicateExternDeclaration(_))));
     }
 
     #[test]
     fn test_code_without_api_calls_doesnt_need_unsafe() {
         let code = r#"
             {
                 let x = 42;
                 let y = x + 1;
             }
         "#;
         assert!(CodeValidator::validate_technique_code(code).is_ok());
     }
 
     #[test]
     fn test_validation_with_context() {
         let code = r#"
             {
                 let ptr = VirtualAlloc(ptr::null_mut(), 1024, MEM_COMMIT, PAGE_READWRITE);
             }
         "#;
         let result = CodeValidator::validate_with_context(code, "CreateFiberGenerator");
         assert!(result.is_err());
         let errors = result.unwrap_err();
         assert_eq!(errors[0].0, "CreateFiberGenerator");
     }
 }

© 2024 - 2026 djnn, inc • generated with stagit