echoip-slatecave/src/main.rs

use axum::{
	extract::Query,
	extract::State,
	http::StatusCode,
	response::Html,
	response::IntoResponse,
	response::Response,
	response::Json,
	Router,
	routing::get,
};
use tera::Tera;
use trust_dns_resolver::{
	TokioAsyncResolver,
	config::ResolverOpts,
	config::ResolverConfig,
};


use std::net::{IpAddr, Ipv4Addr};
use std::sync::Arc;

mod simple_dns;

#[derive(serde::Deserialize, serde::Serialize, Clone, Copy)]
enum ResponseFormat {
	#[serde(rename="text/plain", alias="text")]
	TextPlain,
	#[serde(rename="text/html", alias="html")]
	TextHtml,
	#[serde(rename="application/json", alias="json")]
	ApplicationJson,
}

impl ToString for ResponseFormat {
	fn to_string(&self) -> String {
		match self {
			ResponseFormat::TextPlain => "text/plain",
			ResponseFormat::TextHtml => "text/html",
			ResponseFormat::ApplicationJson => "application/json",
		}.to_string()
	}
}

#[derive(serde::Deserialize, serde::Serialize)]
struct IpQuery {
	ip:     Option<IpAddr>,
	format: Option<ResponseFormat>,
}

#[derive(serde::Deserialize, serde::Serialize)]
struct DigQuery {
	name:   String,
	format: Option<ResponseFormat>,
}

#[derive(serde::Deserialize, serde::Serialize)]
struct IpResult {
	hostname: Option<String>,
}

struct TemplatingEngine {
	tera: Tera,
}

struct ServiceSharedState {
	templating_engine: TemplatingEngine,
	dns_resolver:      TokioAsyncResolver,
}

#[derive(serde::Deserialize, serde::Serialize)]
#[serde(untagged)]
enum EchoipView {
	Dig { query: DigQuery, result: simple_dns::DnsLookupResult },
	Index { query: IpQuery, result: IpResult },
	Ip { query: IpQuery, result: IpResult },
	Message(String),
	#[serde(rename="404")]
	NotFound,
}

impl EchoipView {
	fn template_name(&self) -> String {
		match self {
			EchoipView::Dig{..} => "dig",
			EchoipView::Index{..} => "index",
			EchoipView::Ip{..} => "ip",
			EchoipView::Message(..) => "message",
			EchoipView::NotFound => "404",
		}.to_string()
	}
}


#[tokio::main]
async fn main() {
	// Initalize Tera templates
	// TODO: don't hardcode template directory
	println!("Parsing Templates ...");
	let res = Tera::new("templates/*.html");
	let tera = match res {
		Ok(t) => t,
		Err(e) => {
			println!("Template parsing error(s): {}", e);
			::std::process::exit(1);
		}
	};
	
	// Initalize DNS resolver with os defaults
	println!("Initalizing dns resolver ...");
	let res = TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default());
	let dns_resolver = match res {
		Ok(resolver) => resolver,
		Err(e) => {
			println!("Error while setting up dns resolver: {e}");
			::std::process::exit(1);
		}
	};
	
	// Initialize shared state
	let shared_state = Arc::new(ServiceSharedState{
		templating_engine: TemplatingEngine{
			tera: tera,
		},
		dns_resolver: dns_resolver,
	});
	
	// Initalize axum server
    let app = Router::new()
	    .route("/", get(handle_default_route))
	    .route("/dig", get(handle_dig_route))
    	.route("/hi", get(hello_world_handler))
    	.with_state(shared_state)
    	;

	println!("Starting Server ...");

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn hello_world_handler(
	State(arc_state): State<Arc<ServiceSharedState>>,
) -> Response {
	let state = Arc::clone(&arc_state);

	state.templating_engine.render_view(
		ResponseFormat::TextPlain,
		EchoipView::Message("Hello! There, You, Awesome Creature!".to_string())
	).await	
}

impl TemplatingEngine {
	async fn render_view(
		&self,
		format: ResponseFormat,
		view: EchoipView,
	) -> Response {
		match format {
			ResponseFormat::TextHtml => {
				let template_name = view.template_name();
				
				let mut context = tera::Context::new();
				context.insert("view", &template_name);
				//intented for shared macros
				context.insert("format", &format.to_string());
				context.insert("data", &view);

				match self.tera.render(&(template_name+".html"), &context) {
					Ok(html) => Html(html).into_response(),
					Err(e) => {
						println!("There was an error while rendering index.html: {e:?}");
						StatusCode::INTERNAL_SERVER_ERROR.into_response()
					}
				}
			}
			//TODO: Plain Text should have its own matcher 
			ResponseFormat::ApplicationJson | ResponseFormat::TextPlain => {
				match view {
					EchoipView::Dig{result, ..} => {
						Json(result).into_response()
					},
					EchoipView::Index{result, ..} | EchoipView::Ip{result, ..} => {
						Json(result).into_response()
					},
					_ => Json(view).into_response(),
				}
			}
		}
	}
}

async fn handle_default_route(
	Query(ip_query): Query<IpQuery>,
	State(arc_state): State<Arc<ServiceSharedState>>,
) -> Response {

	let address = ip_query.ip.unwrap_or(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)));
	let format = ip_query.format.unwrap_or(ResponseFormat::TextHtml);

	let state = Arc::clone(&arc_state);

	// do reverse lookup
	let hostname = simple_dns::reverse_lookup(&state.dns_resolver, &address);

	let result = IpResult{
		hostname: hostname.await,
	};

	state.templating_engine.render_view(
		format,
		EchoipView::Index{query: ip_query, result: result}
	).await
}

async fn handle_dig_route(
	Query(dig_query): Query<DigQuery>,
	State(arc_state): State<Arc<ServiceSharedState>>,
) -> Response {
	
	let state = Arc::clone(&arc_state);
	let name = &dig_query.name;
	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,
		EchoipView::Dig{ query: dig_query, result: dig_result}
	).await

}