account/src/json_api/oauth/generate_token_v3.rs

use crate::Pool;
use serde::Serialize;
use rocket::FromForm;
use sha2::{Sha256, Digest};
use bytemuck::bytes_of;
use rocket::{post, State, form::Form, http::Status, serde::json::Json};
use sqlx::Row;
use std::env;
use once_cell::sync::Lazy;
use crate::nnid::oauth::generate_token::create_token;
use crate::nnid::oauth::generate_token::token_type::AUTH_TOKEN;
use chrono::Utc;

pub static CLIENT_SECRET: Lazy<String> = Lazy::new(|| {
    env::var("OAUTH_CLIENT_SECRET").expect("OAUTH_CLIENT_SECRET not set")
});

#[derive(Serialize)]
pub struct OAuthTokenResponse {
    pub access_token: String,
    pub token_type: String,
    pub expires_in: i64,
}

#[derive(FromForm)]
pub struct TokenRequest<'r> {
    pub grant_type: &'r str,
    pub username: Option<&'r str>,
    pub password: Option<&'r str>,
    pub client_id: Option<&'r str>,
    pub client_secret: Option<&'r str>,
    pub code: Option<&'r str>,
    pub redirect_uri: Option<&'r str>,
}

pub fn verify_nintendo_password(pid: i32, text_password: &str, db_bcrypt_hash: &str) -> bool {
    // maple: binder, there's already a function for this, why duplicate code?
    // binder: dear maple, with this function i can just hand it a pid, the password from the user and the db_bcrypt_hash and it gives me a true/false value.
    let mut sha = Sha256::new();
    sha.update(bytes_of(&pid));
    sha.update(&[0x02, 0x65, 0x43, 0x46]);
    sha.update(text_password.as_bytes());
    let hashed_password_hex = hex::encode(sha.finalize());
    match bcrypt::verify(hashed_password_hex, db_bcrypt_hash) {
        Ok(valid) => valid,
        Err(_) => false,
    }
}

// dummy error responses
#[derive(Serialize)]
pub struct OAuthErrorResponse {
    pub error: String,
}

#[post("/api/v3/oauth2/generate_token", data = "<form_data>")]
pub async fn generate_token(
    pool: &State<Pool>, 
    form_data: Form<TokenRequest<'_>>
) -> Result<Json<OAuthTokenResponse>, (Status, Json<OAuthErrorResponse>)> {
    
    if form_data.client_id != Some("account") {
        return Err((
            Status::BadRequest,
            Json(OAuthErrorResponse { error: "invalid_client".to_string() })
        ));
    }

    let cl_secret: String = CLIENT_SECRET.clone();
    if form_data.client_secret != Some(&cl_secret) {
        return Err((
            Status::Unauthorized,
            Json(OAuthErrorResponse { error: "invalid_client".to_string() })
        ));
    }

    // i'm only supporting the password grant incase someone feels lazy.
    let pid: i32 = match form_data.grant_type {
        "password" => {
            let username = form_data.username.ok_or((
                Status::BadRequest,
                Json(OAuthErrorResponse { error: "invalid_request".to_string() })
            ))?;
            let password = form_data.password.ok_or((
                Status::BadRequest,
                Json(OAuthErrorResponse { error: "invalid_request".to_string() })
            ))?;

            let user_row = match sqlx::query("SELECT pid, password FROM users WHERE username = $1")
                .bind(username)
                .fetch_optional(pool.inner())
                .await 
            {
                Ok(Some(row)) => row,
                _ => return Err((
                    Status::BadRequest,
                    Json(OAuthErrorResponse { error: "invalid_grant".to_string() })
                )),
            };

            let db_pid: i32 = user_row.get("pid");
            let db_bcrypt_hash: String = user_row.get("password");

            if !verify_nintendo_password(db_pid, password, &db_bcrypt_hash) {
                return Err((
                    Status::Unauthorized,
                    Json(OAuthErrorResponse { error: "invalid_grant".to_string() })
                ));
            }

            db_pid
        },

        "authorization_code" => {
            let incoming_code = form_data.code.ok_or((
                Status::BadRequest, 
                Json(OAuthErrorResponse { error: "invalid_request".to_string() })
            ))?;
            
            let code_row = match sqlx::query(
                "SELECT pid, expires_at, used FROM oauth_auth_codes WHERE code = $1"
            )
            .bind(incoming_code)
            .fetch_optional(pool.inner())
            .await
            {
                Ok(Some(row)) => row,
                _ => return Err((
                    Status::BadRequest, 
                    Json(OAuthErrorResponse { error: "invalid_grant".to_string() })
                )),
            };

            let is_used: bool = code_row.get("used");
            let expires_at: chrono::NaiveDateTime = code_row.get("expires_at");
            let target_pid: i32 = code_row.get("pid");

            if is_used || expires_at < Utc::now().naive_utc() {
                return Err((
                    Status::BadRequest, 
                    Json(OAuthErrorResponse { error: "invalid_grant".to_string() })
                ));
            }

            let _ = sqlx::query("UPDATE oauth_auth_codes SET used = TRUE WHERE code = $1")
                .bind(incoming_code)
                .execute(pool.inner())
                .await;

            target_pid
        },

        _ => return Err((
            Status::BadRequest,
            Json(OAuthErrorResponse { error: "unsupported_grant_type".to_string() })
        ))
    };

    let token = create_token(pool.inner(), pid, AUTH_TOKEN, None).await;

    Ok(Json(OAuthTokenResponse {
        access_token: token,
        token_type: "Bearer".to_string(),
        expires_in: 3600,
    }))
}