summaryrefslogtreecommitdiff
path: root/scripts/maint/geoip/geoip-db-tool/src/db.rs
blob: 316182d8237865d57f0de66e68b7b571597cd744 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
/// Code to parse a dump file
use std::collections::HashMap;
use std::convert::TryInto;
use std::iter::Peekable;

use super::{AsBlock, NetBlock};

pub struct BlockReader<I>
where
    I: Iterator<Item = std::io::Result<String>>,
{
    iter: Peekable<I>,
}

pub enum AnyBlock {
    NetBlock(NetBlock),
    AsBlock(AsBlock),
    OtherBlock,
}

impl<I> BlockReader<I>
where
    I: Iterator<Item = std::io::Result<String>>,
{
    pub fn new(iter: I) -> Self {
        BlockReader {
            iter: iter.peekable(),
        }
    }

    /// Extract the initial header from the file.
    pub fn extract_header(&mut self) -> String {
        let mut res: String = "".to_string();

        while let Some(Ok(line)) = self.iter.peek() {
            if !line.starts_with('#') {
                break;
            }
            res.push_str(line.as_str());
            res.push('\n');
            let _ = self.iter.next();
        }

        res
    }

    /// Extract the next empty-line-delimited block from the file.
    ///
    /// This isn't terribly efficient, but it's "fast enough".
    fn get_block(&mut self) -> Option<std::io::Result<AnyBlock>> {
        let mut kv = HashMap::new();

        while let Some(line) = self.iter.next() {
            //dbg!(&line);
            if let Err(e) = line {
                return Some(Err(e));
            }
            let line_orig = line.unwrap();
            let line = line_orig.splitn(2, '#').next().unwrap().trim();
            if line.is_empty() {
                if kv.is_empty() {
                    continue;
                } else {
                    break;
                }
            }
            let kwds: Vec<_> = line.splitn(2, ':').collect();
            if kwds.len() != 2 {
                return None; // XXXX handle the error better.
            }
            kv.insert(kwds[0].trim().to_string(), kwds[1].trim().to_string());
        }

        if kv.is_empty() {
            return None;
        }

        if let Some(name) = kv.remove("name") {
            // This is an AS block.
            let asn = kv.get("aut-num").unwrap(); // XXXX handle error better
            assert!(asn.starts_with("AS"));
            let asn = asn[2..].parse().unwrap();
            return Some(Ok(AnyBlock::AsBlock(AsBlock { name, asn })));
        }

        let net = if let Some(net) = kv.get("net") {
            net.parse().unwrap() //XXXX handle the error better.
        } else {
            return Some(Ok(AnyBlock::OtherBlock));
        };

        let asn = if let Some(asn) = kv.get("aut-num") {
            asn.parse().ok()
        } else {
            None
        };

        let cc = if let Some(country) = kv.get("country") {
            assert!(country.as_bytes().len() == 2);
            country.as_bytes()[0..2].try_into().unwrap()
        } else {
            *b"??"
        };

        fn is_true(v: Option<&String>) -> bool {
            match v {
                Some(s) => s == "true",
                None => false,
            }
        }

        let is_anon_proxy = is_true(kv.get("is-anonymous-proxy"));
        let is_anycast = is_true(kv.get("is-anycast-proxy"));
        let is_satellite = is_true(kv.get("is-satellite-provider"));

        Some(Ok(AnyBlock::NetBlock(NetBlock {
            net,
            asn,
            cc,
            is_anon_proxy,
            is_anycast,
            is_satellite,
        })))
    }
}

impl<I> Iterator for BlockReader<I>
where
    I: Iterator<Item = std::io::Result<String>>,
{
    type Item = AnyBlock;
    fn next(&mut self) -> Option<Self::Item> {
        match self.get_block() {
            Some(Ok(b)) => Some(b),
            _ => None,
        }
    }
}