captcha/src/main.rs

// tricaptcha
// Copyright (C) 2023 trimill <trimill012@gmail.com>
//
// This program is licensed under the GNU GPLv3 free
// software license. You are free to modify and distribute
// it under the terms of the license. If you did not recieve
// a copy of the license with this program, it can be found
// at https://www.gnu.org/licenses/gpl-3.0.html.

use std::{fmt::Write, collections::HashMap, sync::Arc, process::Stdio, time::{SystemTime, Duration}, net::ToSocketAddrs};

use axum::{
	routing::{get, post},
	response::{Html, IntoResponse, Response},
	Router, extract::{Path, Query, State}, Form, http::{header, StatusCode}, Json,
};
use log::{info, debug, trace};
use rand::{seq::SliceRandom, Rng, distributions::{Alphanumeric, DistString}};
use serde::{Deserialize, Serialize};
use tokio::{sync::Mutex, process::Command};

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

lazy_static::lazy_static! {
	static ref INDEX: String = std::fs::read_to_string("resources/index.html").unwrap();
	static ref SUBMIT_OK: String = std::fs::read_to_string("resources/submit_ok.html").unwrap();
	static ref SUBMIT_ERR: String = std::fs::read_to_string("resources/submit_err.html").unwrap();

	static ref IMAGES: [Vec<String>; 10] = (0u8..=9).map(|i| {
		let mut path = "resources/images/".to_owned();
		path.push((i + b'0') as char);
		std::fs::read_dir(path.clone()).unwrap()
			.map(|x| path.clone() + "/" + x.unwrap().file_name().to_str().unwrap()).collect()
	}).collect::<Vec<Vec<String>>>().try_into().unwrap();

	static ref AUDIO: [Vec<String>; 10] = (0u8..=9).map(|i| {
		let mut path = "resources/audio/".to_owned();
		path.push((i + b'0') as char);
		std::fs::read_dir(path.clone()).unwrap()
			.map(|x| path.clone() + "/" + x.unwrap().file_name().to_str().unwrap()).collect()
	}).collect::<Vec<Vec<String>>>().try_into().unwrap();
}

struct IncompleteCaptcha {
	digits: String,
	issued: u64,
	got_image: bool,
	got_audio: bool,
	userdata: String,
}

#[derive(Serialize)]
struct CompleteCaptcha {
	count: usize,
	issued: u64,
	completed: u64,
	userdata: String,
}

struct AppState {
	incomplete: Mutex<HashMap<u64, IncompleteCaptcha>>,
	complete:   Mutex<HashMap<String, CompleteCaptcha>>,
}

type AState = Arc<AppState>;

const fn default_count() -> usize { 7 }

async fn next_id(state: &AppState) -> u64 {
	loop {
		let id = rand::thread_rng().gen();
		if !state.incomplete.lock().await.contains_key(&id) {
			debug!("generated id {id}");
			return id
		}
	}
}

async fn mk_captcha_img(digits: &str) -> Result<Vec<u8>> {
	let images: Vec<String> = digits.bytes().map(|b|
		IMAGES[(b - b'0') as usize].choose(&mut rand::thread_rng()).unwrap().clone()
	).collect();

	let cmd = Command::new("convert")
		.arg("+append")
		.args(images)
		.args([
			"-attenuate", "55.0", "+noise", "Uniform",
			"-resize", "x84",
			"-attenuate", "105.0", "+noise", "Uniform",
			"PNG:-"
		]).stdout(Stdio::piped())
		.spawn()?;

	let output = cmd.wait_with_output().await?;

	if !output.status.success() {
		let msg = String::from_utf8(output.stderr)?;
		return Err(format!("imagemagick error: {msg}").into())
	}

	Ok(output.stdout)
}

async fn mk_captcha_audio(digits: &str) -> Result<Vec<u8>> {
	let audio_files: Vec<&str> = digits.bytes().map(|b|
		AUDIO[(b - b'0') as usize].choose(&mut rand::thread_rng()).unwrap().as_str()
	).collect();

	let mut input_args: Vec<&str> = Vec::with_capacity(audio_files.len() * 2);
	for file in &audio_files {
		input_args.push("-i");
		input_args.push(file);
	}

	let mut filter = String::new();
	write!(filter, "aevalsrc=0:d=0.2[x];aevalsrc=0:d=0.4,asplit={}", audio_files.len()).unwrap();
	for i in 0..audio_files.len() {
		write!(filter, "[s{i}]").unwrap();
	}
	write!(filter, ";[x]").unwrap();
	for i in 0..audio_files.len() {
		write!(filter, "[{i}:a][s{i}]").unwrap();
	}
	write!(filter, "concat=n={}:v=0:a=1", audio_files.len()*2 + 1).unwrap();

	let cmd = Command::new("ffmpeg")
		.args(["-loglevel", "16"])
		.args(input_args)
		.args(["-filter_complex", &filter, "-f", "wav", "-"])
		.stdout(Stdio::piped())
		.spawn()?;

	let output = cmd.wait_with_output().await?;

	if !output.status.success() {
		let msg = String::from_utf8(output.stderr)?;
		return Err(format!("ffmpeg error: {msg}").into())
	}

	Ok(output.stdout)
}

#[derive(Deserialize)]
struct IndexQuery {
	#[serde(default="default_count")]
	count: usize,
	#[serde(default)]
	userdata: String,
}

async fn page_index(State(state): State<AState>, Query(query): Query<IndexQuery>) -> Html<String> {
	debug!("loading index page, count {} userdata {:?}", query.count, query.userdata);
	let id = next_id(&state).await;
	let id_string = id.to_string();
	let id_str: &str = id_string.as_ref();

	let count = query.count.clamp(1, 16);

	let digits: String = (0..count)
		.map(|_| (rand::thread_rng().gen_range(0..=9) + b'0') as char)
		.collect();

	let time = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs();

	let cap = IncompleteCaptcha {
		issued: time,
		digits,
		userdata: query.userdata,
		got_image: false,
		got_audio: false,
	};

	state.incomplete.lock().await.insert(id, cap);

	let body = INDEX.replace("{{id}}", id_str);

	Html(body)
}

async fn page_image(State(state): State<AState>, Path(id): Path<u64>) -> Response {
	debug!("loading image {id}");
	let mut incomplete = state.incomplete.lock().await;
	let Some(cap) = incomplete.get_mut(&id) else {
		debug!("invalid image {id}");
		return (StatusCode::BAD_REQUEST, "Invalid ID").into_response()
	};

	if cap.got_image {
		debug!("image {id} already loaded");
		return (StatusCode::BAD_REQUEST, "Image already viewed").into_response()
	}
	cap.got_image = true;

	let digits = cap.digits.clone();
	drop(incomplete);

	tokio::time::sleep(Duration::from_millis(300)).await;

	debug!("generating image");
	let img = match mk_captcha_img(&digits).await {
		Ok(img) => img,
		Err(e) => {
			debug!("image generation failed: {e}");
			return (StatusCode::INTERNAL_SERVER_ERROR, "Internal server error").into_response()
		}
	};
	debug!("image generated");

	let mut res = img.into_response();
	res.headers_mut().insert(
		header::CONTENT_TYPE, "image/png".parse().unwrap()
	);
	res.into_response()
}

async fn page_audio(State(state): State<AState>, Path(id): Path<u64>) -> Response {
	debug!("loading audio {id}");
	let mut incomplete = state.incomplete.lock().await;
	let Some(cap) = incomplete.get_mut(&id) else {
		debug!("invalid audio {id}");
		return (StatusCode::BAD_REQUEST, "Invalid ID").into_response()
	};

	if cap.got_audio {
		debug!("audio {id} already loaded");
		return (StatusCode::BAD_REQUEST, "Audio already accessed").into_response()
	}
	cap.got_audio = true;

	let digits = cap.digits.clone();
	drop(incomplete);

	tokio::time::sleep(Duration::from_millis(300)).await;

	debug!("generating audio");
	let audio = match mk_captcha_audio(&digits).await {
		Ok(audio) => audio,
		Err(e) => {
			debug!("audio generation failed: {e}");
			return (StatusCode::INTERNAL_SERVER_ERROR, "Internal server error").into_response()
		}
	};
	debug!("audio generated");

	let mut res = audio.into_response();
	res.headers_mut().insert(
		header::CONTENT_TYPE, "audio/wav".parse().unwrap()
	);
	res.into_response()
}

#[derive(Deserialize)]
struct SubmitForm {
	id: u64,
	digits: String,
}

async fn page_submit(State(state): State<AState>, Form(form): Form<SubmitForm>) -> impl IntoResponse {
	debug!("submitting id {} digits {}", form.id, form.digits);
	let Some(cap) = state.incomplete.lock().await.remove(&form.id) else {
		debug!("id {} expired", form.id);
		let body = SUBMIT_ERR.replace("{{msg}}", "CAPTCHA expired or ID invalid. Consider trying again.");
		return Html(body)
	};

	if form.digits != cap.digits {
		debug!("id {} failed: expected {}, got {}", form.id, cap.digits, form.digits);
		let body = SUBMIT_ERR.replace("{{msg}}", "CAPTCHA failed. Consider trying again.");
		return Html(body)
	}


	let token = Alphanumeric.sample_string(&mut rand::thread_rng(), 24);

	let count = cap.digits.len();
	let completed = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs();

	debug!("id {} succeeded, generated token {}", form.id, token);

	let body = SUBMIT_OK.replace("{{token}}", &token);

	state.complete.lock().await.insert(token, CompleteCaptcha { count, issued: cap.issued, completed, userdata: cap.userdata });

	Html(body)
}

#[derive(Deserialize)]
struct VerifyQuery {
	token: String,
}

async fn page_verify(State(state): State<AState>, Query(query): Query<VerifyQuery>) -> Response {
	debug!("verifying token {}", query.token);
	let Some(info) = state.complete.lock().await.remove(&query.token) else {
		debug!("verification failed for {}", query.token);
		return (StatusCode::BAD_REQUEST, "Invalid or expired token").into_response()
	};
	debug!("verification succeeded for {}", query.token);

	Json(info).into_response()
}

#[tokio::main]
async fn main() {
	env_logger::init();

	let addr = std::env::var("TCAP_ADDR")
		.unwrap_or_else(|_| "localhost:8000".to_owned())
		.to_socket_addrs().unwrap().next().unwrap();

	let state = Arc::new(AppState {
		incomplete: Mutex::new(HashMap::new()),
		complete: Mutex::new(HashMap::new()),
	});

	let app = Router::new()
		.route("/", get(page_index))
		.route("/image/:id", get(page_image))
		.route("/audio/:id", get(page_audio))
		.route("/submit", post(page_submit))
		.route("/verify", get(page_verify))
		.with_state(state.clone());

	let server = tokio::spawn(async move {
		info!("Listening on http://{addr}");
		axum::Server::bind(&addr)
			.serve(app.into_make_service())
			.await
			.unwrap();
	});

	let timeout = tokio::spawn(async move {
        let mut interval = tokio::time::interval(Duration::from_secs(15));
        loop {
            interval.tick().await;
			trace!("cleaning up");

			let now = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs();

			state.incomplete.lock().await
				.retain(|_, v| now - v.issued < 60*3);
			state.complete.lock().await
				.retain(|_, v| now - v.completed < 60*3);
        }
    });

	let (a, b) = tokio::join!(server, timeout);
	a.unwrap();
	b.unwrap();
}