Mercurial > lbo > hg > analyrics

changeset 36:0c5f8caf6736

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Implement geoip lookups
author Lewin Bormann <lbo@spheniscida.de>
date Sat, 16 Jul 2022 10:45:58 -0700
parents 7a64e348786d
children 4347956773fb
files Cargo.toml Rocket.toml src/geoip.rs src/logsdb.rs src/main.rs usertool/src/main.rs
diffstat 6 files changed, 104 insertions(+), 11 deletions(-) [+]
line wrap: on
line diff
--- a/Cargo.toml	Sat Jul 16 10:45:28 2022 -0700
+++ b/Cargo.toml	Sat Jul 16 10:45:58 2022 -0700
@@ -17,6 +17,7 @@
 sha256 = "1.0.3"
 time = { version = "0.3.11", features = ["serde", "serde-well-known"] }
 tokio = { version = "1.19.2", features = ["fs"] }
+maxminddb = "0.23.0"
 
 [features]
 default = ["sqlite"]
--- a/Rocket.toml	Sat Jul 16 10:45:28 2022 -0700
+++ b/Rocket.toml	Sat Jul 16 10:45:58 2022 -0700
@@ -14,6 +14,9 @@
 [default]
 port = 8000
 secret_key = "edeaaff3f1774ad2888673770c6d64097e391bc362d7d6fb34982ddf0efd18cb"
+
+
 template_dir = "assets"
 asset_path = "assets/static"
 deploy_path = "/"
+geodb_path = "/usr/share/GeoIP/GeoLite2-City.mmdb"
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/src/geoip.rs	Sat Jul 16 10:45:58 2022 -0700
@@ -0,0 +1,46 @@
+use anyhow::{Context, Error};
+use log::warn;
+use maxminddb::{
+    geoip2::{City, Country},
+    Reader,
+};
+
+use std::collections::BTreeMap;
+use std::net::IpAddr;
+use std::path::Path;
+
+pub struct GeoIP {
+    r: Reader<Vec<u8>>,
+}
+
+impl GeoIP {
+    pub fn new<P: AsRef<Path>>(path: P) -> Result<GeoIP, Error> {
+        Ok(GeoIP {
+            r: Reader::open_readfile(path)?,
+        })
+    }
+
+    pub fn lookup(&self, addr: IpAddr) -> Option<(String, String)> {
+        match self.r.lookup::<City>(addr) {
+            Ok(c) => Some((
+                c.country
+                    .map(|c| c.iso_code.unwrap_or(""))
+                    .unwrap_or("")
+                    .to_string(),
+                c.city
+                    .map(|c| {
+                        c.names
+                            .unwrap_or_else(BTreeMap::new)
+                            .get("en")
+                            .unwrap_or(&"")
+                            .to_string()
+                    })
+                    .unwrap_or(String::new()),
+            )),
+            Err(e) => {
+                warn!("GeoIP lookup failed: {}", e);
+                None
+            }
+        }
+    }
+}
--- a/src/logsdb.rs	Sat Jul 16 10:45:28 2022 -0700
+++ b/src/logsdb.rs	Sat Jul 16 10:45:58 2022 -0700
@@ -72,11 +72,25 @@
         Ok(ntags)
     }
 
-    pub async fn start_session(&mut self, domain: &str) -> Result<u32, Error> {
-        Ok(sqlx::query("INSERT INTO Sessions (start, last, domain) VALUES (strftime('%s', 'now'), strftime('%s', 'now'), ?) RETURNING id")
-            .bind(domain)
-            .fetch_one(&mut *self.0)
-            .await?.get(0))
+    // Takes domain and (country, city) tuple.
+    pub async fn start_session(
+        &mut self,
+        domain: &str,
+        orig: Option<(String, String)>,
+    ) -> Result<u32, Error> {
+        let (country, city) = orig.unwrap_or((Default::default(), Default::default()));
+        Ok(sqlx::query(
+            r#"
+INSERT INTO Sessions (start, last, domain, origin_country, origin_city)
+VALUES (strftime('%s', 'now'), strftime('%s', 'now'), ?, ?, ?)
+RETURNING id"#,
+        )
+        .bind(domain)
+        .bind(country)
+        .bind(city)
+        .fetch_one(&mut *self.0)
+        .await?
+        .get(0))
     }
 
     pub async fn update_session_time(&mut self, id: u32) -> Result<(), Error> {
--- a/src/main.rs	Sat Jul 16 10:45:28 2022 -0700
+++ b/src/main.rs	Sat Jul 16 10:45:58 2022 -0700
@@ -1,10 +1,12 @@
 mod configdb;
 mod db;
 mod fromparam;
+mod geoip;
 mod guards;
 mod logsdb;
 
 use crate::configdb::{ConfigDB, ConfigDBSession};
+use crate::geoip::GeoIP;
 use crate::guards::{HeadersGuard, LoggedInGuard, USER_ID_COOKIE_KEY};
 use crate::logsdb::{LogsDB, LogsDBSession};
 
@@ -21,6 +23,7 @@
 use rocket::response::{self, Flash, Redirect, Responder};
 use rocket::serde::json::{json, Value};
 use rocket::serde::{self, Serialize};
+use rocket::State;
 
 use rocket_db_pools::sqlx::{Executor, Row, Sqlite, SqlitePool};
 use rocket_db_pools::{Connection, Database, Pool};
@@ -31,6 +34,7 @@
 use std::collections::HashMap;
 use std::net::IpAddr;
 use std::path::Path;
+use std::str::FromStr;
 
 #[derive(Responder)]
 enum LoginResponse {
@@ -285,7 +289,16 @@
         Err(e) => "undefined".to_string(),
     };
 
-    let begin_nav = [("-30d", -30), ("-14d", -14), ("-7d", -7), ("-1d", -1), ("+1d", 1), ("+7d", 7), ("+14d", 14), ("+30d", 30)];
+    let begin_nav = [
+        ("-30d", -30),
+        ("-14d", -14),
+        ("-7d", -7),
+        ("-1d", -1),
+        ("+1d", 1),
+        ("+7d", 7),
+        ("+14d", 14),
+        ("+30d", 30),
+    ];
 
     Template::render(
         "index",
@@ -315,6 +328,8 @@
 async fn route_log(
     mut conn: Connection<LogsDB>,
     cookies: &CookieJar<'_>,
+    geoipdb: &State<Option<GeoIP>>,
+
     host: Option<String>,
     status: Option<u32>,
     path: Option<String>,
@@ -339,9 +354,15 @@
                 }
             }
         }
+
+        let sip = ip.to_string();
+        let sip = headers.0.get_one("x-real-ip").unwrap_or(&sip);
+        let numip = IpAddr::from_str(sip).unwrap_or(ip);
+
         if session_id.is_none() {
+            let orig = geoipdb.as_ref().and_then(|g| g.lookup(numip));
             match conn
-                .start_session(host.as_ref().map(String::as_str).unwrap_or(""))
+                .start_session(host.as_ref().map(String::as_str).unwrap_or(""), orig)
                 .await
             {
                 Ok(id) => {
@@ -357,8 +378,6 @@
 
         let ntags = tags.len() as u32;
         let ua = headers.0.get_one("user-agent").unwrap_or("");
-        let ip = ip.to_string();
-        let ip = headers.0.get_one("x-real-ip").unwrap_or(&ip);
         let host: String = host.unwrap_or(
             headers
                 .0
@@ -372,7 +391,7 @@
         match conn
             .log_request(
                 session_id,
-                ip,
+                sip,
                 host,
                 path.unwrap_or(String::new()),
                 status.unwrap_or(200),
@@ -402,6 +421,7 @@
 struct CustomConfig {
     asset_path: String,
     deploy_path: String,
+    geodb_path: String,
 }
 
 #[rocket::launch]
@@ -420,11 +440,20 @@
         .expect("absolute path to asset directory")
         .to_path_buf();
 
+    let geoip = match crate::geoip::GeoIP::new(cfg.geodb_path) {
+        Ok(g) => Some(g),
+        Err(e) => {
+            error!("Couldn't open geoip database: {}", e);
+            None
+        }
+    };
+
     let r = r
         .attach(ConfigDB::init())
         .attach(LogsDB::init())
         .attach(Template::fairing())
         .attach(rocket::fairing::AdHoc::config::<CustomConfig>())
+        .manage(geoip)
         .mount("/static", FileServer::from(ap))
         .mount(
             "/",
--- a/usertool/src/main.rs	Sat Jul 16 10:45:28 2022 -0700
+++ b/usertool/src/main.rs	Sat Jul 16 10:45:58 2022 -0700
@@ -8,7 +8,7 @@
     let a = args().collect::<Vec<String>>();
 
     let ip = IpAddr::from_str(&a[1]).expect("IP address from argument");
-    let citydb = Reader::open_readfile("/usr/share/GeoIP/GeoLite2-City.mmdb").expect("city db");
+    let citydb = Reader::open_readfile("/usr/share/GeoIP/GeoLite2-City.mmdb").expect("country db");
 
     let cy: maxminddb::geoip2::City = citydb.lookup(ip).expect("city from ip");
     println!("{:?}", cy);