prisma_auth 0.2.1

// src/multi_user.rs
use rocket::request::{FromRequest, Outcome, Request};
use rocket::http::Status;
use rocket::serde::json::Json;
use rocket::State;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use uuid::Uuid;
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
use argon2::password_hash::{rand_core::OsRng, SaltString};

use crate::{LoginRequest, RegisterRequest, TokenResponse, User, Role, AuthError, DEFAULT_TOKEN_TTL_SECONDS};
use crate::storage::{UserStorage, InMemoryStorage, current_timestamp};

/// Token store for managing JWT tokens with user association
pub struct MultiUserTokenStore {
    tokens: RwLock<HashMap<String, (String, Instant)>>, // token -> (username, timestamp)
}

impl MultiUserTokenStore {
    /// Create a new empty token store
    pub fn new() -> Self {
        Self {
            tokens: RwLock::new(HashMap::new()),
        }
    }

    /// Insert a token for a specific user
    pub fn insert(&self, token: String, username: String) {
        self.tokens.write().unwrap().insert(token, (username, Instant::now()));
    }

    /// Validate token and return username if valid
    pub fn validate(&self, token: &str) -> Option<String> {
        let mut map = self.tokens.write().unwrap();
        if let Some((username, timestamp)) = map.get(token) {
            let ttl = get_token_ttl();
            if timestamp.elapsed() <= ttl {
                return Some(username.clone());
            } else {
                // Token expired, remove it
                map.remove(token);
            }
        }
        None
    }

    /// Remove expired tokens
    pub fn cleanup_expired(&self) {
        let mut map = self.tokens.write().unwrap();
        let ttl = get_token_ttl();
        map.retain(|_, (_, timestamp)| timestamp.elapsed() <= ttl);
    }

    /// Remove all tokens for a specific user (logout all sessions)
    pub fn revoke_user_tokens(&self, username: &str) {
        let mut map = self.tokens.write().unwrap();
        map.retain(|_, (user, _)| user != username);
    }
}

impl Default for MultiUserTokenStore {
    fn default() -> Self {
        Self::new()
    }
}

/// Authentication service for user management and password handling
pub struct AuthService<S: UserStorage> {
    storage: Arc<S>,
    argon2: Argon2<'static>,
}

impl<S: UserStorage> AuthService<S> {
    /// Create a new authentication service with the given storage
    pub fn new(storage: S) -> Self {
        Self {
            storage: Arc::new(storage),
            argon2: Argon2::default(),
        }
    }

    /// Hash a password using Argon2
    pub fn hash_password(&self, password: &str) -> Result<String, AuthError> {
        let salt = SaltString::generate(&mut OsRng);
        let password_hash = self.argon2
            .hash_password(password.as_bytes(), &salt)
            .map_err(|e| AuthError::StorageError(format!("Password hashing failed: {}", e)))?;
        Ok(password_hash.to_string())
    }

    /// Check if a password matches the stored hash
    pub fn verify_password(&self, password: &str, hash: &str) -> Result<bool, AuthError> {
        let parsed_hash = PasswordHash::new(hash)
            .map_err(|e| AuthError::StorageError(format!("Invalid password hash: {}", e)))?;

        Ok(self.argon2.verify_password(password.as_bytes(), &parsed_hash).is_ok())
    }

    /// Create a new user account with the given details
    pub fn register_user(&self, request: RegisterRequest) -> Result<User, AuthError> {
        // Check if user already exists
        if self.storage.get_user(&request.username)?.is_some() {
            return Err(AuthError::StorageError("Username already exists".to_string()));
        }

        let password_hash = self.hash_password(&request.password)?;
        let role = request.role.unwrap_or_else(|| "user".to_string());

        // Verify role exists if not default
        if role != "user" && self.storage.get_role(&role)?.is_none() {
            return Err(AuthError::StorageError(format!("Role '{}' does not exist", role)));
        }

        let user = User {
            username: request.username,
            password_hash,
            role,
            permissions: vec![], // Role-based permissions by default
            created_at: current_timestamp(),
            last_login: None,
        };

        self.storage.create_user(user.clone())?;
        Ok(user)
    }

    /// Authenticate user and return user info
    pub fn authenticate(&self, request: &LoginRequest) -> Result<User, AuthError> {
        let username = match &request.username {
            Some(u) => u.clone(),
            None => {
                // Backward compatibility: try admin authentication
                return self.authenticate_admin(&request.password);
            }
        };

        let user = self.storage.get_user(&username)?
            .ok_or(AuthError::InvalidCredentials)?;

        if !self.verify_password(&request.password, &user.password_hash)? {
            return Err(AuthError::InvalidCredentials);
        }

        // Update last login
        let mut updated_user = user.clone();
        updated_user.last_login = Some(current_timestamp());
        self.storage.update_user(updated_user.clone())?;

        Ok(updated_user)
    }

    /// Backward compatibility: authenticate against admin password
    fn authenticate_admin(&self, password: &str) -> Result<User, AuthError> {
        let expected = std::env::var("ADMIN_PASSWORD").unwrap_or_default();
        if password == expected {
            // Create a virtual admin user for backward compatibility
            Ok(User {
                username: "admin".to_string(),
                password_hash: String::new(), // Not used in legacy mode
                role: "admin".to_string(),
                permissions: vec!["*".to_string()], // Full permissions
                created_at: current_timestamp(),
                last_login: Some(current_timestamp()),
            })
        } else {
            Err(AuthError::InvalidCredentials)
        }
    }

    /// Create or update a role
    pub fn manage_role(&self, role: Role) -> Result<(), AuthError> {
        if self.storage.get_role(&role.name)?.is_some() {
            self.storage.update_role(role)
        } else {
            self.storage.create_role(role)
        }
    }

    /// Check if user has specific permission
    pub fn check_permission(&self, username: &str, permission: &str) -> Result<bool, AuthError> {
        self.storage.user_has_permission(username, permission)
    }

    /// Get user by username
    pub fn get_user(&self, username: &str) -> Result<Option<User>, AuthError> {
        self.storage.get_user(username)
    }

    /// List all users (admin only)
    pub fn list_users(&self) -> Result<Vec<User>, AuthError> {
        self.storage.list_users()
    }

    /// List all roles
    pub fn list_roles(&self) -> Result<Vec<Role>, AuthError> {
        self.storage.list_roles()
    }

    /// Delete a user (admin only)
    pub fn delete_user(&self, username: &str) -> Result<(), AuthError> {
        self.storage.delete_user(username)
    }
}

/// Enhanced authentication guard with user context
pub struct MultiUserAuthGuard {
    pub username: String,
    pub role: String,
    pub permissions: Vec<String>,
}

impl MultiUserAuthGuard {
    /// Check if user has specific permission
    pub fn has_permission(&self, permission: &str) -> bool {
        self.permissions.contains(&permission.to_string()) ||
        self.permissions.contains(&"*".to_string())
    }

    /// Check if user has admin role (wildcard permission)
    pub fn is_admin(&self) -> bool {
        self.has_permission("*")
    }
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for MultiUserAuthGuard {
    type Error = AuthError;

    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, AuthError> {
        let token_store = match req.guard::<&State<MultiUserTokenStore>>().await {
            Outcome::Success(store) => store,
            _ => return Outcome::Error((Status::InternalServerError, AuthError::StorageError("Token store not available".to_string()))),
        };

        let auth_service = match req.guard::<&State<AuthService<InMemoryStorage>>>().await {
            Outcome::Success(service) => service,
            _ => return Outcome::Error((Status::InternalServerError, AuthError::StorageError("Auth service not available".to_string()))),
        };

        match req.headers().get_one("Authorization") {
            Some(auth_header) => {
                let token = if auth_header.starts_with("Bearer ") {
                    &auth_header[7..]
                } else {
                    auth_header
                };

                if let Some(username) = token_store.validate(token) {
                    match auth_service.get_user(&username) {
                        Ok(Some(user)) => {
                            // Get role permissions
                            let mut all_permissions = user.permissions.clone();
                            if let Ok(Some(role)) = auth_service.storage.get_role(&user.role) {
                                all_permissions.extend(role.permissions);
                            }

                            Outcome::Success(MultiUserAuthGuard {
                                username: user.username,
                                role: user.role,
                                permissions: all_permissions,
                            })
                        },
                        _ => Outcome::Error((Status::Unauthorized, AuthError::InvalidCredentials)),
                    }
                } else {
                    Outcome::Error((Status::Unauthorized, AuthError::TokenExpired))
                }
            },
            None => Outcome::Error((Status::Unauthorized, AuthError::TokenNotFound)),
        }
    }
}

/// Permission-based guard for specific permissions
pub struct RequirePermission(pub String);

#[rocket::async_trait]
impl<'r> FromRequest<'r> for RequirePermission {
    type Error = AuthError;

    async fn from_request(_req: &'r Request<'_>) -> Outcome<Self, AuthError> {
        // This is a placeholder - actual permission checking would be done in route handlers
        // using the MultiUserAuthGuard
        Outcome::Error((Status::Unauthorized, AuthError::StorageError("Use MultiUserAuthGuard with permission checking".to_string())))
    }
}

/// Enhanced login handler for multi-user authentication
pub fn multi_user_login_handler(
    body: Json<LoginRequest>,
    token_store: &State<MultiUserTokenStore>,
    auth_service: &State<AuthService<InMemoryStorage>>,
) -> Result<Json<TokenResponse>, Status> {
    match auth_service.authenticate(&body) {
        Ok(user) => {
            let token = Uuid::new_v4().to_string();
            token_store.insert(token.clone(), user.username.clone());

            Ok(Json(TokenResponse {
                token,
                username: Some(user.username),
                role: Some(user.role),
            }))
        },
        Err(_) => Err(Status::Unauthorized),
    }
}

/// User registration handler
pub fn register_handler(
    body: Json<RegisterRequest>,
    auth_service: &State<AuthService<InMemoryStorage>>,
    _admin: MultiUserAuthGuard, // Require authentication to register users
) -> Result<Json<User>, Status> {
    // Check if requester has admin permissions
    if !_admin.has_permission("user_management") && !_admin.is_admin() {
        return Err(Status::Forbidden);
    }

    match auth_service.register_user(body.into_inner()) {
        Ok(user) => Ok(Json(user)),
        Err(_) => Err(Status::BadRequest),
    }
}

fn get_token_ttl() -> Duration {
    let ttl_seconds = std::env::var("TOKEN_TTL_SECONDS")
        .unwrap_or_else(|_| DEFAULT_TOKEN_TTL_SECONDS.to_string())
        .parse::<u64>()
        .unwrap_or(DEFAULT_TOKEN_TTL_SECONDS);
    Duration::from_secs(ttl_seconds)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_password_hashing() {
        let storage = InMemoryStorage::new();
        let auth_service = AuthService::new(storage);

        let password = "test_password123";
        let hash = auth_service.hash_password(password).unwrap();

        assert!(auth_service.verify_password(password, &hash).unwrap());
        assert!(!auth_service.verify_password("wrong_password", &hash).unwrap());
    }

    #[test]
    fn test_user_registration() {
        let storage = InMemoryStorage::new().with_default_roles();
        let auth_service = AuthService::new(storage);

        let register_request = RegisterRequest {
            username: "testuser".to_string(),
            password: "password123".to_string(),
            role: Some("user".to_string()),
        };

        let user = auth_service.register_user(register_request).unwrap();
        assert_eq!(user.username, "testuser");
        assert_eq!(user.role, "user");
    }

    #[test]
    fn test_authentication() {
        let storage = InMemoryStorage::new().with_default_roles();
        let auth_service = AuthService::new(storage);

        // Register user first
        let register_request = RegisterRequest {
            username: "testuser".to_string(),
            password: "password123".to_string(),
            role: Some("user".to_string()),
        };
        auth_service.register_user(register_request).unwrap();

        // Test authentication
        let login_request = LoginRequest {
            username: Some("testuser".to_string()),
            password: "password123".to_string(),
        };

        let user = auth_service.authenticate(&login_request).unwrap();
        assert_eq!(user.username, "testuser");

        // Test wrong password
        let wrong_login = LoginRequest {
            username: Some("testuser".to_string()),
            password: "wrongpassword".to_string(),
        };

        assert!(auth_service.authenticate(&wrong_login).is_err());
    }

    #[test]
    fn test_backward_compatibility() {
        std::env::set_var("ADMIN_PASSWORD", "admin123");

        let storage = InMemoryStorage::new();
        let auth_service = AuthService::new(storage);

        // Test legacy login (no username)
        let legacy_login = LoginRequest {
            username: None,
            password: "admin123".to_string(),
        };

        let user = auth_service.authenticate(&legacy_login).unwrap();
        assert_eq!(user.username, "admin");
        assert!(user.has_permission("*"));
    }
}

impl User {
    /// Helper method to check permissions
    pub fn has_permission(&self, permission: &str) -> bool {
        self.permissions.contains(&permission.to_string()) ||
        self.permissions.contains(&"*".to_string())
    }
}