fire_scope/
process.rs

1use crate::common::OutputFormat;
2use crate::output::write_ip_lists_to_files; // これが async に
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
10/// 指定された国コードと、ダウンロード済みのRIRファイル文字列から
11/// IPアドレスをパースしてファイル書き込みまで実行する
12pub 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    // 結果をファイルに書き出す
20    write_ip_lists_to_files(country_code, &ipv4_set, &ipv6_set, mode, output_format).await?;
21    Ok(())
22}
23
24/// 全RIRテキストから、指定国コードに合致するIPアドレスをすべて集約し、BTreeSetとして返す
25pub 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
54/// パース済みの国コードマップから特定の国コードのIPアドレスを取得し、ファイルに書き出す
55pub 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    // マップから該当する国コードのIPアドレスを取得
64    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    // ソートとBTreeSetへの変換
76    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    // 非同期ファイル出力
85    write_ip_lists_to_files(&upper_code, &ipv4_set, &ipv6_set, mode, output_format).await?;
86
87    Ok(())
88}
89
90/// 全RIRテキストから全ての国コードに対するIPアドレスをパースし、
91/// 指定された国コードのリストに対応するIPアドレスをファイルに書き出す
92pub 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    // 全ての国コードに対するIPアドレスを一度にパース
99    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    // 各国コードに対して非同期タスクを起動
110    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    // すべてのタスクが完了するのを待つ
130    for t in tasks {
131        let _ = t.await?;
132    }
133
134    Ok(())
135}