nebulous 0.1.86

A globally distributed container orchestrator
Documentation
use crate::auth::db;
use crate::auth::models;
use crate::auth::models::SanitizedApiKey;
use argon2::{
    password_hash::{PasswordHash, PasswordVerifier, SaltString},
    Argon2, PasswordHasher,
};
use base64::{engine::general_purpose, Engine as _};
use rand::rngs::OsRng;
use rand::RngCore;
use sea_orm::entity::*;
use sea_orm::DatabaseConnection;
use uuid::Uuid;

pub async fn get_api_key(
    db_conn: &DatabaseConnection,
    id: &str,
) -> Result<SanitizedApiKey, Box<dyn std::error::Error>> {
    let result = db::Entity::find_by_id(id).one(db_conn).await?;
    match result {
        Some(api_key) => Ok(models::ApiKey::from(api_key).into()),
        None => Err("API key not found".into()),
    }
}

pub async fn get_sanitized_api_key(
    db_conn: &DatabaseConnection,
    id: &str,
) -> Result<SanitizedApiKey, Box<dyn std::error::Error>> {
    let api_key = get_api_key(db_conn, id).await?;
    Ok(models::SanitizedApiKey::from(api_key))
}

pub async fn list_api_keys(
    db_conn: &DatabaseConnection,
) -> Result<Vec<SanitizedApiKey>, Box<dyn std::error::Error>> {
    let result = db::Entity::find().all(db_conn).await?;
    Ok(result
        .into_iter()
        .map(models::ApiKey::from)
        .map(models::SanitizedApiKey::from)
        .collect())
}

pub async fn generate_api_key(
    db_conn: &DatabaseConnection,
) -> Result<String, Box<dyn std::error::Error>> {
    let mut raw_key = [0u8; 32];
    OsRng.fill_bytes(&mut raw_key);
    let key = general_purpose::STANDARD.encode(&raw_key);
    let salt = SaltString::generate(&mut OsRng);

    let argon2 = Argon2::default();
    let hash = match argon2.hash_password(key.as_str().as_bytes(), salt.as_salt()) {
        Ok(hash) => hash.to_string(),
        Err(_) => return Err("Failed to hash API key.".to_string().into()),
    };

    let id = Uuid::new_v4().to_string();
    let api_key = models::ApiKey::new(id.clone(), hash);
    let new_api_key: db::ActiveModel = db::Model::from(api_key).into();
    new_api_key.insert(db_conn).await?;

    Ok(format!("nebu-{}.{}", id, key.as_str().to_string()))
}

pub async fn validate_api_key(
    db_conn: &DatabaseConnection,
    provided_key: &str,
) -> Result<bool, Box<dyn std::error::Error + Sync + Send>> {
    if let Some(full_key) = provided_key.strip_prefix("nebu-") {
        let parts: Vec<&str> = full_key.split('.').collect();
        if parts.len() == 2 {
            let (id, key) = (parts[0], parts[1]);
            if let Some(mut api_key) = db::Entity::find_by_id(id).one(db_conn).await? {
                if api_key.revoked_at.is_some() {
                    return Ok(false);
                }
                let parsed_hash = match PasswordHash::new(&api_key.hash) {
                    Ok(hash) => hash,
                    Err(_) => return Ok(false),
                };
                let argon2 = Argon2::default();
                return match argon2.verify_password(key.as_bytes(), &parsed_hash) {
                    Ok(_) => {
                        let mut active_api_key: db::ActiveModel = api_key.into();
                        active_api_key.last_used_at = Set(Some(chrono::Utc::now()));
                        active_api_key.update(db_conn).await?;
                        Ok(true)
                    }
                    Err(_) => Ok(false),
                };
            }
        }
    }
    Ok(false)
}

pub async fn revoke_api_key(
    db_conn: &DatabaseConnection,
    id: &str,
) -> Result<SanitizedApiKey, Box<dyn std::error::Error>> {
    if let Some(mut api_key) = db::Entity::find_by_id(id).one(db_conn).await? {
        let mut current_api_key: db::ActiveModel = api_key.into();
        current_api_key.revoked_at = Set(Some(chrono::Utc::now()));
        current_api_key.hash = Set(String::new());
        let result = current_api_key.update(db_conn).await?;
        Ok(models::ApiKey::from(result).into())
    } else {
        Err("API key not found".into())
    }
}