path: root/src/main.rs

#[macro_use] extern crate rocket;

use std::{collections::HashSet, process::exit};

use rocket::{http::Status, serde::json::Json, State};
use rocket_cors::CorsOptions;
use serde::{Deserialize, Serialize};

mod uid;
use tokio::sync::Mutex;
use uid::UID;

mod db;
use db::{DB, User, Config, Token, Color};

#[get("/")]
async fn index() -> &'static str {
    "This is a POKO server"
}

#[derive(Serialize)]
struct AuthUserOut {
    name: String,
    id: UID,
    is_admin: bool,
    tokovec: Vec<Token>,
}
impl Into<AuthUserOut> for User {
    fn into(self) -> AuthUserOut {
        AuthUserOut { name: self.get_name(), id: self.get_id(), is_admin: self.is_admin(), tokovec: self.get_tokovec() }
    }
}
#[derive(Serialize)]
struct UserOut {
    name: String,
    id: UID,
}
impl Into<UserOut> for User {
    fn into(self) -> UserOut {
        UserOut { name: self.get_name(), id: self.get_id() }
    }
}
#[post("/get", data="<data>", format="json")]
async fn get_user_authenticated(data: Json<LogoutForm>, db: &State<Mutex<DB>>) -> (Status, Result<Json<AuthUserOut>, Json<String>>) {
    let db = db.lock().await;
    match db.get_user_authenticated(data.id, &data.session).await {
        Ok(n) => (Status::Ok, Ok(Json(n.clone().into()))),
        Err(n) => (Status::NotFound, Err(n.into()))
    } 
}
#[get("/<user>")]
async fn get_user(user: UID, db: &State<Mutex<DB>>) -> (Status, Result<Json<UserOut>, Json<String>>) {
    let db = db.lock().await;
    match db.get_user(user) {
        Ok(n) => (Status::Ok, Ok(Json(n.clone().into()))),
        Err(n) => (Status::NotFound, Err(n.into()))
    }
}
#[get("/by_name/<user>")]
async fn get_users_by_name(user: &str, db: &State<Mutex<DB>>) -> (Status, Result<Json<Vec<UserOut>>, Json<String>>) {
    let db = db.lock().await;
    match db.get_user_by_name(user).await {
        Ok(n) => (Status::Ok, Ok(Json(n.into_iter().map(|v| v.into()).collect()))),
        Err(n) => (Status::NotFound, Err(n.into()))
    }
}
#[get("/list")]
async fn get_all_users(db: &State<Mutex<DB>>) -> (Status, Result<Json<Vec<UserOut>>, Json<String>>) {
    let db = db.lock().await;
    match db.get_all_users().await {
        Ok(n) => (Status::Ok, Ok(Json(n.into_iter().map(|v| v.into()).collect()))),
        //Nothing found
        Err(n) => (Status::NotFound, Err(n.into())),
    }
}

fn default_id() -> UID {
    uid::EMPTY_UID
}
#[derive(Deserialize)]
struct RegisterForm {
    name: String,
    password: String,
    #[serde(default = "default_id")]
    id: UID,
    key: String,
}
#[post("/register", data="<data>", format="json")]
async fn new_user(data: Json<RegisterForm>, db: &State<Mutex<DB>>) -> (Status, Result<Json<UserOut>, Json<String>>) {
    let mut db = db.lock().await;
    let id;

    let name = data.name.to_string();
    let name = name.trim();
    if name.is_empty() {
        return (Status::BadRequest, Err("Invalid Username".to_string().into()))
    }

    if data.id == uid::EMPTY_UID {
        match db.uid_generator.new_uid() {
            // Gotten a new UID to make a new user with
            Ok(n) => id = Some(n),
            // Did not get a new UID
            Err(n) => return (Status::InternalServerError, Err(n.clone().into())),
        };
    } else {
        if db.uid_generator.is_taken(data.id) {
            // UID is taken
            return (Status::BadRequest, Err("UID is taken".to_string().into()));
        } else {
            // UID is not taken
            match db.uid_generator.add_uid(data.id) {
                // Could not add UID for some reason
                Err(n) => return (Status::InternalServerError, Err(n.into())),
                // Made 
                Ok(n) => id = Some(n),
            };
        }
    }

    match id {
        None => return (Status::InternalServerError, Err("".to_string().into())),
        Some(n) => {
            match db.new_user(name.into(), data.password.clone(), n, &data.key).await {
                // User has been created
                Ok(u) => return (Status::Created, Ok(Json(u.into()))),
                // Could not create user
                Err(e) => {
                    db.uid_generator.delete_uid(n);
                    return (Status::InternalServerError, Err(e.into()))
                },
            }
        }
    }
}

#[derive(Deserialize)]
struct LoginForm {
    id: UID,
    password: String,
    clientid: String,
}
#[post("/login", data="<data>", format="json", rank=2)]
async fn login(data: Json<LoginForm>, db: &State<Mutex<DB>>) -> (Status, Json<String>) {
    let mut db = db.lock().await;
    match db.login(data.id, &data.password, &data.clientid).await {
        Ok(n) => (Status::Ok, n.into()),
        Err(n) => (Status::Unauthorized, n.into()),
    }
}

#[derive(Deserialize)]
struct LogoutForm {
    id: UID,
    session: String,
}
#[post("/logout", data="<data>", format="json")]
async fn logout(data: Json<LogoutForm>, db: &State<Mutex<DB>>) -> (Status, Result<(), Json<String>>) {
    let mut db = db.lock().await;
    match db.logout(data.id, &data.session).await {
        Ok(_) => (Status::Ok, Ok(())),
        Err(n) => (Status::Unauthorized, Err(n.into())),
    }
}
#[post("/logout/all", data="<data>", format="json")]
async fn logout_all(data: Json<LogoutForm>, db: &State<Mutex<DB>>) -> (Status, Result<Json<String>, Json<String>>) {
    let mut db = db.lock().await;
    match db.logout_all(data.id, &data.session).await {
        Ok(n) => (Status::Ok, Ok(n.into())),
        Err(n) => (Status::Unauthorized, Err(n.into())),
    }
}
#[post("/sessions", data="<data>", format="json")]
async fn get_sessions(data: Json<LogoutForm>, db: &State<Mutex<DB>>) -> (Status, Result<Json<Vec<(String, String)>>, Json<String>>) {
    let db = db.lock().await;
    match db.get_sessions(data.id, &data.session).await {
        Ok(n) => (Status::Ok, Ok(n.into())),
        Err(n) => (Status::Unauthorized, Err(n.into())),
    }
}

#[derive(Deserialize)]
struct DeleteForm {
    id: UID,
    session: String,
    password: String,
}

#[post("/delete", data="<data>", format="json")]
async fn delete(data: Json<DeleteForm>, db: &State<Mutex<DB>>) -> (Status, Result<Json<String>, Json<String>>) {
    let mut db = db.lock().await;
    match db.delete_user(data.id, &data.session, &data.password).await {
        Ok(n) => (Status::Ok, Ok(n.into())),
        Err(n) => (Status::Unauthorized, Err(n.into())),
    }
}

#[derive(Deserialize)]
struct UpdateNameForm {
    id: UID,
    session: String,
    name: String,
}
#[derive(Deserialize)]
struct UpdatePasswordForm {
    id: UID,
    session: String,
    old_password: String,
    new_password: String,
}
#[post("/name", data="<data>", format="json")]
async fn update_password(data: Json<UpdateNameForm>, db: &State<Mutex<DB>>) -> (Status, Result<(), Json<String>>) {
    let mut db = db.lock().await;
    match db.update_user_name(data.id, &data.session, &data.name).await {
        Ok(_) => (Status::Ok, Ok(())),
        Err(n) => (Status::InternalServerError, Err(n.into()))
    }
}

#[post("/password", data="<data>", format="json")]
async fn update_name(data: Json<UpdatePasswordForm>, db: &State<Mutex<DB>>) -> (Status, Result<(), Json<String>>) {
    let mut db = db.lock().await;
    match db.update_user_password(data.id, &data.session, &data.old_password, &data.new_password).await {
        Ok(_) => (Status::Ok, Ok(())),
        Err(n) => (Status::InternalServerError, Err(n.into()))
    }
}

#[derive(Deserialize)]
struct NewRegistrationKeyForm {
    id: UID,
    session: String,
    key: String,
}

#[post("/regkey/new", data="<data>", format="json")]
async fn new_registration_key(data: Json<NewRegistrationKeyForm>, db: &State<Mutex<DB>>) -> (Status, Result<(), Json<String>>) {
    let mut db = db.lock().await;
    match db.new_registration_key(data.id, &data.session, &data.key).await {
        Ok(_) => (Status::Ok, Ok(())),
        Err(n) => (Status::InternalServerError, Err(n.into()))
    }
}

#[post("/regkey/del", data="<data>", format="json")]
async fn del_registration_key(data: Json<NewRegistrationKeyForm>, db: &State<Mutex<DB>>) -> (Status, Result<(), Json<String>>) {
    let mut db = db.lock().await;
    match db.del_registration_key(data.id, &data.session, &data.key).await {
        Ok(_) => (Status::Ok, Ok(())),
        Err(n) => (Status::InternalServerError, Err(n.into()))
    }
}

#[post("/regkey/list", data="<data>", format="json")]
async fn list_registration_keys(data: Json<LogoutForm>, db: &State<Mutex<DB>>) -> (Status, Result<Json<Vec<String>>, Json<String>>) {
    let db = db.lock().await;
    match db.list_registration_keys(data.id, &data.session).await {
        Ok(n) => (Status::Ok, Ok(n.into())),
        Err(n) => (Status::InternalServerError, Err(n.into())),
    }
}

#[derive(Deserialize, Debug)]
struct TransferForm {
    id: UID,
    session: String,
    to: UID,
    color: Color,
    amount: usize,
}
#[post("/out", data="<data>", format="json")]
async fn transfer_out(data: Json<TransferForm>, db: &State<Mutex<DB>>) -> (Status, Result<(), Json<String>>) {
    let mut db = db.lock().await;
    match db.transfer(data.id, data.to, &data.session, data.color, data.amount).await {
        Ok(n) => (Status::Ok, Ok(n.into())),
        Err(n) => (Status::InternalServerError, Err(n.into()))
    }
}

#[launch]
fn rocket() -> _ {
    let cors = CorsOptions{
        allowed_origins: rocket_cors::AllOrSome::All,
        allowed_methods: HashSet::from([rocket_cors::Method::from(rocket::http::Method::Post), rocket_cors::Method::from(rocket::http::Method::Get)]),
        allow_credentials: false,
        allowed_headers: rocket_cors::AllOrSome::All,
        expose_headers: HashSet::new(),
        max_age: None,
        send_wildcard: false,
        fairing_route_base: "/cors".to_string(),
        fairing_route_rank: 0,
    };
    let cors = cors.to_cors().unwrap();

    let mut config_path = "config.json".to_string();

    let args: Vec<String> = std::env::args().collect();
    for a in args {
        let b: Vec<&str> = a.split('=').collect();
        match b[0] {
            "config" => {
                config_path = b[1].to_string();
            },
            "genconfig" => {
                Config::genconfig();
                exit(0);
            }
            _ => ()
        }
    }

    let config = match Config::load(config_path.into()) {
        Ok(n) => n,
        Err(_) => Config::new(),
    };

    let rocket_config = rocket::config::Config::figment()
        .merge(("address", config.address()))
        .merge(("port", config.port()));

    rocket::custom(rocket_config).manage(Mutex::new(DB::load(&config)))
    .mount("/", routes![index])
    .mount("/user", routes![login, get_users_by_name, get_user_authenticated, get_user, new_user, get_all_users, logout, logout_all, get_sessions, delete])
    .mount("/user/update", routes![update_name, update_password])
    .mount("/transfer", routes![transfer_out])
    .mount("/admin", routes![new_registration_key, list_registration_keys, del_registration_key])
    .attach(cors)
}