Merge pull request #15 from bgpkit/feature/hegemony

add support for loading hegemony score from bgpkit mirror

examples/asnames.rs

@@ -2,6 +2,7 @@ use bgpkit_commons::asnames::{get_asnames, AsName};
 use std::collections::HashMap;
 
 fn main() {
+    tracing_subscriber::fmt::init();
     let asnames: HashMap<u32, AsName> = get_asnames().unwrap();
     println!(
         "{}",

src/asnames/hegemony.rs

@@ -0,0 +1,80 @@
+use anyhow::{anyhow, Result};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use tracing::{info, warn};
+
+const IIJ_IHR_HEGEMONY_IPV4_GLOBAL: &str =
+    "https://data.bgpkit.com/ihr/hegemony/ipv4/global/latest-simplified.csv.gz";
+const IIJ_IHR_HEGEMONY_IPV6_GLOBAL: &str =
+    "https://data.bgpkit.com/ihr/hegemony/ipv6/global/latest-simplified.csv.gz";
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Hegemony {
+    hegemony_map: HashMap<u32, HegemonyData>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct HegemonyData {
+    pub asn: u32,
+    pub ipv4: f64,
+    pub ipv6: f64,
+}
+
+fn load_hegemony(path: &str) -> Result<Vec<(u32, f64)>> {
+    info!("loading hegemony scores from {}", path);
+    // load global hegemony scores data CSV file with header where there are two columns: ASN, score
+    let mut hegemony = Vec::new();
+    for line in oneio::read_lines(path)? {
+        let text = line.ok().ok_or(anyhow!("error reading line"))?;
+        if text.trim() == "" || text.starts_with("asn") {
+            continue;
+        }
+        let splits: Vec<&str> = text.split(',').map(|x| x.trim()).collect();
+        if splits.len() != 2 {
+            return Err(anyhow!("row missing fields: {}", text.as_str()));
+        }
+        let asn = match splits[0].parse::<u32>() {
+            Ok(asn) => asn,
+            Err(_) => {
+                warn!("invalid ASN: {}", text);
+                continue;
+            }
+        };
+        let score = splits[1].parse::<f64>()?;
+        hegemony.push((asn, score));
+    }
+    Ok(hegemony)
+}
+
+impl Hegemony {
+    pub fn new() -> Result<Self> {
+        let ipv4 = load_hegemony(IIJ_IHR_HEGEMONY_IPV4_GLOBAL)?;
+        let ipv6 = load_hegemony(IIJ_IHR_HEGEMONY_IPV6_GLOBAL)?;
+        let mut hegemony_map = HashMap::new();
+        for (asn, score) in ipv4 {
+            hegemony_map.insert(
+                asn,
+                HegemonyData {
+                    asn,
+                    ipv4: score,
+                    ipv6: 0.0,
+                },
+            );
+        }
+        for (asn, score) in ipv6 {
+            hegemony_map
+                .entry(asn)
+                .or_insert_with(|| HegemonyData {
+                    asn,
+                    ipv4: 0.0,
+                    ipv6: 0.0,
+                })
+                .ipv6 = score;
+        }
+        Ok(Self { hegemony_map })
+    }
+
+    pub fn get_score(&self, asn: u32) -> Option<&HegemonyData> {
+        self.hegemony_map.get(&asn)
+    }
+}

src/asnames/mod.rs

@@ -45,8 +45,10 @@
 //! assert_eq!(asnames.get(&400644).unwrap().country, "US");
 //! ```
 
+mod hegemony;
 mod population;
 
+use crate::asnames::hegemony::HegemonyData;
 use crate::asnames::population::AsnPopulationData;
 use anyhow::Result;
 use serde::{Deserialize, Serialize};
@@ -60,6 +62,7 @@ pub struct AsName {
     pub country: String,
     pub as2org: Option<As2orgInfo>,
     pub population: Option<AsnPopulationData>,
+    pub hegemony: Option<HegemonyData>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -79,6 +82,8 @@ pub fn get_asnames() -> Result<HashMap<u32, AsName>> {
     let as2org = as2org_rs::As2org::new(None)?;
     info!("loading ASN population data from APNIC...");
     let population = population::AsnPopulation::new()?;
+    info!("loading IIJ IHR hegemony score data from BGPKIT mirror...");
+    let hegemony = hegemony::Hegemony::new()?;
 
     let asnames = text
         .lines()
@@ -105,6 +110,7 @@ pub fn get_asnames() -> Result<HashMap<u32, AsName>> {
                 country: country_str.to_string(),
                 as2org,
                 population,
+                hegemony: hegemony.get_score(asn).cloned(),
             })
         })
         .collect::<Vec<AsName>>();