Write some major geolocation shuffling around code

src/main.rs

@@ -5,6 +5,8 @@ use axum::{
 	Router,
 	routing::get,
 };
+use maxminddb;
+use maxminddb::geoip2;
 use tera::Tera;
 use trust_dns_resolver::{
 	TokioAsyncResolver,
@@ -15,8 +17,13 @@ use trust_dns_resolver::{
 use std::net::{IpAddr, Ipv4Addr};
 use std::sync::Arc;
 
+
 mod simple_dns;
 mod templating_engine;
+mod geoip;
+
+use geoip::AsnResult;
+use geoip::LocationResult;
 
 use crate::templating_engine::View;
 use crate::templating_engine::ResponseFormat;
@@ -36,11 +43,14 @@ pub struct DigQuery {
 #[derive(serde::Deserialize, serde::Serialize)]
 pub struct IpResult {
 	hostname: Option<String>,
+	asn:      Option<AsnResult>,
 }
 
 struct ServiceSharedState {
 	templating_engine: templating_engine::Engine,
 	dns_resolver:      TokioAsyncResolver,
+	mmdb_asn:          Option<maxminddb::Reader<Vec<u8>>>,
+	mmdb_location:     Option<maxminddb::Reader<Vec<u8>>>,
 }
 
 #[tokio::main]
@@ -56,6 +66,23 @@ async fn main() {
 			::std::process::exit(1);
 		}
 	};
+
+	// Initalize GeoIP Database
+
+	println!("Opening GeoIP ASN Databse ...");
+	let mmdb_asn = maxminddb::Reader::open_readfile("mmdb/GeoLite2-ASN.mmdb");
+	match mmdb_asn {
+		Ok(_) => { /* NOP */ },
+		Err(ref e) => println!("Error while opening GeoIP ASN Databse: {e}"),
+	} 
+
+	println!("Opening GeoIP Location Databse ...");
+	let mmdb_location = maxminddb::Reader::open_readfile("mmdb/GeoLite2-City.mmdb");
+	match mmdb_location {
+		Ok(_) => { /* NOP */ },
+		Err(ref e) => println!("Error while opening GeoIP Location Databse: {e}"),
+	} 
+
 	
 	// Initalize DNS resolver with os defaults
 	println!("Initalizing dns resolver ...");
@@ -74,6 +101,8 @@ async fn main() {
 			tera: tera,
 		},
 		dns_resolver: dns_resolver,
+		mmdb_asn: mmdb_asn.ok(),
+		mmdb_location: mmdb_location.ok(),
 	});
 	
 	// Initalize axum server
@@ -115,9 +144,47 @@ async fn handle_default_route(
 
 	// do reverse lookup
 	let hostname = simple_dns::reverse_lookup(&state.dns_resolver, &address);
+	
+	let asn_result = match &state.mmdb_asn {
+		Some(mmdb_asn) => {
+			match mmdb_asn.lookup::<geoip2::Asn>(address) {
+				Ok(res) => {
+					Some(AsnResult {
+						asn: res.autonomous_system_number,
+						name: res.autonomous_system_organization.map(ToString::to_string),
+					})
+				}
+				Err(e) => {
+					println!("Error while looking up ASN for {address}: {e}");
+					Default::default()
+				}
+			}
+		},
+		None => None,
+	};
+
+	let location_result = match &state.mmdb_location {
+		Some(mmdb_location) => {
+			match mmdb_location.lookup::<geoip2::City>(address) {
+				Ok(res) => {
+					Some(AsnResult {
+						asn: res.autonomous_system_number,
+						name: res.autonomous_system_organization.map(ToString::to_string),
+					})
+				}
+				Err(e) => {
+					println!("Error while looking up ASN for {address}: {e}");
+					Default::default()
+				}
+			}
+		},
+		None => None,
+	};
+	
 
 	let result = IpResult{
 		hostname: hostname.await,
+		asn: asn_result,
 	};
 
 	state.templating_engine.render_view(
@@ -136,7 +203,7 @@ async fn handle_dig_route(
 	let format = dig_query.format.unwrap_or(ResponseFormat::TextHtml);
 
 	let dig_result = simple_dns::lookup(&state.dns_resolver, name, true).await;
-
+	
 	state.templating_engine.render_view(
 		format,
 		View::Dig{ query: dig_query, result: dig_result}