trippy_tui/
geoip.rs

1use anyhow::Context;
2use itertools::Itertools;
3use maxminddb::Reader;
4use std::cell::RefCell;
5use std::collections::HashMap;
6use std::net::IpAddr;
7use std::path::Path;
8use std::rc::Rc;
9use std::str::FromStr;
10
11#[derive(Debug, Clone, Default)]
12pub struct GeoIpCity {
13    latitude: Option<f64>,
14    longitude: Option<f64>,
15    accuracy_radius: Option<u16>,
16    city: Option<String>,
17    subdivision: Option<String>,
18    subdivision_code: Option<String>,
19    country: Option<String>,
20    country_code: Option<String>,
21    continent: Option<String>,
22}
23
24impl GeoIpCity {
25    pub fn short_name(&self) -> String {
26        [
27            self.city.as_ref(),
28            self.subdivision_code.as_ref(),
29            self.country_code.as_ref(),
30        ]
31        .into_iter()
32        .flatten()
33        .join(", ")
34    }
35
36    pub fn long_name(&self) -> String {
37        [
38            self.city.as_ref(),
39            self.subdivision.as_ref(),
40            self.country.as_ref(),
41            self.continent.as_ref(),
42        ]
43        .into_iter()
44        .flatten()
45        .join(", ")
46    }
47
48    pub fn location(&self) -> String {
49        format!(
50            "{}, {} (~{}km)",
51            self.latitude.unwrap_or_default(),
52            self.longitude.unwrap_or_default(),
53            self.accuracy_radius.unwrap_or_default(),
54        )
55    }
56
57    pub const fn coordinates(&self) -> Option<(f64, f64, u16)> {
58        match (self.latitude, self.longitude, self.accuracy_radius) {
59            (Some(lat), Some(long), Some(radius)) => Some((lat, long, radius)),
60            _ => None,
61        }
62    }
63}
64
65mod ipinfo {
66    use serde::{Deserialize, Serialize};
67    use serde_with::serde_as;
68
69    /// The IPinfo mmdb database format.
70    ///
71    /// Support both the "IP to Geolocation Extended" and "IP to Country + ASN" database formats.
72    ///
73    /// IP to Geolocation Extended Database:
74    /// See <https://ipinfo.io/developers/ip-to-geolocation-extended/>
75    ///
76    /// IP to Country + ASN Database;
77    /// See <https://ipinfo.io/developers/ip-to-country-asn-database/>
78    #[serde_as]
79    #[derive(Debug, Serialize, Deserialize)]
80    pub struct IpInfoGeoIp {
81        /// "42.48948"
82        #[serde(default)]
83        #[serde_as(as = "serde_with::NoneAsEmptyString")]
84        pub latitude: Option<String>,
85        /// "-83.14465"
86        #[serde(default)]
87        #[serde_as(as = "serde_with::NoneAsEmptyString")]
88        pub longitude: Option<String>,
89        /// "500"
90        #[serde(default)]
91        #[serde_as(as = "serde_with::NoneAsEmptyString")]
92        pub radius: Option<String>,
93        /// "Royal Oak"
94        #[serde(default)]
95        #[serde_as(as = "serde_with::NoneAsEmptyString")]
96        pub city: Option<String>,
97        /// "Michigan"
98        #[serde(default)]
99        #[serde_as(as = "serde_with::NoneAsEmptyString")]
100        pub region: Option<String>,
101        /// "48067"
102        #[serde(default)]
103        #[serde_as(as = "serde_with::NoneAsEmptyString")]
104        pub postal_code: Option<String>,
105        /// "US"
106        #[serde(default)]
107        #[serde_as(as = "serde_with::NoneAsEmptyString")]
108        pub country: Option<String>,
109        /// "Japan"
110        #[serde(default)]
111        #[serde_as(as = "serde_with::NoneAsEmptyString")]
112        pub country_name: Option<String>,
113        /// "Asia"
114        #[serde(default)]
115        #[serde_as(as = "serde_with::NoneAsEmptyString")]
116        pub continent_name: Option<String>,
117    }
118
119    #[cfg(test)]
120    mod tests {
121        use super::*;
122
123        #[test]
124        fn test_empty() {
125            let json = "{}";
126            let value: IpInfoGeoIp = serde_json::from_str(json).unwrap();
127            assert_eq!(None, value.latitude);
128            assert_eq!(None, value.longitude);
129            assert_eq!(None, value.radius);
130            assert_eq!(None, value.city);
131            assert_eq!(None, value.region);
132            assert_eq!(None, value.postal_code);
133            assert_eq!(None, value.country.as_deref());
134            assert_eq!(None, value.country_name.as_deref());
135            assert_eq!(None, value.continent_name.as_deref());
136        }
137
138        #[test]
139        fn test_country_asn_db_format() {
140            let json = r#"
141                {
142                    "start_ip": "40.96.54.192",
143                    "end_ip": "40.96.54.255",
144                    "country": "JP",
145                    "country_name": "Japan",
146                    "continent": "AS",
147                    "continent_name": "Asia",
148                    "asn": "AS8075",
149                    "as_name": "Microsoft Corporation",
150                    "as_domain": "microsoft.com"
151                }
152                "#;
153            let value: IpInfoGeoIp = serde_json::from_str(json).unwrap();
154            assert_eq!(None, value.latitude);
155            assert_eq!(None, value.longitude);
156            assert_eq!(None, value.radius);
157            assert_eq!(None, value.city);
158            assert_eq!(None, value.region);
159            assert_eq!(None, value.postal_code);
160            assert_eq!(Some("JP"), value.country.as_deref());
161            assert_eq!(Some("Japan"), value.country_name.as_deref());
162            assert_eq!(Some("Asia"), value.continent_name.as_deref());
163        }
164
165        #[test]
166        fn test_extended_db_format() {
167            let json = r#"
168                {
169                    "start_ip": "60.127.10.249",
170                    "end_ip": "60.127.10.249",
171                    "join_key": "60.127.0.0",
172                    "city": "Yokohama",
173                    "region": "Kanagawa",
174                    "country": "JP",
175                    "latitude": "35.43333",
176                    "longitude": "139.65",
177                    "postal_code": "220-8588",
178                    "timezone": "Asia/Tokyo",
179                    "geoname_id": "1848354",
180                    "radius": "500"
181                }
182                "#;
183            let value: IpInfoGeoIp = serde_json::from_str(json).unwrap();
184            assert_eq!(Some("35.43333"), value.latitude.as_deref());
185            assert_eq!(Some("139.65"), value.longitude.as_deref());
186            assert_eq!(Some("500"), value.radius.as_deref());
187            assert_eq!(Some("Yokohama"), value.city.as_deref());
188            assert_eq!(Some("Kanagawa"), value.region.as_deref());
189            assert_eq!(Some("220-8588"), value.postal_code.as_deref());
190            assert_eq!(Some("JP"), value.country.as_deref());
191            assert_eq!(None, value.country_name.as_deref());
192            assert_eq!(None, value.continent_name.as_deref());
193        }
194    }
195}
196
197impl From<ipinfo::IpInfoGeoIp> for GeoIpCity {
198    fn from(value: ipinfo::IpInfoGeoIp) -> Self {
199        Self {
200            latitude: value.latitude.and_then(|val| f64::from_str(&val).ok()),
201            longitude: value.longitude.and_then(|val| f64::from_str(&val).ok()),
202            accuracy_radius: value.radius.and_then(|val| u16::from_str(&val).ok()),
203            city: value.city,
204            subdivision: value.region,
205            subdivision_code: value.postal_code,
206            country: value.country_name,
207            country_code: value.country,
208            continent: value.continent_name,
209        }
210    }
211}
212
213impl From<(maxminddb::geoip2::City<'_>, &str)> for GeoIpCity {
214    fn from((value, locale): (maxminddb::geoip2::City<'_>, &str)) -> Self {
215        let city = value
216            .city
217            .as_ref()
218            .and_then(|city| city.names.as_ref())
219            .and_then(|names| names.get(locale).or_else(|| names.get(FALLBACK_LOCALE)))
220            .map(ToString::to_string);
221        let subdivision = value
222            .subdivisions
223            .as_ref()
224            .and_then(|c| c.first())
225            .and_then(|c| c.names.as_ref())
226            .and_then(|names| names.get(locale).or_else(|| names.get(FALLBACK_LOCALE)))
227            .map(ToString::to_string);
228        let subdivision_code = value
229            .subdivisions
230            .as_ref()
231            .and_then(|c| c.first())
232            .and_then(|c| c.iso_code.as_ref())
233            .map(ToString::to_string);
234        let country = value
235            .country
236            .as_ref()
237            .and_then(|country| country.names.as_ref())
238            .and_then(|names| names.get(locale).or_else(|| names.get(FALLBACK_LOCALE)))
239            .map(ToString::to_string);
240        let country_code = value
241            .country
242            .as_ref()
243            .and_then(|country| country.iso_code.as_ref())
244            .map(ToString::to_string);
245        let continent = value
246            .continent
247            .as_ref()
248            .and_then(|continent| continent.names.as_ref())
249            .and_then(|names| names.get(locale).or_else(|| names.get(FALLBACK_LOCALE)))
250            .map(ToString::to_string);
251        let latitude = value
252            .location
253            .as_ref()
254            .and_then(|location| location.latitude);
255        let longitude = value
256            .location
257            .as_ref()
258            .and_then(|location| location.longitude);
259        let accuracy_radius = value
260            .location
261            .as_ref()
262            .and_then(|location| location.accuracy_radius);
263        Self {
264            latitude,
265            longitude,
266            accuracy_radius,
267            city,
268            subdivision,
269            subdivision_code,
270            country,
271            country_code,
272            continent,
273        }
274    }
275}
276
277/// The fallback locale.
278///
279/// The `MaxMind` support documentation says:
280///
281/// > Our geolocation name data includes the names of the continent, country, city, and
282/// > subdivisions of the location of the IP address. We include the country names in
283/// > English, Simplified Chinese, Spanish, Brazilian Portuguese, Russian, Japanese, French,
284/// > and German.
285/// >
286/// > Please note: Not every place name is always available in each language. We recommend checking
287/// > English names as a default for cases where a localized name is not available in your preferred
288/// > language.
289const FALLBACK_LOCALE: &str = "en";
290
291/// Alias for a cache of `GeoIp` data.
292type Cache = RefCell<HashMap<IpAddr, Rc<GeoIpCity>>>;
293
294/// Lookup `GeoIpCity` data form an `IpAddr`.
295#[derive(Debug)]
296pub struct GeoIpLookup {
297    reader: Option<Reader<Vec<u8>>>,
298    cache: Cache,
299    locale: String,
300}
301
302impl GeoIpLookup {
303    /// Create a new `GeoIpLookup` from a `MaxMind` DB file.
304    pub fn from_file<P: AsRef<Path>>(path: P, locale: String) -> anyhow::Result<Self> {
305        let reader = maxminddb::Reader::open_readfile(path.as_ref())
306            .context(format!("{}", path.as_ref().display()))?;
307        Ok(Self {
308            reader: Some(reader),
309            cache: RefCell::new(HashMap::new()),
310            locale,
311        })
312    }
313
314    /// Create a `GeoIpLookup` that returns `None` for all `IpAddr` lookups.
315    pub fn empty() -> Self {
316        Self {
317            reader: None,
318            cache: RefCell::new(HashMap::new()),
319            locale: FALLBACK_LOCALE.to_string(),
320        }
321    }
322
323    /// Lookup an `GeoIpCity` for an `IpAddr`.
324    ///
325    /// If an entry is found it is cached and returned, otherwise None is returned.
326    pub fn lookup(&self, addr: IpAddr) -> anyhow::Result<Option<Rc<GeoIpCity>>> {
327        if let Some(reader) = &self.reader {
328            if let Some(geo) = self.cache.borrow().get(&addr) {
329                return Ok(Some(geo.clone()));
330            }
331            let city_data = if reader.metadata.database_type.starts_with("ipinfo") {
332                GeoIpCity::from(reader.lookup::<ipinfo::IpInfoGeoIp>(addr)?)
333            } else {
334                GeoIpCity::from((
335                    reader.lookup::<maxminddb::geoip2::City<'_>>(addr)?,
336                    self.locale.as_ref(),
337                ))
338            };
339            let geo = self.cache.borrow_mut().insert(addr, Rc::new(city_data));
340            Ok(geo)
341        } else {
342            Ok(None)
343        }
344    }
345}