use crate::error::AppError;
use ipnet::IpNet;
use std::{collections::BTreeSet, path::{Path, PathBuf}};
use tokio::fs::{self, OpenOptions};
use tokio::io::AsyncWriteExt;
pub fn sanitize_identifier(input: &str) -> String {
let mut s = String::with_capacity(input.len());
for ch in input.chars() {
if ch.is_ascii_alphanumeric() {
s.push(ch);
} else {
s.push('_');
}
}
let s = s.trim_matches('_').to_string();
let s = if s.len() > 64 { s[..64].to_string() } else { s };
if s.is_empty() { "UNKNOWN".to_string() } else { s }
}
pub fn make_header(now_str: &str, country_code: &str, as_number: &str) -> String {
format!(
"# Generated at: {}\n# Country Code: {}\n# AS Number: {}\n\n",
now_str, country_code, as_number
)
}
pub async fn write_list_txt<P: AsRef<Path>>(
path: P,
ipnets: &BTreeSet<IpNet>,
header: &str,
) -> Result<(), AppError> {
let body = ipnets
.iter()
.map(|net| net.to_string())
.collect::<Vec<_>>()
.join("\n");
let content = format!("{}{}\n", header, body);
atomic_write(path.as_ref(), content.as_bytes()).await?;
Ok(())
}
pub async fn write_list_nft<P: AsRef<Path>>(
path: P,
ipnets: &BTreeSet<IpNet>,
header: &str,
) -> Result<(), AppError> {
let file_path = path.as_ref();
let define_name_raw = file_path
.file_stem()
.and_then(|os| os.to_str())
.unwrap_or("unknown_define");
let define_name = sanitize_identifier(define_name_raw);
let mut content = String::new();
content.push_str(header);
content.push_str(&format!("define {} = {{\n", define_name));
if !ipnets.is_empty() {
let lines: Vec<String> = ipnets.iter().map(|n| format!(" {}", n)).collect();
content.push_str(&lines.join(",\n"));
content.push('\n');
}
content.push_str("}\n");
atomic_write(file_path, content.as_bytes()).await?;
Ok(())
}
async fn atomic_write(path: &Path, content: &[u8]) -> Result<(), AppError> {
let dir = path.parent().unwrap_or_else(|| Path::new("."));
let mut tmp_path = PathBuf::from(dir);
let fname = path
.file_name()
.and_then(|s| s.to_str())
.unwrap_or("output");
let suffix: u64 = rand::random();
tmp_path.push(format!(".{}.tmp.{}", fname, suffix));
{
let mut file = OpenOptions::new()
.create_new(true)
.write(true)
.open(&tmp_path)
.await?;
file.write_all(content).await?;
file.sync_all().await?;
}
match fs::rename(&tmp_path, path).await {
Ok(_) => Ok(()),
Err(_) => {
let _ = fs::remove_file(path).await;
fs::rename(&tmp_path, path).await?;
Ok(())
}
}
}