pub mod db; use anyhow::Result; use axum::{ extract::Path, body::Body, response::{IntoResponse, Response}, routing::{get, post}, Json, Router }; use futures::stream::{self, StreamExt, TryStreamExt}; use serde::{Deserialize, Serialize}; use crate::ApiResult; use db::{get_database, get_users, LogEntry, Mole, User}; use http::StatusCode; use mongodb::bson::{doc, oid::ObjectId, to_document, Document}; use serde_json::json; pub fn router() -> Router { Router::new() .route("/sign-in/:username", get(get_sign_in)) .route("/sign-in", post(post_sign_in) ) .route("/sign-up", post(post_sign_up)) .route("/sign-out/:api", get(get_sign_out)) .route("/backup/:api", post(post_backup) ) .route("/restore/:api", get(get_restore) ) } pub async fn auth(api: &str) -> Result> { let db = get_users().await?; let query = doc! { "$expr": { "$eq": ["$_auth._api", ObjectId::parse_str(api)?] } }; Ok(db.find_one(query, None).await?) } #[derive(Serialize, Deserialize)] struct BackupData { moles: Vec, logs: Vec, } async fn get_sign_in(Path(username): Path) -> ApiResult { let db = get_users().await?; let query = doc! { "$expr": { "$eq": ["$username", username] } }; match db.find_one(query, None).await? { Some(user) => { Ok(Response::builder() .status(StatusCode::CREATED) .body(Json(json!({ "_salt": user.auth.unwrap().salt.unwrap() })).into_response().into_body())?) }, None => { Ok(Response::builder() .status(StatusCode::NOT_FOUND) .body(Json(json!({ "_error": "User does not exist" })).into_response().into_body())?) } } } async fn post_sign_in(Json(body): Json) -> ApiResult { let db = get_users().await?; let api = ObjectId::new(); let query = doc! { "$expr": { "$and": [ { "$eq": ["$username", body.username] }, { "$eq": ["$_auth._hash", body.auth.unwrap_or_default().hash] } ] } }; let update = doc! { "$set": { "_auth._api": api } }; match db.count_documents(query.clone(), None).await? { count if count != 0 => { db.update_one(query, update, None).await?; Ok(Response::builder() .status(StatusCode::OK) .body(Json(json!({ "_api_key": api.to_string() })).into_response().into_body())?) }, _ => { Ok(Response::builder() .status(StatusCode::UNAUTHORIZED) .body(Json(json!({ "_error": "Username or password is incorrect" })).into_response().into_body())?) }, } } async fn post_sign_up(Json(body): Json) -> ApiResult { let db = get_users().await?; let query = doc! { "$expr": { "$eq": ["$username", &body.username] } }; match db.count_documents(query, None).await? { count if count != 0 => { Ok(Response::builder() .status(StatusCode::NOT_ACCEPTABLE) .body(Json(json!({ "_error": "Username is already taken" })).into_response().into_body())?) }, _ => { db.insert_one(body, None).await?; Ok(Response::builder() .status(StatusCode::CREATED) .body(Json(json!({ "_success": "Account created successfully" })).into_response().into_body())?) } } } async fn get_sign_out(Path(api): Path) -> ApiResult { let db = get_users().await?; let query = doc! { "$expr": { "$eq": ["$_auth._api", ObjectId::parse_str(&api)?] } }; let update = doc! { "$unset": { "_auth._api": ObjectId::new() } }; match auth(&api).await? { Some(_user) => { db.update_one(query, update, None).await?; Ok(Response::builder() .status(StatusCode::OK) .body(Json(json!({ "_success": "Sign out successful" })).into_response().into_body())?) }, None => { Ok(Response::builder() .status(StatusCode::NOT_ACCEPTABLE) .body(Json(json!({ "_error": "User does not exist" })).into_response().into_body())?) }, } } async fn get_restore(Path(api): Path) -> ApiResult { let db = &get_database().await?; match auth(&api).await? { Some(user) => { let moles: Vec = db.collection::("moles") .find(doc! { "_user_id": user.id }, None).await?.try_collect().await?; let logs: Vec = stream::iter(&moles) .then(|mole| async move { db.collection::("logs") .find(doc! { "_mole_id": mole.id}, None) .await.unwrap() .try_collect() .await.unwrap_or_else(|_| vec![]) }) .flat_map(stream::iter).collect().await; Ok(Response::builder() .status(StatusCode::OK) .body(Json(json!({ "_moles": moles, "_logs": logs })).into_response().into_body()).unwrap()) }, None => Ok(Response::builder() .status(StatusCode::UNAUTHORIZED) .body(Body::from("API Key is incorrect"))?), } } async fn post_backup(Path(api): Path, Json(user_data): Json) -> ApiResult { match auth(&api).await? { Some(_user) => { let _ = backup_docs_if_new(user_data.moles, "moles").await; let _ = backup_docs_if_new(user_data.logs, "logs").await; Ok(Response::builder() .status(StatusCode::CREATED) .body(Body::from("Account backup successful"))?) }, None => Ok(Response::builder() .status(StatusCode::UNAUTHORIZED) .body(Body::from("API Key is incorrect"))?), } } async fn backup_docs_if_new(docs: Vec, collection: &'static str) -> Result<()> { let _ = docs.iter() .map(|doc| to_document(doc).unwrap()) .map(|doc| async move { let _ = insert_if_new(doc, collection).await; }); Ok(()) } async fn insert_if_new(doc: Document, collection: &str) -> Result<()> { let db = get_database().await?.collection(collection); match db.count_documents(doc.clone(), None).await? { 0 => Some(db.insert_one(doc, None)), _ => None, }; Ok(()) }