PostfixAdminRS/src/auth.rs

use std::{
    collections::VecDeque,
    net::IpAddr,
    sync::Arc,
    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};

use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use dashmap::DashMap;
use data_encoding::BASE32_NOPAD;
use hmac::{Hmac, Mac};
use rand::RngCore;
use sha1::Sha1;
use subtle::ConstantTimeEq;

use crate::{
    config::AppConfig,
    error::{AppError, AppResult},
    models::{AccountKind, Role},
};

#[derive(Clone)]
pub struct Session {
    pub username: String,
    pub kind: AccountKind,
    pub role: Role,
    pub allowed_domains: Arc<[String]>,
    pub csrf: String,
    mail_password: Option<Arc<str>>,
    expires_at: Instant,
}

impl Session {
    pub fn can_manage_domain(&self, domain: &str) -> bool {
        self.role.is_global()
            || self
                .allowed_domains
                .iter()
                .any(|d| d == domain || d == "ALL")
    }

    pub fn require_admin(&self) -> AppResult<()> {
        if self.role.is_admin() {
            Ok(())
        } else {
            Err(AppError::forbidden("Administrator role required"))
        }
    }

    pub fn require_global(&self) -> AppResult<()> {
        if self.role.is_global() {
            Ok(())
        } else {
            Err(AppError::forbidden("Global administrator role required"))
        }
    }

    pub fn require_domain(&self, domain: &str) -> AppResult<()> {
        self.require_admin()?;
        if self.can_manage_domain(domain) {
            Ok(())
        } else {
            Err(AppError::forbidden(
                "Domain is outside this administrator's scope",
            ))
        }
    }

    pub fn mail_password(&self) -> Option<&str> {
        self.mail_password.as_deref()
    }
}

#[derive(Clone)]
pub struct PendingLogin {
    pub username: String,
    pub kind: AccountKind,
    pub role: Role,
    pub allowed_domains: Arc<[String]>,
    pub ip: IpAddr,
    pub mail_password: Option<Arc<str>>,
    pub expires_at: Instant,
}

#[derive(Clone)]
pub struct AuthStore {
    sessions: Arc<DashMap<String, Session>>,
    pending: Arc<DashMap<String, PendingLogin>>,
    attempts: Arc<DashMap<String, VecDeque<Instant>>>,
    timeout: Duration,
    login_window: Duration,
    max_attempts: usize,
}

impl AuthStore {
    pub fn new(config: &AppConfig) -> Self {
        Self {
            sessions: Arc::new(DashMap::new()),
            pending: Arc::new(DashMap::new()),
            attempts: Arc::new(DashMap::new()),
            timeout: Duration::from_secs(
                config.security.session_timeout_minutes.saturating_mul(60),
            ),
            login_window: Duration::from_secs(config.security.login_window_seconds),
            max_attempts: config.security.login_attempts,
        }
    }

    pub fn create_session(
        &self,
        username: String,
        kind: AccountKind,
        role: Role,
        allowed_domains: Arc<[String]>,
        mail_password: Option<Arc<str>>,
    ) -> (String, Session) {
        let token = random_token(32);
        let session = Session {
            username,
            kind,
            role,
            allowed_domains,
            csrf: random_token(32),
            mail_password,
            expires_at: Instant::now() + self.timeout,
        };
        self.sessions.insert(token.clone(), session.clone());
        (token, session)
    }

    pub fn get_session(&self, token: &str) -> Option<Session> {
        let item = self.sessions.get(token)?;
        if Instant::now() >= item.expires_at {
            drop(item);
            self.sessions.remove(token);
            return None;
        }
        Some(item.clone())
    }

    pub fn revoke(&self, token: &str) {
        self.sessions.remove(token);
    }

    pub fn update_mail_password(&self, token: &str, password: String) {
        if let Some(mut session) = self.sessions.get_mut(token)
            && session.kind == AccountKind::User
        {
            session.mail_password = Some(Arc::<str>::from(password));
        }
    }

    pub fn create_pending(
        &self,
        username: String,
        kind: AccountKind,
        role: Role,
        allowed_domains: Arc<[String]>,
        ip: IpAddr,
        mail_password: Option<Arc<str>>,
    ) -> String {
        let token = random_token(32);
        self.pending.insert(
            token.clone(),
            PendingLogin {
                username,
                kind,
                role,
                allowed_domains,
                ip,
                mail_password,
                expires_at: Instant::now() + Duration::from_secs(300),
            },
        );
        token
    }

    pub fn take_pending(&self, token: &str, ip: IpAddr) -> Option<PendingLogin> {
        let (_, pending) = self.pending.remove(token)?;
        if pending.ip != ip || Instant::now() >= pending.expires_at {
            return None;
        }
        Some(pending)
    }

    pub fn login_allowed(&self, ip: IpAddr, username: &str) -> bool {
        let now = Instant::now();
        let key = format!("{ip}|{}", username.to_ascii_lowercase());
        let mut entry = self.attempts.entry(key).or_default();
        while entry
            .front()
            .is_some_and(|t| now.duration_since(*t) > self.login_window)
        {
            entry.pop_front();
        }
        if entry.len() >= self.max_attempts {
            return false;
        }
        entry.push_back(now);
        true
    }

    pub fn clear_attempts(&self, ip: IpAddr, username: &str) {
        self.attempts
            .remove(&format!("{ip}|{}", username.to_ascii_lowercase()));
    }
}

pub fn verify_csrf(session: &Session, provided: Option<&str>) -> AppResult<()> {
    let Some(provided) = provided else {
        return Err(AppError::forbidden("Missing CSRF token"));
    };
    let a = session.csrf.as_bytes();
    let b = provided.as_bytes();
    if a.len() == b.len() && bool::from(a.ct_eq(b)) {
        Ok(())
    } else {
        Err(AppError::forbidden("Invalid CSRF token"))
    }
}

pub fn verify_totp(secret: &str, code: &str) -> bool {
    if code.len() != 6 || !code.bytes().all(|b| b.is_ascii_digit()) {
        return false;
    }
    let normalized = secret.trim().replace([' ', '-'], "").to_ascii_uppercase();
    let Ok(key) = BASE32_NOPAD.decode(normalized.as_bytes()) else {
        return false;
    };
    let now = match SystemTime::now().duration_since(UNIX_EPOCH) {
        Ok(v) => v.as_secs(),
        Err(_) => return false,
    };
    let target = code.parse::<u32>().unwrap_or(u32::MAX);
    for offset in [-1i64, 0, 1] {
        let counter = (now / 30) as i64 + offset;
        if counter < 0 {
            continue;
        }
        if totp_code(&key, counter as u64) == target {
            return true;
        }
    }
    false
}

fn totp_code(key: &[u8], counter: u64) -> u32 {
    let mut mac = Hmac::<Sha1>::new_from_slice(key).expect("HMAC accepts arbitrary key size");
    mac.update(&counter.to_be_bytes());
    let digest = mac.finalize().into_bytes();
    let offset = (digest[19] & 0x0f) as usize;
    let binary = ((u32::from(digest[offset]) & 0x7f) << 24)
        | (u32::from(digest[offset + 1]) << 16)
        | (u32::from(digest[offset + 2]) << 8)
        | u32::from(digest[offset + 3]);
    binary % 1_000_000
}

pub fn random_token(bytes: usize) -> String {
    let mut buf = vec![0u8; bytes];
    rand::rng().fill_bytes(&mut buf);
    URL_SAFE_NO_PAD.encode(buf)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn rfc6238_sha1_vector_internal() {
        let key = b"12345678901234567890";
        assert_eq!(totp_code(key, 59 / 30), 287082);
    }
}