dermy-api/src/account.rs

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<Option<User>> {
    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<Mole>,
    logs: Vec<LogEntry>,
}

async fn get_sign_in(Path(username): Path<String>) -> 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<User>) -> 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<User>) -> 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<String>) -> 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<String>) -> ApiResult {
   let db = &get_database().await?;

   match auth(&api).await? {
       Some(user) => {
           let moles: Vec<Mole> = db.collection::<Mole>("moles")
               .find(doc! { "_user_id": user.id }, None).await?.try_collect().await?;
           
           let logs: Vec<LogEntry> = stream::iter(&moles)
               .then(|mole| async move {
                   db.collection::<LogEntry>("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<String>, Json(user_data): Json<BackupData>) -> 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<T: serde::Serialize>(docs: Vec<T>, 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(())
}