1use crate::common::OutputFormat;
2use crate::output::write_ip_lists_to_files;
3use 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 std::error::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)?;
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 std::error::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 std::error::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
79 ipv4_sorted.sort();
80 ipv6_sorted.sort();
81
82 let ipv4_set = ipv4_sorted.into_iter().collect::<BTreeSet<_>>();
83 let ipv6_set = ipv6_sorted.into_iter().collect::<BTreeSet<_>>();
84
85 write_ip_lists_to_files(&upper_code, &ipv4_set, &ipv6_set, mode, output_format)?;
87
88 Ok(())
89}
90
91pub async fn process_all_country_codes(
94 country_codes: &[String],
95 rir_texts: &[String],
96 mode: &str,
97 output_format: OutputFormat,
98) -> Result<(), Box<dyn Error + Send + Sync>> {
99 let country_map = match parse_all_country_codes(rir_texts) {
101 Ok(map) => map,
102 Err(e) => {
103 eprintln!("Country code parsing failed.\n{}", e);
104 return Err(e);
105 }
106 };
107
108 let country_map = Arc::new(country_map);
110
111 let mut tasks: Vec<JoinHandle<Result<(), Box<dyn Error + Send + Sync>>>> = Vec::new();
113 for code in country_codes {
114 let code_clone = code.clone();
115 let mode_clone = mode.to_string();
116 let format_clone = output_format;
117 let map_arc = Arc::clone(&country_map);
119
120 let handle = tokio::spawn(async move {
121 if let Err(e) =
122 process_country_code_from_map(&code_clone, &map_arc, &mode_clone, format_clone)
123 .await
124 {
125 eprintln!("Error (country={}): {}", code_clone, e);
126 }
127 Ok(())
128 });
129 tasks.push(handle);
130 }
131
132 for t in tasks {
134 let _ = t.await?;
135 }
136
137 Ok(())
138}