boc/aamos-ledger-rust/crates/aamos-auth/src/lib.rs

use aamos_core::{LedgerError, LedgerResult, TenantId};
use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;

// ── Roller ────────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Role {
    Admin,
    Accountant,
    Viewer,
    Auditor,
    Payroll,
}

impl Role {
    pub fn as_str(&self) -> &'static str {
        match self {
            Role::Admin => "admin",
            Role::Accountant => "accountant",
            Role::Viewer => "viewer",
            Role::Auditor => "auditor",
            Role::Payroll => "payroll",
        }
    }
}

impl std::str::FromStr for Role {
    type Err = LedgerError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "admin" => Ok(Role::Admin),
            "accountant" => Ok(Role::Accountant),
            "viewer" => Ok(Role::Viewer),
            "auditor" => Ok(Role::Auditor),
            "payroll" => Ok(Role::Payroll),
            _ => Err(LedgerError::Validation(format!("Unknown role: {}", s))),
        }
    }
}

// ── JWT Claims ────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Claims {
    pub sub: String,           // user_id
    pub tenant_id: String,
    pub roles: Vec<String>,
    pub exp: usize,
    pub iat: usize,
}

// ── AuthUser ──────────────────────────────────────────────────────────────
#[derive(Debug, Clone)]
pub struct AuthUser {
    pub user_id: String,
    pub tenant_id: TenantId,
    pub roles: HashSet<Role>,
}

impl AuthUser {
    pub fn has_role(&self, role: Role) -> bool {
        self.roles.contains(&role)
    }
    pub fn has_any_role(&self, roles: &[Role]) -> bool {
        roles.iter().any(|r| self.roles.contains(r))
    }
    pub fn require_role(&self, roles: &[Role]) -> LedgerResult<()> {
        if self.has_any_role(roles) {
            Ok(())
        } else {
            Err(LedgerError::Forbidden(format!(
                "Requires one of: {}",
                roles.iter().map(|r| r.as_str()).collect::<Vec<_>>().join(", ")
            )))
        }
    }
}

// ── JWT-validering ────────────────────────────────────────────────────────
pub struct JwtValidator {
    decoding_key: DecodingKey,
    validation: Validation,
}

impl JwtValidator {
    pub fn from_secret(secret: &str) -> Self {
        let mut validation = Validation::new(Algorithm::HS256);
        validation.set_required_spec_claims(&["exp", "sub", "tenant_id"]);
        Self {
            decoding_key: DecodingKey::from_secret(secret.as_bytes()),
            validation,
        }
    }

    pub fn from_rsa_pem(pem: &[u8]) -> Result<Self, jsonwebtoken::errors::Error> {
        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_required_spec_claims(&["exp", "sub", "tenant_id"]);
        Ok(Self {
            decoding_key: DecodingKey::from_rsa_pem(pem)?,
            validation,
        })
    }

    pub fn validate(&self, token: &str) -> LedgerResult<AuthUser> {
        let token_data = decode::<Claims>(token, &self.decoding_key, &self.validation)
            .map_err(|e| LedgerError::Unauthorized(e.to_string()))?;

        let claims = token_data.claims;
        let roles: HashSet<Role> = claims
            .roles
            .iter()
            .filter_map(|r| r.parse().ok())
            .collect();

        Ok(AuthUser {
            user_id: claims.sub,
            tenant_id: TenantId::new(claims.tenant_id),
            roles,
        })
    }
}

// ── Tester ────────────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
    use super::*;
    use jsonwebtoken::{encode, EncodingKey, Header};

    #[test]
    fn test_role_from_str() {
        assert_eq!("admin".parse::<Role>().unwrap(), Role::Admin);
        assert_eq!("accountant".parse::<Role>().unwrap(), Role::Accountant);
        assert_eq!("viewer".parse::<Role>().unwrap(), Role::Viewer);
        assert_eq!("auditor".parse::<Role>().unwrap(), Role::Auditor);
        assert_eq!("payroll".parse::<Role>().unwrap(), Role::Payroll);
        assert!("unknown".parse::<Role>().is_err());
    }

    #[test]
    fn test_role_as_str() {
        assert_eq!(Role::Admin.as_str(), "admin");
        assert_eq!(Role::Accountant.as_str(), "accountant");
        assert_eq!(Role::Viewer.as_str(), "viewer");
        assert_eq!(Role::Auditor.as_str(), "auditor");
        assert_eq!(Role::Payroll.as_str(), "payroll");
    }

    #[test]
    fn test_auth_user_has_role() {
        let user = AuthUser {
            user_id: "user-1".to_string(),
            tenant_id: TenantId::new("tenant-1"),
            roles: {
                let mut set = HashSet::new();
                set.insert(Role::Admin);
                set.insert(Role::Viewer);
                set
            },
        };

        assert!(user.has_role(Role::Admin));
        assert!(user.has_role(Role::Viewer));
        assert!(!user.has_role(Role::Accountant));
    }

    #[test]
    fn test_auth_user_has_any_role() {
        let user = AuthUser {
            user_id: "user-1".to_string(),
            tenant_id: TenantId::new("tenant-1"),
            roles: {
                let mut set = HashSet::new();
                set.insert(Role::Viewer);
                set
            },
        };

        assert!(user.has_any_role(&[Role::Viewer, Role::Admin]));
        assert!(!user.has_any_role(&[Role::Admin, Role::Accountant]));
    }

    #[test]
    fn test_auth_user_require_role_ok() {
        let user = AuthUser {
            user_id: "user-1".to_string(),
            tenant_id: TenantId::new("tenant-1"),
            roles: {
                let mut set = HashSet::new();
                set.insert(Role::Admin);
                set
            },
        };

        assert!(user.require_role(&[Role::Admin]).is_ok());
        assert!(user.require_role(&[Role::Admin, Role::Viewer]).is_ok());
    }

    #[test]
    fn test_auth_user_require_role_err() {
        let user = AuthUser {
            user_id: "user-1".to_string(),
            tenant_id: TenantId::new("tenant-1"),
            roles: {
                let mut set = HashSet::new();
                set.insert(Role::Viewer);
                set
            },
        };

        let result = user.require_role(&[Role::Admin]);
        assert!(result.is_err());
        match result {
            Err(LedgerError::Forbidden(msg)) => {
                assert!(msg.contains("admin"));
            }
            _ => panic!("Expected Forbidden error"),
        }
    }

    #[test]
    fn test_jwt_validator_from_secret() {
        let validator = JwtValidator::from_secret("my-secret");
        // Just verify it constructs without panic
        let _ = validator;
    }

    #[test]
    fn test_jwt_validator_validate_valid_token() {
        let secret = "test-secret-key";
        let validator = JwtValidator::from_secret(secret);

        let claims = Claims {
            sub: "user-123".to_string(),
            tenant_id: "tenant-abc".to_string(),
            roles: vec!["admin".to_string(), "viewer".to_string()],
            exp: usize::MAX,
            iat: 0,
        };

        let token = encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(secret.as_bytes()),
        )
        .unwrap();

        let auth_user = validator.validate(&token).unwrap();
        assert_eq!(auth_user.user_id, "user-123");
        assert_eq!(auth_user.tenant_id.0, "tenant-abc");
        assert!(auth_user.has_role(Role::Admin));
        assert!(auth_user.has_role(Role::Viewer));
        assert!(!auth_user.has_role(Role::Accountant));
    }

    #[test]
    fn test_jwt_validator_validate_invalid_signature() {
        let validator = JwtValidator::from_secret("correct-secret");

        let token = "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJ1c2VyLTEiLCJ0ZW5hbnRfaWQiOiJ0ZW5hbnQtMSIsInJvbGVzIjpbImFkbWluIl0sImV4cCI6OTk5OTk5OTk5OSwiaWF0IjowfQ.wrong_signature";

        let result = validator.validate(token);
        assert!(result.is_err());
    }

    #[test]
    fn test_jwt_validator_validate_expired_token() {
        let secret = "test-secret-key";
        let validator = JwtValidator::from_secret(secret);

        let claims = Claims {
            sub: "user-123".to_string(),
            tenant_id: "tenant-abc".to_string(),
            roles: vec!["admin".to_string()],
            exp: 0, // expired
            iat: 0,
        };

        let token = encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(secret.as_bytes()),
        )
        .unwrap();

        let result = validator.validate(&token);
        assert!(result.is_err());
    }

    #[test]
    fn test_jwt_validator_validate_unknown_roles_ignored() {
        let secret = "test-secret-key";
        let validator = JwtValidator::from_secret(secret);

        let claims = Claims {
            sub: "user-123".to_string(),
            tenant_id: "tenant-abc".to_string(),
            roles: vec!["admin".to_string(), "superuser".to_string(), "viewer".to_string()],
            exp: usize::MAX,
            iat: 0,
        };

        let token = encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(secret.as_bytes()),
        )
        .unwrap();

        let auth_user = validator.validate(&token).unwrap();
        assert!(auth_user.has_role(Role::Admin));
        assert!(auth_user.has_role(Role::Viewer));
        // "superuser" is unknown and silently ignored
        assert_eq!(auth_user.roles.len(), 2);
    }
}