use anyhow::{bail, Result};
use parking_lot::RwLock;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::{error, info, warn};
pub struct CommandRateLimiter {
limits: Arc<RwLock<RateLimitState>>,
}
#[derive(Debug)]
struct RateLimitState {
command_counts: std::collections::HashMap<String, WindowedCounter>,
global_counter: WindowedCounter,
}
#[derive(Debug)]
struct WindowedCounter {
count: u64,
window_start: Instant,
window_duration: Duration,
max_count: u64,
}
impl WindowedCounter {
fn new(max_count: u64, window_duration: Duration) -> Self {
Self {
count: 0,
window_start: Instant::now(),
window_duration,
max_count,
}
}
fn check_and_increment(&mut self) -> Result<()> {
if self.window_start.elapsed() > self.window_duration {
self.count = 0;
self.window_start = Instant::now();
}
if self.count >= self.max_count {
bail!(
"Rate limit exceeded: {} requests per {:?}",
self.max_count,
self.window_duration
);
}
self.count += 1;
Ok(())
}
}
impl Default for CommandRateLimiter {
fn default() -> Self {
Self::new()
}
}
impl CommandRateLimiter {
pub fn new() -> Self {
Self {
limits: Arc::new(RwLock::new(RateLimitState {
command_counts: std::collections::HashMap::new(),
global_counter: WindowedCounter::new(100, Duration::from_secs(60)),
})),
}
}
pub fn check_command(&self, command: &str) -> Result<()> {
let mut state = self.limits.write();
state.global_counter.check_and_increment()?;
let limit = match command {
"scan" => (10, Duration::from_secs(60)), "dashboard" => (5, Duration::from_secs(300)), "status" => (60, Duration::from_secs(60)), _ => (30, Duration::from_secs(60)), };
let counter = state
.command_counts
.entry(command.to_string())
.or_insert_with(|| WindowedCounter::new(limit.0, limit.1));
counter.check_and_increment()
}
}
pub struct ResourceMonitor {
max_memory_mb: usize,
#[allow(dead_code)] max_cpu_percent: f32,
}
impl Default for ResourceMonitor {
fn default() -> Self {
Self::new()
}
}
impl ResourceMonitor {
pub const fn new() -> Self {
Self {
max_memory_mb: 512, max_cpu_percent: 80.0, }
}
pub fn check_resources(&self) -> Result<()> {
let memory_usage = self.get_memory_usage_mb();
if memory_usage > self.max_memory_mb {
bail!(
"Memory usage too high: {}MB (max: {}MB)",
memory_usage,
self.max_memory_mb
);
}
if memory_usage > self.max_memory_mb * 80 / 100 {
warn!("Memory usage approaching limit: {}MB", memory_usage);
}
Ok(())
}
const fn get_memory_usage_mb(&self) -> usize {
50 }
}
#[derive(Clone)]
pub struct SecurityContext {
pub user_id: Option<String>,
pub source: CommandSource,
pub timestamp: chrono::DateTime<chrono::Utc>,
pub request_id: String,
pub neutralization: NeutralizationContext,
}
impl SecurityContext {
pub fn new(source: CommandSource) -> Self {
Self {
user_id: None,
source,
timestamp: chrono::Utc::now(),
request_id: uuid::Uuid::new_v4().to_string(),
neutralization: NeutralizationContext::default(),
}
}
pub fn with_user(mut self, user_id: String) -> Self {
self.user_id = Some(user_id);
self
}
pub const fn with_neutralization_mode(mut self, mode: NeutralizationMode) -> Self {
self.neutralization.mode = mode;
self
}
pub const fn with_enhanced_mode(mut self, enhanced: bool) -> Self {
self.neutralization.enhanced_mode = enhanced;
self
}
pub fn record_neutralization(&mut self, success: bool) {
if success {
self.neutralization.record_success();
} else {
self.neutralization.record_failure();
}
}
pub const fn should_neutralize(&self) -> bool {
match self.neutralization.mode {
NeutralizationMode::Automatic => true,
NeutralizationMode::Interactive => self.neutralization.auto_neutralize,
NeutralizationMode::ReportOnly => false,
}
}
}
#[derive(Debug, Clone)]
pub enum CommandSource {
Cli,
WebDashboard,
Api,
Unknown,
}
#[derive(Debug, Clone)]
pub struct NeutralizationContext {
pub threats_neutralized: u32,
pub neutralization_failures: u32,
pub auto_neutralize: bool,
pub mode: NeutralizationMode,
pub enhanced_mode: bool,
pub last_neutralization: Option<chrono::DateTime<chrono::Utc>>,
}
impl Default for NeutralizationContext {
fn default() -> Self {
Self {
threats_neutralized: 0,
neutralization_failures: 0,
auto_neutralize: false,
mode: NeutralizationMode::ReportOnly,
enhanced_mode: false,
last_neutralization: None,
}
}
}
impl NeutralizationContext {
pub fn record_success(&mut self) {
self.threats_neutralized += 1;
self.last_neutralization = Some(chrono::Utc::now());
}
pub fn record_failure(&mut self) {
self.neutralization_failures += 1;
}
pub fn success_rate(&self) -> f64 {
let total = self.threats_neutralized + self.neutralization_failures;
if total == 0 {
1.0
} else {
f64::from(self.threats_neutralized) / f64::from(total)
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NeutralizationMode {
ReportOnly,
Interactive,
Automatic,
}
pub struct SecurityAuditLogger {
log_path: Option<std::path::PathBuf>,
}
impl SecurityAuditLogger {
pub const fn new(log_path: Option<std::path::PathBuf>) -> Self {
Self { log_path }
}
pub fn log_command(
&self,
context: &SecurityContext,
command: &str,
args: &serde_json::Value,
result: &Result<()>,
) {
let event = serde_json::json!({
"timestamp": context.timestamp,
"request_id": context.request_id,
"user_id": context.user_id,
"source": format!("{:?}", context.source),
"command": command,
"args": args,
"success": result.is_ok(),
"error": result.as_ref().err().map(std::string::ToString::to_string),
});
if let Ok(()) = result {
info!(event = %event, "Command executed")
} else {
warn!(event = %event, "Command failed")
}
if let Some(ref path) = self.log_path {
if let Err(e) = self.write_to_file(path, &event) {
error!("Failed to write audit log: {}", e);
}
}
}
fn write_to_file(&self, path: &std::path::Path, event: &serde_json::Value) -> Result<()> {
use std::fs::OpenOptions;
use std::io::Write;
let mut file = OpenOptions::new().create(true).append(true).open(path)?;
writeln!(file, "{event}")?;
Ok(())
}
}
pub struct FileSandbox {
allowed_paths: Vec<std::path::PathBuf>,
}
impl FileSandbox {
pub const fn new(allowed_paths: Vec<std::path::PathBuf>) -> Self {
Self { allowed_paths }
}
pub fn check_path(&self, path: &std::path::Path) -> Result<()> {
let canonical = path
.canonicalize()
.map_err(|e| anyhow::anyhow!("Invalid path: {}", e))?;
for allowed in &self.allowed_paths {
if canonical.starts_with(allowed) {
return Ok(());
}
}
bail!(
"Access denied: path '{}' is outside allowed directories",
path.display()
);
}
}
pub mod injection {
use super::{bail, Result};
use regex::Regex;
static DANGEROUS_PATTERNS: std::sync::LazyLock<Vec<Regex>> = std::sync::LazyLock::new(|| {
vec![
Regex::new(r"[;&|]").unwrap(), Regex::new(r"\$\(.*\)").unwrap(), Regex::new(r"`.*`").unwrap(), Regex::new(r"<<.*>>").unwrap(), Regex::new(r"[<>]").unwrap(), ]
});
pub fn check_command_injection(input: &str) -> Result<()> {
for pattern in DANGEROUS_PATTERNS.iter() {
if pattern.is_match(input) {
bail!("Potential command injection detected");
}
}
Ok(())
}
}
pub mod info_disclosure {
pub fn sanitize_error(error: anyhow::Error) -> String {
if cfg!(debug_assertions) {
format!("{error:#}")
} else {
match error.to_string().to_lowercase() {
s if s.contains("permission") => "Permission denied".to_string(),
s if s.contains("not found") => "Resource not found".to_string(),
s if s.contains("timeout") => "Operation timed out".to_string(),
s if s.contains("rate limit") => "Rate limit exceeded".to_string(),
_ => "An error occurred. Please try again.".to_string(),
}
}
}
pub fn mask_sensitive(key: &str, value: &str) -> String {
let sensitive_keys = ["password", "token", "secret", "key", "auth"];
if sensitive_keys
.iter()
.any(|&k| key.to_lowercase().contains(k))
{
"***MASKED***".to_string()
} else {
value.to_string()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_rate_limiter() {
let limiter = CommandRateLimiter::new();
for _ in 0..5 {
assert!(limiter.check_command("scan").is_ok());
}
let mut hit_limit = false;
for _ in 0..20 {
if limiter.check_command("scan").is_err() {
hit_limit = true;
break;
}
}
assert!(hit_limit);
}
#[test]
fn test_command_injection_detection() {
use injection::check_command_injection;
assert!(check_command_injection("normal text").is_ok());
assert!(check_command_injection("/path/to/file.txt").is_ok());
assert!(check_command_injection("test; rm -rf /").is_err());
assert!(check_command_injection("$(cat /etc/passwd)").is_err());
assert!(check_command_injection("`whoami`").is_err());
assert!(check_command_injection("test > /dev/null").is_err());
}
#[test]
fn test_file_sandbox() {
use tempfile::tempdir;
let temp_dir = tempdir().unwrap();
let allowed_path = temp_dir.path().to_path_buf();
let test_file = allowed_path.join("test.txt");
std::fs::write(&test_file, "test").unwrap();
let sandbox = FileSandbox::new(vec![allowed_path.clone()]);
assert!(sandbox.check_path(&test_file).is_ok());
let outside_path = std::env::temp_dir().join("outside.txt");
std::fs::write(&outside_path, "test").unwrap();
if !outside_path.starts_with(&allowed_path) {
assert!(sandbox.check_path(&outside_path).is_err());
}
let _ = std::fs::remove_file(outside_path);
}
#[test]
fn test_info_disclosure_prevention() {
use info_disclosure::{mask_sensitive, sanitize_error};
let error = anyhow::anyhow!("Connection to database at 192.168.1.1:5432 failed");
let sanitized = sanitize_error(error);
if cfg!(debug_assertions) {
assert!(sanitized.contains("database"));
} else {
assert_eq!(sanitized, "An error occurred. Please try again.");
}
let perm_error = anyhow::anyhow!("Permission denied for user");
assert!(sanitize_error(perm_error).contains("Permission"));
assert_eq!(mask_sensitive("password", "secret123"), "***MASKED***");
assert_eq!(mask_sensitive("api_token", "xyz"), "***MASKED***");
assert_eq!(mask_sensitive("username", "john"), "john");
}
}