fire_scope/
output_common.rs1use crate::error::AppError;
2use ipnet::IpNet;
3use std::{collections::BTreeSet, path::{Path, PathBuf}};
4use tokio::fs::{self, OpenOptions};
5use tokio::io::AsyncWriteExt;
6
7pub fn sanitize_identifier(input: &str) -> String {
13 let mut s = String::with_capacity(input.len());
14 for ch in input.chars() {
15 if ch.is_ascii_alphanumeric() {
16 s.push(ch);
17 } else {
18 s.push('_');
19 }
20 }
21
22 let s = s.trim_matches('_').to_string();
24 let s = if s.len() > 64 { s[..64].to_string() } else { s };
25 if s.is_empty() { "UNKNOWN".to_string() } else { s }
26}
27
28pub fn make_header(now_str: &str, country_code: &str, as_number: &str) -> String {
30 format!(
31 "# Generated at: {}\n# Country Code: {}\n# AS Number: {}\n\n",
32 now_str, country_code, as_number
33 )
34}
35
36pub async fn write_list_txt<P: AsRef<Path>>(
37 path: P,
38 ipnets: &BTreeSet<IpNet>,
39 header: &str,
40) -> Result<(), AppError> {
41 let body = ipnets
42 .iter()
43 .map(|net| net.to_string())
44 .collect::<Vec<_>>()
45 .join("\n");
46
47 let content = format!("{}{}\n", header, body);
48
49 atomic_write(path.as_ref(), content.as_bytes()).await?;
51
52 Ok(())
53}
54
55pub async fn write_list_nft<P: AsRef<Path>>(
56 path: P,
57 ipnets: &BTreeSet<IpNet>,
58 header: &str,
59) -> Result<(), AppError> {
60 let file_path = path.as_ref();
61 let define_name_raw = file_path
62 .file_stem()
63 .and_then(|os| os.to_str())
64 .unwrap_or("unknown_define");
65 let define_name = sanitize_identifier(define_name_raw);
66
67 let mut content = String::new();
68 content.push_str(header);
69 content.push_str(&format!("define {} = {{\n", define_name));
70
71 if !ipnets.is_empty() {
72 let lines: Vec<String> = ipnets.iter().map(|n| format!(" {}", n)).collect();
73 content.push_str(&lines.join(",\n"));
74 content.push('\n');
75 }
76
77 content.push_str("}\n");
78
79 atomic_write(file_path, content.as_bytes()).await?;
81
82 Ok(())
83}
84
85async fn atomic_write(path: &Path, content: &[u8]) -> Result<(), AppError> {
87 let dir = path.parent().unwrap_or_else(|| Path::new("."));
88 let mut tmp_path = PathBuf::from(dir);
89 let fname = path
90 .file_name()
91 .and_then(|s| s.to_str())
92 .unwrap_or("output");
93 let suffix: u64 = rand::random();
94 tmp_path.push(format!(".{}.tmp.{}", fname, suffix));
95
96 {
98 let mut file = OpenOptions::new()
99 .create_new(true)
100 .write(true)
101 .open(&tmp_path)
102 .await?;
103 file.write_all(content).await?;
104 file.sync_all().await?;
106 }
107
108 match fs::rename(&tmp_path, path).await {
110 Ok(_) => Ok(()),
111 Err(_) => {
112 let _ = fs::remove_file(path).await;
113 fs::rename(&tmp_path, path).await?;
114 Ok(())
115 }
116 }
117}