fire_scope/
process.rs

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
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 std::error::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)?;
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 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
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 std::error::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
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    // 結果をファイルに書き出す
86    write_ip_lists_to_files(&upper_code, &ipv4_set, &ipv6_set, mode, output_format)?;
87
88    Ok(())
89}
90
91/// 全RIRテキストから全ての国コードに対するIPアドレスをパースし、
92/// 指定された国コードのリストに対応するIPアドレスをファイルに書き出す。
93pub 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    // 全ての国コードに対するIPアドレスを一度にパース
100    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    // Arc で包む(大きなマップを参照カウントで共有する)
109    let country_map = Arc::new(country_map);
110
111    // 各国コードに対して非同期タスクを起動
112    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        // Arcの参照カウントを増やしてクローン
118        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    // すべてのタスクが完了するのを待つ
133    for t in tasks {
134        let _ = t.await?;
135    }
136
137    Ok(())
138}