use std::collections::hash_map::HashMap; use std::net::SocketAddr; use axum::extract::{ConnectInfo, Path}; use axum::http::HeaderMap; use axum::http::StatusCode; use axum::response::{Html, IntoResponse, Redirect}; use axum::Extension; use info_utils::prelude::*; use crate::ServerState; use crate::TrackingRow; use crate::UdsConnectInfo; use crate::UrlRow; enum TrackingParameter { Ip, Referrer, UserAgent, } pub async fn index(Extension(state): Extension) -> impl IntoResponse { if let Some(redirect) = state.main_page_redirect { return Redirect::temporary(redirect.as_str()).into_response(); } Html(format!( r#" {} URL Shortener

{} URL shortener

create "#, state.host, state.host )) .into_response() } pub async fn id_unix( headers: HeaderMap, ConnectInfo(addr): ConnectInfo, Extension(state): Extension, Path(id): Path, ) -> impl IntoResponse { let ip = if state.behind_proxy { match headers.get("x-real-ip") { Some(it) => { if let Ok(i) = it.to_str() { Some(i.to_string()) } else { None } } None => None, } } else { Some(format!("{:?}", addr.peer_addr)) } .unwrap_or_default(); run_id(headers, ip, state, id).await } /// # Panics /// Will panic if `parse()` fails pub async fn id( headers: HeaderMap, ConnectInfo(addr): ConnectInfo, Extension(state): Extension, Path(id): Path, ) -> impl IntoResponse { let ip = get_ip(&headers, addr, &state).unwrap_or_default(); run_id(headers, ip, state, id).await } async fn run_id( headers: HeaderMap, ip: String, state: ServerState, id: String, ) -> impl IntoResponse { let mut show_request = false; log!("Request for '{}' from {}", id.clone(), ip); let mut use_id = id; if use_id.ends_with('+') { show_request = true; use_id.pop(); } let item: Result = sqlx::query_as("SELECT * FROM chela.urls WHERE id = $1") .bind(use_id.clone()) .fetch_one(&state.db_pool) .await; if let Ok(it) = item { if url::Url::parse(&it.url).is_ok() { if show_request { return Html(format!( r#"

http://{}/{} -> {}

"#, state.host, it.id, it.url, it.url )) .into_response(); } log!("Redirecting {} -> {}", it.id, it.url); save_analytics(headers, it.clone(), ip, state).await; let mut response_headers = HeaderMap::new(); response_headers.insert("Cache-Control", "private, max-age=90".parse().unwrap()); response_headers.insert("Location", it.url.parse().unwrap()); return ( StatusCode::MOVED_PERMANENTLY, response_headers, Html(format!( r#"Redirecting to {}"#, it.url, it.url )), ) .into_response(); } } else { warn!("'{}' not found.", use_id); return ( StatusCode::NOT_FOUND, Html("

Not found.

".to_string()), ) .into_response(); } (StatusCode::NOT_FOUND, Html("

Not found.

")).into_response() } async fn save_analytics(headers: HeaderMap, item: UrlRow, ip: String, state: ServerState) { let id = item.id; let referer = match headers.get("referer") { Some(it) => { if let Ok(i) = it.to_str() { Some(i) } else { None } } None => None, }; let user_agent = match headers.get("user-agent") { Some(it) => { if let Ok(i) = it.to_str() { Some(i) } else { None } } None => None, }; let res = sqlx::query( " INSERT INTO chela.tracking (id,ip,referrer,user_agent) VALUES ($1,$2,$3,$4) ", ) .bind(id.clone()) .bind(ip.clone()) .bind(referer) .bind(user_agent) .execute(&state.db_pool) .await; if res.is_ok() { log!("Saved analytics for '{id}' from {}", ip); } } fn get_ip(headers: &HeaderMap, addr: SocketAddr, state: &ServerState) -> Option { if state.behind_proxy { match headers.get("x-real-ip") { Some(it) => { if let Ok(i) = it.to_str() { Some(i.to_string()) } else { None } } None => None, } } else { Some(addr.ip().to_string()) } } pub async fn create_id(Extension(state): Extension) -> Html { Html(format!( r#" {} URL Shortener "#, state.host )) } pub async fn tracking(Extension(state): Extension) -> impl IntoResponse { let url_rows: Vec = sqlx::query_as("SELECT * FROM chela.urls") .fetch_all(&state.db_pool) .await .unwrap(); let html = format!( r#" {} Tracking {} "#, state.host, table_css(), make_table_from_urls(&url_rows) ); return Html(html).into_response(); } pub async fn tracking_id( Extension(state): Extension, Path(id): Path, ) -> impl IntoResponse { let tracking_rows: Vec = sqlx::query_as("SELECT * FROM chela.tracking WHERE id = $1") .bind(id.clone()) .fetch_all(&state.db_pool) .await .unwrap(); let url: UrlRow = sqlx::query_as("SELECT * FROM chela.urls WHERE id = $1") .bind(id.clone()) .fetch_one(&state.db_pool) .await .unwrap(); let html = format!( r#" {} Tracking {}

Tracking for {} from ID '{}'

Visited {} times

{}

By IP

{}

By Referrer

{}

By User Agent

{} "#, state.host, id, table_css(), url.url, url.url, url.id, tracking_rows.len(), make_table_from_tracking(&tracking_rows), make_grouped_table_from_tracking(&tracking_rows, TrackingParameter::Ip), make_grouped_table_from_tracking(&tracking_rows, TrackingParameter::Referrer), make_grouped_table_from_tracking(&tracking_rows, TrackingParameter::UserAgent) ); return Html(html).into_response(); } fn make_table_from_tracking(rows: &Vec) -> String { let mut html = r#" "# .to_string(); for row in rows { html += &format!( r#" "#, row.timestamp, row.id, row.ip.as_ref().unwrap_or(&String::default()), row.referrer.as_ref().unwrap_or(&String::default()), row.user_agent.as_ref().unwrap_or(&String::default()) ); } html += r#"

Timestamp	ID	IP	Referrer	User Agent
{}	{}	{}	{}	{}

"#; html } fn make_grouped_table_from_tracking(rows: &Vec, group: TrackingParameter) -> String { let column_name = match group { TrackingParameter::Ip => "IP", TrackingParameter::Referrer => "Referrer", TrackingParameter::UserAgent => "User Agent", } .to_string(); let mut html = format!( r#" "#, column_name ); let mut aggregate: HashMap = HashMap::new(); for row in rows { let tracker = match group { TrackingParameter::Ip => { let v = match &row.ip { Some(val) => val, None => continue, }; v } TrackingParameter::Referrer => { let v = match &row.referrer { Some(val) => val, None => continue, }; v } TrackingParameter::UserAgent => { let v = match &row.user_agent { Some(val) => val, None => continue, }; v } }; let count = aggregate.get(tracker).unwrap_or(&0); aggregate.insert(tracker.to_string(), count + 1); } for (key, val) in aggregate { html += &format!( r#" "#, val, key ); } html += r#"

Occurrences	{}
{}	{}

"#; html } fn make_table_from_urls(urls: &Vec) -> String { let mut html = r#" "# .to_string(); for url in urls { html += &format!( r#" "#, url.index, url.id, url.id, url.url, url.url, url.custom_id ); } html += r#"

Index	ID	URL	Custom ID
{}	{}	{}	{}

"#; html } fn table_css() -> String { r#" tr:nth-child(even) { background: #f2f2f2; } table, th, td { border: 1px solid black; } "# .to_string() }