use super::{CommandValidationError, ExecutionMode};
use crate::client::ApiClient;
use chrono::{Datelike, Timelike};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecurityEvent {
pub timestamp: SystemTime,
pub user: String,
pub command: String,
pub action: SecurityAction,
pub result: SecurityResult,
pub details: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SecurityAction {
CommandValidation,
KnowledgeGraphCheck,
PermissionCheck,
RateLimitCheck,
BlacklistCheck,
TimeRestrictionCheck,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SecurityResult {
Allowed,
Denied(String),
Warning(String),
}
#[derive(Debug, Clone)]
pub struct RateLimit {
pub max_requests: usize,
pub window: Duration,
pub current_requests: Vec<SystemTime>,
}
#[derive(Debug, Clone)]
pub struct TimeRestrictions {
pub allowed_hours: Vec<u8>, pub allowed_days: Vec<u8>, pub maintenance_windows: Vec<MaintenanceWindow>,
}
#[derive(Debug, Clone)]
pub struct MaintenanceWindow {
pub start_day: u8,
pub start_hour: u8,
pub duration_hours: u8,
pub reason: String,
}
pub struct CommandValidator {
api_client: Option<Arc<ApiClient>>,
role_permissions: HashMap<String, Vec<String>>,
concept_cache: HashMap<String, Vec<String>>,
audit_log: Vec<SecurityEvent>,
rate_limits: HashMap<String, RateLimit>,
blacklisted_commands: Vec<String>,
time_restrictions: TimeRestrictions,
}
impl CommandValidator {
pub fn new() -> Self {
let mut role_permissions = HashMap::new();
role_permissions.insert(
"Default".to_string(),
vec!["read".to_string(), "search".to_string(), "help".to_string()],
);
role_permissions.insert(
"Terraphim Engineer".to_string(),
vec![
"read".to_string(),
"write".to_string(),
"execute".to_string(),
"search".to_string(),
"configure".to_string(),
"help".to_string(),
],
);
let mut rate_limits = HashMap::new();
rate_limits.insert(
"search".to_string(),
RateLimit {
max_requests: 100,
window: Duration::from_secs(60),
current_requests: Vec::new(),
},
);
rate_limits.insert(
"deploy".to_string(),
RateLimit {
max_requests: 5,
window: Duration::from_secs(3600),
current_requests: Vec::new(),
},
);
let blacklisted_commands = vec![
"rm -rf /".to_string(),
"dd if=/dev/zero".to_string(),
"mkfs".to_string(),
"fdisk".to_string(),
];
let time_restrictions = if cfg!(test) {
TimeRestrictions {
allowed_hours: vec![], allowed_days: vec![], maintenance_windows: vec![],
}
} else {
TimeRestrictions {
allowed_hours: (9..=17).collect(), allowed_days: (1..=5).collect(), maintenance_windows: Vec::new(),
}
};
Self {
api_client: None,
role_permissions,
concept_cache: HashMap::new(),
audit_log: Vec::new(),
rate_limits,
blacklisted_commands,
time_restrictions,
}
}
pub fn set_rate_limit(&mut self, command: &str, max_requests: u32, window: Duration) {
self.rate_limits.insert(
command.to_string(),
RateLimit {
max_requests: max_requests as usize,
window,
current_requests: Vec::new(),
},
);
}
pub fn with_api_client(api_client: Arc<ApiClient>) -> Self {
let mut validator = Self::new();
validator.api_client = Some(api_client);
validator
}
pub async fn validate_command_execution(
&mut self,
command: &str,
role: &str,
_parameters: &HashMap<String, String>,
) -> Result<ExecutionMode, CommandValidationError> {
self.validate_command_execution_with_mode(command, role, _parameters, None)
.await
}
pub async fn validate_command_execution_with_mode(
&mut self,
command: &str,
role: &str,
_parameters: &HashMap<String, String>,
definition_execution_mode: Option<ExecutionMode>,
) -> Result<ExecutionMode, CommandValidationError> {
if let Some(permissions) = self.role_permissions.get(role) {
if !self.has_required_permissions(command, permissions) {
return Err(CommandValidationError::InsufficientPermissions(format!(
"'{}' role lacks required permissions for command",
role
)));
}
}
if self.api_client.is_some() {
let kg_concepts = self.get_knowledge_graph_concepts(role).await?;
let command_concept = self.extract_command_concept(command);
if !kg_concepts.is_empty() && !kg_concepts.contains(&command_concept) {
return Err(CommandValidationError::MissingKnowledgeGraphConcepts(
command.to_string(),
vec![command_concept],
));
}
}
let execution_mode =
self.determine_execution_mode_with_override(command, role, definition_execution_mode);
Ok(execution_mode)
}
async fn get_knowledge_graph_concepts(
&mut self,
role: &str,
) -> Result<Vec<String>, CommandValidationError> {
if let Some(cached_concepts) = self.concept_cache.get(role) {
return Ok(cached_concepts.clone());
}
if let Some(api_client) = &self.api_client {
match api_client.get_rolegraph_edges(Some(role)).await {
Ok(rolegraph_response) => {
let concepts: Vec<String> = rolegraph_response
.nodes
.into_iter()
.map(|node| node.label.to_lowercase())
.collect();
self.concept_cache
.insert(role.to_string(), concepts.clone());
Ok(concepts)
}
Err(e) => {
eprintln!("Warning: Failed to fetch rolegraph for '{}': {}", role, e);
Ok(vec![])
}
}
} else {
Ok(vec![])
}
}
fn is_write_operation(&self, command: &str) -> bool {
let write_commands = vec![
"rm", "mv", "cp", "touch", "mkdir", "rmdir", "chmod", "chown", "write", "create",
"delete", "update", "modify", "edit",
];
write_commands.iter().any(|cmd| command.contains(cmd))
}
fn extract_command_concept(&self, command: &str) -> String {
command
.split_whitespace()
.next()
.unwrap_or("unknown")
.to_lowercase()
}
#[allow(dead_code)]
fn determine_execution_mode(&self, command: &str, role: &str) -> ExecutionMode {
self.determine_execution_mode_with_override(command, role, None)
}
fn determine_execution_mode_with_override(
&self,
command: &str,
role: &str,
definition_mode: Option<ExecutionMode>,
) -> ExecutionMode {
if self.is_high_risk_command(command) {
return ExecutionMode::Firecracker;
}
if let Some(mode) = definition_mode {
return mode;
}
if role == "Terraphim Engineer" && self.is_safe_command(command) {
return ExecutionMode::Local;
}
ExecutionMode::Hybrid
}
fn is_high_risk_command(&self, command: &str) -> bool {
let high_risk_patterns = [
"rm -rf",
"dd if=",
"mkfs",
"fdisk",
"iptables",
"systemctl",
"shutdown",
"reboot",
"passwd",
"chown root",
"chmod 777",
];
high_risk_patterns
.iter()
.any(|pattern| command.contains(pattern))
}
fn is_safe_command(&self, command: &str) -> bool {
let safe_commands = [
"ls", "cat", "echo", "pwd", "date", "whoami", "grep", "find", "head", "tail", "wc",
"sort",
];
safe_commands.iter().any(|cmd| command.starts_with(cmd))
}
fn is_system_command(&self, command: &str) -> bool {
let system_commands = [
"systemctl",
"shutdown",
"reboot",
"passwd",
"chown",
"chmod",
"iptables",
"fdisk",
"mkfs",
];
system_commands.iter().any(|cmd| command.starts_with(cmd))
}
pub fn add_role_permissions(&mut self, role: String, permissions: Vec<String>) {
self.role_permissions.insert(role, permissions);
}
pub fn is_blacklisted(&self, command: &str) -> bool {
self.blacklisted_commands
.iter()
.any(|blacklisted| command.contains(blacklisted) || command.starts_with(blacklisted))
}
pub fn check_rate_limit(&mut self, command: &str) -> Result<(), CommandValidationError> {
let command_name = self.extract_command_concept(command);
if let Some(rate_limit) = self.rate_limits.get_mut(&command_name) {
let now = SystemTime::now();
rate_limit.current_requests.retain(|×tamp| {
now.duration_since(timestamp).unwrap_or(Duration::MAX) < rate_limit.window
});
if rate_limit.current_requests.len() >= rate_limit.max_requests {
return Err(CommandValidationError::ValidationFailed(format!(
"Rate limit exceeded for command '{}': {}/{} requests per {:?}",
command_name,
rate_limit.current_requests.len(),
rate_limit.max_requests,
rate_limit.window
)));
}
rate_limit.current_requests.push(now);
}
Ok(())
}
pub fn check_time_restrictions(&self) -> Result<(), CommandValidationError> {
let now = std::time::SystemTime::now();
let datetime = chrono::DateTime::<chrono::Utc>::from(now);
let local_time = datetime.with_timezone(&chrono::Local);
if !self.time_restrictions.allowed_days.is_empty()
&& !self
.time_restrictions
.allowed_days
.contains(&(local_time.weekday().num_days_from_sunday() as u8))
{
return Err(CommandValidationError::ValidationFailed(
"Commands not allowed on this day".to_string(),
));
}
if !self.time_restrictions.allowed_hours.is_empty()
&& !self
.time_restrictions
.allowed_hours
.contains(&(local_time.hour() as u8))
{
return Err(CommandValidationError::ValidationFailed(format!(
"Commands not allowed at this time: {}:00",
local_time.hour()
)));
}
Ok(())
}
pub fn log_security_event(
&mut self,
user: &str,
command: &str,
action: SecurityAction,
result: SecurityResult,
details: &str,
) {
let event = SecurityEvent {
timestamp: SystemTime::now(),
user: user.to_string(),
command: command.to_string(),
action,
result: result.clone(),
details: details.to_string(),
};
self.audit_log.push(event);
if self.audit_log.len() > 1000 {
self.audit_log.drain(0..100);
}
}
pub fn get_recent_events(&self, limit: usize) -> Vec<&SecurityEvent> {
self.audit_log.iter().rev().take(limit).collect()
}
pub async fn validate_command_security(
&mut self,
command: &str,
role: &str,
user: &str,
) -> Result<(), CommandValidationError> {
if self.is_blacklisted(command) {
let details = format!("Command '{}' is blacklisted for security reasons", command);
self.log_security_event(
user,
command,
SecurityAction::BlacklistCheck,
SecurityResult::Denied(details.clone()),
&details,
);
return Err(CommandValidationError::ValidationFailed(details));
}
self.log_security_event(
user,
command,
SecurityAction::BlacklistCheck,
SecurityResult::Allowed,
"Command not blacklisted",
);
if let Err(e) = self.check_rate_limit(command) {
self.log_security_event(
user,
command,
SecurityAction::RateLimitCheck,
SecurityResult::Denied(e.to_string()),
"Rate limit exceeded",
);
return Err(e);
}
self.log_security_event(
user,
command,
SecurityAction::RateLimitCheck,
SecurityResult::Allowed,
"Rate limit check passed",
);
if let Err(e) = self.check_time_restrictions() {
self.log_security_event(
user,
command,
SecurityAction::TimeRestrictionCheck,
SecurityResult::Denied(e.to_string()),
"Time restriction violation",
);
return Err(e);
}
self.log_security_event(
user,
command,
SecurityAction::TimeRestrictionCheck,
SecurityResult::Allowed,
"Time restrictions satisfied",
);
if let Some(permissions) = self.role_permissions.get(role) {
if !self.has_required_permissions(command, permissions) {
let details = format!("Role '{}' lacks required permissions for command", role);
self.log_security_event(
user,
command,
SecurityAction::PermissionCheck,
SecurityResult::Denied(details.clone()),
&details,
);
return Err(CommandValidationError::InsufficientPermissions(details));
}
}
self.log_security_event(
user,
command,
SecurityAction::PermissionCheck,
SecurityResult::Allowed,
"Role permissions verified",
);
Ok(())
}
fn has_required_permissions(&self, command: &str, permissions: &[String]) -> bool {
if self.is_write_operation(command) && !permissions.contains(&"write".to_string()) {
return false;
}
if self.is_high_risk_command(command) && !permissions.contains(&"execute".to_string()) {
return false;
}
if self.is_system_command(command) && !permissions.contains(&"execute".to_string()) {
return false;
}
true
}
pub fn get_security_stats(&self) -> SecurityStats {
let total_events = self.audit_log.len();
let denied_events = self
.audit_log
.iter()
.filter(|e| matches!(e.result, SecurityResult::Denied(_)))
.count();
let recent_hour = SystemTime::now()
.checked_sub(Duration::from_secs(3600))
.unwrap_or(SystemTime::UNIX_EPOCH);
let recent_events = self
.audit_log
.iter()
.filter(|e| e.timestamp > recent_hour)
.count();
SecurityStats {
total_events,
denied_events,
recent_events,
active_rate_limits: self.rate_limits.len(),
blacklisted_commands: self.blacklisted_commands.len(),
}
}
}
impl Default for CommandValidator {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecurityStats {
pub total_events: usize,
pub denied_events: usize,
pub recent_events: usize,
pub active_rate_limits: usize,
pub blacklisted_commands: usize,
}