From e3054e0158262df15e8acb573b5c261d64c9c0c1 Mon Sep 17 00:00:00 2001
From: Slatian <baschdel@disroot.org>
Date: Wed, 22 Feb 2023 23:07:43 +0100
Subject: [PATCH] implemented text query logic

---
 Cargo.lock               |  2 ++
 Cargo.toml               |  2 ++
 src/main.rs              | 47 ++++++++++++++++++++++++++++++++--------
 src/templating_engine.rs |  8 +++----
 4 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5eb9c99..82a3e72 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1190,7 +1190,9 @@ dependencies = [
  "axum",
  "axum-client-ip",
  "clap",
+ "lazy_static",
  "maxminddb",
+ "regex",
  "serde",
  "tera",
  "tokio",
diff --git a/Cargo.toml b/Cargo.toml
index dcc7beb..2aa8a08 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,6 +10,8 @@ authors = ["Slatian <baschdel@disroot.org>"]
 axum = { version = "0.6", features = ["macros", "headers"] }
 axum-client-ip = "0.4"
 clap = { version = "4", features = ["derive"] }
+lazy_static = "1.4.0"
+regex = "1.7"
 serde = { version = "1", features = ["derive"] }
 tokio = { version = "1", features = ["full"] }
 tera = "1"
diff --git a/src/main.rs b/src/main.rs
index 4060757..1c0f0c9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,6 +15,8 @@ use axum::{
 };
 use axum_client_ip::SecureClientIp;
 use clap::Parser;
+use lazy_static::lazy_static;
+use regex::Regex;
 use tera::Tera;
 use tower::ServiceBuilder;
 use trust_dns_resolver::{
@@ -60,6 +62,7 @@ pub struct SearchQuery {
 
 #[derive(serde::Deserialize, serde::Serialize, Clone)]
 pub struct IpResult {
+	address:  IpAddr,
 	hostname: Option<String>,
 	asn:      Option<AsnResult>,
 	location: Option<LocationResult>,
@@ -296,6 +299,10 @@ async fn handle_default_route(
 	
 	let state = Arc::clone(&arc_state);
 
+	if let Some(search_query) = search_query.query {
+		return handle_search_request(search_query, false, settings, state).await;
+	}
+
 	let result = get_ip_result(&address, &settings.lang, &state).await;
 
 	let user_agent: Option<String> = match user_agent_header {
@@ -306,26 +313,46 @@ async fn handle_default_route(
 	state.templating_engine.render_view(
 		&settings,
 		&View::Index{
-		query: address,
 			result: result,
 			user_agent: user_agent,
 		}
 	).await
 }
 
+
 async fn handle_search_request(
 	search_query: String,
-	this_should_have_been_an_ip: bool,
+	_this_should_have_been_an_ip: bool,
 	settings: TemplateSettings,
 	arc_state: Arc<ServiceSharedState>,
 ) -> Response {
 
-	let state = Arc::clone(&arc_state);
-	
-	state.templating_engine.render_view(
-		&settings,
-		&View::Message{title: "Your Search query was:".to_string(), message: search_query}
-	).await	
+	let search_query = search_query.trim();
+
+	lazy_static!{
+		static ref ASN_REGEX: Regex = Regex::new(r"^[Aa][Ss][Nn]?\s*(\d{1,7})$").unwrap();
+	}
+
+	//If someone asked for an asn, give an asn answer
+	if let Some(asn_cap) = ASN_REGEX.captures(&search_query) {
+		if let Some(asn) = asn_cap.get(1).map_or(None, |m| m.as_str().parse::<u32>().ok()) {
+			// Render a dummy template that can at least link to other pages
+			let state = Arc::clone(&arc_state);
+			return state.templating_engine.render_view(
+				&settings,
+				&View::Asn{asn: asn},
+			).await
+		}
+	}
+
+	// Try to interpret as an IP-Address
+	if let Ok(address) = search_query.parse() {
+		return handle_ip_request(address, settings, arc_state).await;
+	}
+
+	// Fall back to treating it as a hostname
+	return handle_dig_request(search_query.to_string(), settings, arc_state).await
+
 }
 
 async fn handle_ip_route_with_path(
@@ -351,7 +378,7 @@ async fn handle_ip_request(
 
 	state.templating_engine.render_view(
 		&settings,
-		&View::Ip{query: address, result: result}
+		&View::Ip{result: result}
 	).await
 }
 
@@ -366,6 +393,7 @@ async fn get_ip_result(
 	if !(ip_info.scope == AddressScope::Global || ip_info.scope == AddressScope::Shared) || ip_info.cast != AddressCast::Unicast {
 		if !((ip_info.scope == AddressScope::Private || ip_info.scope == AddressScope::LinkLocal) && state.config.server.allow_private_ip_lookup) {
 			return IpResult {
+				address: *address,
 				hostname: None,
 				asn: None,
 				location: None,
@@ -403,6 +431,7 @@ async fn get_ip_result(
 	};
 
 	IpResult{
+		address: *address,
 		hostname: final_hostname,
 		asn: asn_result,
 		location: location_result,
diff --git a/src/templating_engine.rs b/src/templating_engine.rs
index ec6a15b..54f3e11 100644
--- a/src/templating_engine.rs
+++ b/src/templating_engine.rs
@@ -13,8 +13,6 @@ use axum::{
 use tera::Tera;
 use toml::Table;
 
-use std::net::IpAddr;
-
 use crate::simple_dns;
 use crate::IpResult;
 
@@ -53,9 +51,10 @@ pub struct TemplateSettings {
 #[derive(serde::Deserialize, serde::Serialize, Clone)]
 #[serde(untagged)]
 pub enum View {
+	Asn { asn: u32 },
 	Dig { query: String, result: simple_dns::DnsLookupResult },
-	Index { query: IpAddr, result: IpResult, user_agent: Option<String> },
-	Ip { query: IpAddr, result: IpResult },
+	Index { result: IpResult, user_agent: Option<String> },
+	Ip { result: IpResult },
 	Message{ title: String, message: String },
 	#[serde(rename="404")]
 	NotFound,
@@ -64,6 +63,7 @@ pub enum View {
 impl View {
 	pub fn template_name(&self) -> String {
 		match self {
+			View::Asn{..} => "asn",
 			View::Dig{..} => "dig",
 			View::Index{..} => "index",
 			View::Ip{..} => "ip",