use serde_json::Value as JsonValue;
use serde_yaml::Value as YamlValue;
use std::path::Path;
use toml::Value as TomlValue;
use crate::error::{MaskitError, Result};
pub const DEFAULT_MASK: &str = "****************";
#[derive(Debug, PartialEq)]
pub enum ConfigFormat {
Json,
Yaml,
Toml,
}
impl ConfigFormat {
pub fn from_path(path: &Path) -> Result<Self> {
let extension = path
.extension()
.and_then(|ext| ext.to_str())
.ok_or_else(|| MaskitError::InvalidPath("No file extension found".to_string()))?;
match extension.to_lowercase().as_str() {
"json" => Ok(ConfigFormat::Json),
"yaml" | "yml" => Ok(ConfigFormat::Yaml),
"toml" => Ok(ConfigFormat::Toml),
ext => Err(MaskitError::UnsupportedFormat(ext.to_string())),
}
}
}
fn should_mask_field(field_name: &str, keywords: &[String]) -> bool {
let field_lower = field_name.to_lowercase();
keywords
.iter()
.any(|keyword| field_lower.contains(&keyword.to_lowercase()))
}
fn mask_json_value(value: &mut JsonValue, keywords: &[String]) {
match value {
JsonValue::Object(map) => {
for (key, val) in map.iter_mut() {
if should_mask_field(key, keywords) {
*val = JsonValue::String(DEFAULT_MASK.to_string());
} else {
mask_json_value(val, keywords);
}
}
}
JsonValue::Array(arr) => {
for item in arr.iter_mut() {
mask_json_value(item, keywords);
}
}
_ => {}
}
}
fn mask_yaml_value(value: &mut YamlValue, keywords: &[String]) {
match value {
YamlValue::Mapping(map) => {
let keys: Vec<_> = map.keys().cloned().collect();
for key in keys {
if let YamlValue::String(key_str) = &key {
if should_mask_field(key_str, keywords) {
map.insert(key.clone(), YamlValue::String(DEFAULT_MASK.to_string()));
} else if let Some(val) = map.get_mut(&key) {
mask_yaml_value(val, keywords);
}
}
}
}
YamlValue::Sequence(seq) => {
for item in seq.iter_mut() {
mask_yaml_value(item, keywords);
}
}
_ => {}
}
}
fn mask_toml_value(value: &mut TomlValue, keywords: &[String]) {
match value {
TomlValue::Table(table) => {
for (key, val) in table.iter_mut() {
if should_mask_field(key, keywords) {
*val = TomlValue::String(DEFAULT_MASK.to_string());
} else {
mask_toml_value(val, keywords);
}
}
}
TomlValue::Array(arr) => {
for item in arr.iter_mut() {
mask_toml_value(item, keywords);
}
}
_ => {}
}
}
pub fn mask_config(content: &str, format: ConfigFormat, keywords: &[String]) -> Result<String> {
match format {
ConfigFormat::Json => {
let mut value: JsonValue =
serde_json::from_str(content).map_err(|e| MaskitError::ParseError {
format: "JSON".to_string(),
message: e.to_string(),
})?;
mask_json_value(&mut value, keywords);
serde_json::to_string_pretty(&value).map_err(|e| MaskitError::SerializeError {
format: "JSON".to_string(),
message: e.to_string(),
})
}
ConfigFormat::Yaml => {
let mut value: YamlValue =
serde_yaml::from_str(content).map_err(|e| MaskitError::ParseError {
format: "YAML".to_string(),
message: e.to_string(),
})?;
mask_yaml_value(&mut value, keywords);
serde_yaml::to_string(&value).map_err(|e| MaskitError::SerializeError {
format: "YAML".to_string(),
message: e.to_string(),
})
}
ConfigFormat::Toml => {
let mut value: TomlValue =
content
.parse()
.map_err(|e: toml::de::Error| MaskitError::ParseError {
format: "TOML".to_string(),
message: e.to_string(),
})?;
mask_toml_value(&mut value, keywords);
toml::to_string_pretty(&value).map_err(|e| MaskitError::SerializeError {
format: "TOML".to_string(),
message: e.to_string(),
})
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_format_detection() {
assert_eq!(
ConfigFormat::from_path(Path::new("config.json")).unwrap(),
ConfigFormat::Json
);
assert_eq!(
ConfigFormat::from_path(Path::new("config.yaml")).unwrap(),
ConfigFormat::Yaml
);
assert_eq!(
ConfigFormat::from_path(Path::new("config.yml")).unwrap(),
ConfigFormat::Yaml
);
assert_eq!(
ConfigFormat::from_path(Path::new("config.toml")).unwrap(),
ConfigFormat::Toml
);
assert!(ConfigFormat::from_path(Path::new("config.txt")).is_err());
}
#[test]
fn test_should_mask_field() {
let keywords = vec!["key".to_string(), "secret".to_string()];
assert!(should_mask_field("api_key", &keywords));
assert!(should_mask_field("API_KEY", &keywords));
assert!(should_mask_field("secret_token", &keywords));
assert!(should_mask_field("database_secret", &keywords));
assert!(!should_mask_field("username", &keywords));
assert!(!should_mask_field("base_url", &keywords));
}
#[test]
fn test_mask_json() {
let json = r#"{
"api_key": "123456",
"base_url": "http://example.com",
"secret_token": "abcdef",
"nested": {
"db_key": "xyz789",
"host": "localhost"
}
}"#;
let keywords = vec!["key".to_string(), "secret".to_string()];
let result = mask_config(json, ConfigFormat::Json, &keywords).unwrap();
let masked: JsonValue = serde_json::from_str(&result).unwrap();
let obj = masked.as_object().unwrap();
assert_eq!(obj.get("api_key").unwrap().as_str().unwrap(), DEFAULT_MASK);
assert_eq!(
obj.get("base_url").unwrap().as_str().unwrap(),
"http://example.com"
);
assert_eq!(
obj.get("secret_token").unwrap().as_str().unwrap(),
DEFAULT_MASK
);
let nested = obj.get("nested").unwrap().as_object().unwrap();
assert_eq!(
nested.get("db_key").unwrap().as_str().unwrap(),
DEFAULT_MASK
);
assert_eq!(nested.get("host").unwrap().as_str().unwrap(), "localhost");
}
#[test]
fn test_mask_yaml() {
let yaml = r#"
api_key: "123456"
base_url: "http://example.com"
secret_token: "abcdef"
nested:
db_key: "xyz789"
host: "localhost"
"#;
let keywords = vec!["key".to_string(), "secret".to_string()];
let result = mask_config(yaml, ConfigFormat::Yaml, &keywords).unwrap();
let masked: YamlValue = serde_yaml::from_str(&result).unwrap();
let map = masked.as_mapping().unwrap();
assert_eq!(
map.get(YamlValue::String("api_key".to_string()))
.unwrap()
.as_str()
.unwrap(),
DEFAULT_MASK
);
assert_eq!(
map.get(YamlValue::String("base_url".to_string()))
.unwrap()
.as_str()
.unwrap(),
"http://example.com"
);
}
#[test]
fn test_mask_toml() {
let toml = r#"
api_key = "123456"
base_url = "http://example.com"
secret_token = "abcdef"
[nested]
db_key = "xyz789"
host = "localhost"
"#;
let keywords = vec!["key".to_string(), "secret".to_string()];
let result = mask_config(toml, ConfigFormat::Toml, &keywords).unwrap();
let masked: TomlValue = result.parse().unwrap();
let table = masked.as_table().unwrap();
assert_eq!(
table.get("api_key").unwrap().as_str().unwrap(),
DEFAULT_MASK
);
assert_eq!(
table.get("base_url").unwrap().as_str().unwrap(),
"http://example.com"
);
assert_eq!(
table.get("secret_token").unwrap().as_str().unwrap(),
DEFAULT_MASK
);
let nested = table.get("nested").unwrap().as_table().unwrap();
assert_eq!(
nested.get("db_key").unwrap().as_str().unwrap(),
DEFAULT_MASK
);
assert_eq!(nested.get("host").unwrap().as_str().unwrap(), "localhost");
}
}