use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::fs::File;
use std::io::{prelude::*, BufReader};
use std::path::Path;
use tracing::{debug, info};
use crate::logger;
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Severity {
Critical,
High,
Medium,
Low,
Info,
}
impl Severity {
pub fn to_value(&self) -> u8 {
match self {
Severity::Critical => 5,
Severity::High => 4,
Severity::Medium => 3,
Severity::Low => 2,
Severity::Info => 1,
}
}
}
impl PartialOrd for Severity {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.to_value().cmp(&other.to_value()))
}
}
impl Ord for Severity {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.to_value().cmp(&other.to_value())
}
}
impl std::fmt::Display for Severity {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Severity::Critical => write!(f, "critical"),
Severity::High => write!(f, "high"),
Severity::Medium => write!(f, "medium"),
Severity::Low => write!(f, "low"),
Severity::Info => write!(f, "info"),
}
}
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Rule {
pub name: String,
pub path: String,
pub signature: String,
#[serde(default)]
pub description: Option<String>,
#[serde(default)]
pub severity: Option<Severity>,
}
impl Rule {
#[allow(dead_code)]
pub fn new(
name: &str,
path: &str,
signature: &str,
description: &str,
severity: Severity,
) -> Self {
Self {
name: name.to_string(),
path: path.to_string(),
signature: signature.to_string(),
description: Some(description.to_string()),
severity: Some(severity),
}
}
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct RuleSet {
pub rules: Vec<Rule>,
}
impl RuleSet {
pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
let file = File::open(path.as_ref()).context(format!(
"Failed to open rules file: {}",
path.as_ref().display()
))?;
let reader = BufReader::new(file);
let mut ruleset: RuleSet = serde_yaml::from_reader(reader).context(format!(
"Failed to parse rules file: {}",
path.as_ref().display()
))?;
ruleset.sort_by_severity();
info!(
"📋 Loaded {} rules from {}",
ruleset.rules.len(),
path.as_ref().display()
);
for rule in &ruleset.rules {
logger::log_rule_loaded(&rule.name, 1); }
Ok(ruleset)
}
pub fn sort_by_severity(&mut self) {
self.rules.sort_by(|a, b| {
match (&a.severity, &b.severity) {
(Some(a_sev), Some(b_sev)) => b_sev.cmp(a_sev), (Some(_), None) => Ordering::Less,
(None, Some(_)) => Ordering::Greater,
(None, None) => Ordering::Equal,
}
});
}
}
pub fn load_rules(rules_file: &str) -> Result<RuleSet> {
RuleSet::from_file(rules_file)
}
pub fn add_rule(yaml_file: &str) -> Result<()> {
let existing_rules_path = "rules.yaml";
let mut existing_ruleset = load_rules(existing_rules_path)?;
let new_ruleset = load_rules(yaml_file)?;
let original_count = existing_ruleset.rules.len();
for new_rule in new_ruleset.rules {
if !existing_ruleset
.rules
.iter()
.any(|r| r.name == new_rule.name)
{
existing_ruleset.rules.push(new_rule);
} else {
debug!("Rule '{}' already exists, skipping", new_rule.name);
}
}
let added_count = existing_ruleset.rules.len() - original_count;
let yaml =
serde_yaml::to_string(&existing_ruleset).context("Failed to serialize rules to YAML")?;
let mut file = File::create(existing_rules_path).context(format!(
"Failed to open rules file for writing: {}",
existing_rules_path
))?;
file.write_all(yaml.as_bytes()).context(format!(
"Failed to write to rules file: {}",
existing_rules_path
))?;
info!(
"✅ Added {} new rules to {}",
added_count, existing_rules_path
);
Ok(())
}
pub fn remove_rule(rule_name: &str) -> Result<()> {
let existing_rules_path = "rules.yaml";
let mut ruleset = load_rules(existing_rules_path)?;
let original_count = ruleset.rules.len();
ruleset.rules.retain(|rule| rule.name != rule_name);
if ruleset.rules.len() == original_count {
info!(
"⚠️ Rule '{}' not found in {}",
rule_name, existing_rules_path
);
return Ok(());
}
let yaml = serde_yaml::to_string(&ruleset).context("Failed to serialize rules to YAML")?;
let mut file = File::create(existing_rules_path).context(format!(
"Failed to open rules file for writing: {}",
existing_rules_path
))?;
file.write_all(yaml.as_bytes()).context(format!(
"Failed to write to rules file: {}",
existing_rules_path
))?;
info!(
"✅ Removed rule '{}' from {}",
rule_name, existing_rules_path
);
Ok(())
}
pub fn list_rules(rules_file: &str) -> Result<()> {
let ruleset = load_rules(rules_file)?;
println!("📋 Rules in {}:", rules_file);
println!("{:<30} {:<15} {:<}", "Name", "Severity", "Description");
println!("{:-<60}", "");
for rule in &ruleset.rules {
let severity = match &rule.severity {
Some(s) => s.to_string(),
None => "N/A".to_string(),
};
let description = rule.description.as_deref().unwrap_or("N/A");
println!("{:<30} {:<15} {:<}", rule.name, severity, description);
}
println!("\nTotal rules: {}", ruleset.rules.len());
Ok(())
}