1use crate::common::OutputFormat;
2use crate::output::write_ip_lists_to_files; use crate::parse::{parse_all_country_codes, parse_ip_lines};
4use ipnet::IpNet;
5use std::collections::{BTreeSet, HashMap};
6use std::error::Error;
7use std::sync::Arc;
8use tokio::task::JoinHandle;
9
10pub async fn process_country_code(
13 country_code: &str,
14 rir_texts: &[String],
15 mode: &str,
16 output_format: OutputFormat,
17) -> Result<(), Box<dyn Error + Send + Sync>> {
18 let (ipv4_set, ipv6_set) = parse_and_collect_ips(country_code, rir_texts)?;
19 write_ip_lists_to_files(country_code, &ipv4_set, &ipv6_set, mode, output_format).await?;
21 Ok(())
22}
23
24pub fn parse_and_collect_ips(
26 country_code: &str,
27 rir_texts: &[String],
28) -> Result<(BTreeSet<IpNet>, BTreeSet<IpNet>), Box<dyn Error + Send + Sync>> {
29 let mut ipv4_vec = Vec::new();
30 let mut ipv6_vec = Vec::new();
31
32 for text in rir_texts {
33 match parse_ip_lines(text, country_code) {
34 Ok((v4, v6)) => {
35 ipv4_vec.extend(v4);
36 ipv6_vec.extend(v6);
37 }
38 Err(e) => eprintln!(
39 "[parse_and_collect_ips] Error parsing for country '{}': {}",
40 country_code, e
41 ),
42 }
43 }
44
45 ipv4_vec.sort();
46 ipv6_vec.sort();
47
48 let ipv4_set = ipv4_vec.into_iter().collect::<BTreeSet<_>>();
49 let ipv6_set = ipv6_vec.into_iter().collect::<BTreeSet<_>>();
50
51 Ok((ipv4_set, ipv6_set))
52}
53
54pub async fn process_country_code_from_map(
56 country_code: &str,
57 country_map: &HashMap<String, (Vec<IpNet>, Vec<IpNet>)>,
58 mode: &str,
59 output_format: OutputFormat,
60) -> Result<(), Box<dyn Error + Send + Sync>> {
61 let upper_code = country_code.to_uppercase();
62
63 let (ipv4_vec, ipv6_vec) = match country_map.get(&upper_code) {
65 Some(ip_lists) => ip_lists,
66 None => {
67 eprintln!(
68 "No IP address corresponding to the country code could be found.\n{}",
69 upper_code
70 );
71 return Ok(());
72 }
73 };
74
75 let mut ipv4_sorted = ipv4_vec.clone();
77 let mut ipv6_sorted = ipv6_vec.clone();
78 ipv4_sorted.sort();
79 ipv6_sorted.sort();
80
81 let ipv4_set = ipv4_sorted.into_iter().collect::<BTreeSet<_>>();
82 let ipv6_set = ipv6_sorted.into_iter().collect::<BTreeSet<_>>();
83
84 write_ip_lists_to_files(&upper_code, &ipv4_set, &ipv6_set, mode, output_format).await?;
86
87 Ok(())
88}
89
90pub async fn process_all_country_codes(
93 country_codes: &[String],
94 rir_texts: &[String],
95 mode: &str,
96 output_format: OutputFormat,
97) -> Result<(), Box<dyn Error + Send + Sync>> {
98 let country_map = match parse_all_country_codes(rir_texts) {
100 Ok(map) => map,
101 Err(e) => {
102 eprintln!("Country code parsing failed.\n{}", e);
103 return Err(e);
104 }
105 };
106
107 let country_map = Arc::new(country_map);
108
109 let mut tasks: Vec<JoinHandle<Result<(), Box<dyn Error + Send + Sync>>>> = Vec::new();
111 for code in country_codes {
112 let code_clone = code.clone();
113 let mode_clone = mode.to_string();
114 let format_clone = output_format;
115 let map_arc = Arc::clone(&country_map);
116
117 let handle = tokio::spawn(async move {
118 if let Err(e) =
119 process_country_code_from_map(&code_clone, &map_arc, &mode_clone, format_clone)
120 .await
121 {
122 eprintln!("Error (country={}): {}", code_clone, e);
123 }
124 Ok(())
125 });
126 tasks.push(handle);
127 }
128
129 for t in tasks {
131 let _ = t.await?;
132 }
133
134 Ok(())
135}