db.rs (3637B)
1 /// Code to parse a dump file 2 use std::collections::HashMap; 3 use std::convert::TryInto; 4 use std::iter::Peekable; 5 6 use super::{AsBlock, NetBlock}; 7 8 pub struct BlockReader<I> 9 where 10 I: Iterator<Item = std::io::Result<String>>, 11 { 12 iter: Peekable<I>, 13 } 14 15 pub enum AnyBlock { 16 Net(NetBlock), 17 As(AsBlock), 18 Other, 19 } 20 21 impl<I> BlockReader<I> 22 where 23 I: Iterator<Item = std::io::Result<String>>, 24 { 25 pub fn new(iter: I) -> Self { 26 BlockReader { 27 iter: iter.peekable(), 28 } 29 } 30 31 /// Extract the initial header from the file. 32 pub fn extract_header(&mut self) -> String { 33 let mut res: String = "".to_string(); 34 35 while let Some(Ok(line)) = self.iter.peek() { 36 if !line.starts_with('#') { 37 break; 38 } 39 res.push_str(line.as_str()); 40 res.push('\n'); 41 let _ = self.iter.next(); 42 } 43 44 res 45 } 46 47 /// Extract the next empty-line-delimited block from the file. 48 /// 49 /// This isn't terribly efficient, but it's "fast enough". 50 fn get_block(&mut self) -> Option<std::io::Result<AnyBlock>> { 51 let mut kv = HashMap::new(); 52 53 for line in self.iter.by_ref() { 54 //dbg!(&line); 55 if let Err(e) = line { 56 return Some(Err(e)); 57 } 58 let line_orig = line.unwrap(); 59 let line = line_orig.split('#').next().unwrap().trim(); 60 if line.is_empty() { 61 if kv.is_empty() { 62 continue; 63 } else { 64 break; 65 } 66 } 67 let kwds: Vec<_> = line.splitn(2, ':').collect(); 68 if kwds.len() != 2 { 69 return None; // XXXX handle the error better. 70 } 71 kv.insert(kwds[0].trim().to_string(), kwds[1].trim().to_string()); 72 } 73 74 if kv.is_empty() { 75 return None; 76 } 77 78 if let Some(name) = kv.remove("name") { 79 // This is an AS block. 80 let asn = kv.get("aut-num").unwrap(); // XXXX handle error better 81 assert!(asn.starts_with("AS")); 82 let asn = asn[2..].parse().unwrap(); 83 return Some(Ok(AnyBlock::As(AsBlock { name, asn }))); 84 } 85 86 let net = if let Some(net) = kv.get("net") { 87 net.parse().unwrap() //XXXX handle the error better. 88 } else { 89 return Some(Ok(AnyBlock::Other)); 90 }; 91 92 let asn = if let Some(asn) = kv.get("aut-num") { 93 asn.parse().ok() 94 } else { 95 None 96 }; 97 98 let cc = if let Some(country) = kv.get("country") { 99 assert!(country.len() == 2); 100 country.as_bytes()[0..2].try_into().unwrap() 101 } else { 102 *b"??" 103 }; 104 105 fn is_true(v: Option<&String>) -> bool { 106 match v { 107 Some(s) => s == "true", 108 None => false, 109 } 110 } 111 112 let is_anon_proxy = is_true(kv.get("is-anonymous-proxy")); 113 let is_anycast = is_true(kv.get("is-anycast-proxy")); 114 let is_satellite = is_true(kv.get("is-satellite-provider")); 115 116 Some(Ok(AnyBlock::Net(NetBlock { 117 net, 118 asn, 119 cc, 120 is_anon_proxy, 121 is_anycast, 122 is_satellite, 123 }))) 124 } 125 } 126 127 impl<I> Iterator for BlockReader<I> 128 where 129 I: Iterator<Item = std::io::Result<String>>, 130 { 131 type Item = AnyBlock; 132 fn next(&mut self) -> Option<Self::Item> { 133 match self.get_block() { 134 Some(Ok(b)) => Some(b), 135 _ => None, 136 } 137 } 138 }