authforge/daemon/src/state.rs

use crate::fido::authenticator::Authenticator;
use crate::lockout::{simulate, UserEnrollment};
use crate::policy_apply::{ApplyError, PolicyApplier};
use crate::storage::credentials::{CredEntry, CredentialsStore, CredsError, CredsPathResolver};
use crate::storage::pending::{PendingError, PendingStore};
use crate::storage::policy::PolicyStore;
use crate::storage::recovery::{RecoveryEntry, RecoveryStore, RecoveryStoreError};
#[cfg(feature = "totp")]
use crate::storage::totp::TotpStore;
use crate::storage::userdb::{UserDb, UserDbError};
use authforge_common::policy::{Policy, PolicyError};
use authforge_common::types::{
    Credential, Method, PendingFlag, PendingStatus, PolicyApplyResult, Transport,
};
use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use thiserror::Error;
use tokio::sync::Mutex;

#[derive(Debug, Error)]
pub enum StateError {
    #[error(transparent)]
    Policy(#[from] PolicyError),
    #[error(transparent)]
    Pending(#[from] PendingError),
    #[error(transparent)]
    Recovery(#[from] RecoveryStoreError),
    #[error(transparent)]
    Creds(#[from] CredsError),
    #[error(transparent)]
    UserDb(#[from] UserDbError),
    #[error(transparent)]
    Apply(#[from] ApplyError),
    #[error("authenticator: {0}")]
    Authn(String),
    #[allow(dead_code)] // wired through D-Bus in Task 5.
    #[error("storage: {0}")]
    Storage(String),
}

pub struct StorageConfig {
    pub policy_dir: PathBuf,
    pub pending_dir: PathBuf,
    pub recovery_dir: PathBuf,
    #[cfg(feature = "totp")]
    pub totp_dir: PathBuf,
    pub userdb_path: PathBuf,
    /// Where pam-auth-update reads the AuthForge profile from. Defaults to
    /// `/usr/share/pam-configs/authforge`; tests / non-root runs override.
    pub pam_profile_path: PathBuf,
    /// Path to the pam-auth-update binary. Tests / non-root runs override
    /// with a no-op shim.
    pub pam_auth_update: PathBuf,
}

impl StorageConfig {
    pub fn from_env_or_defaults() -> Self {
        let env = |k: &str, d: &str| -> PathBuf {
            std::env::var_os(k)
                .map(PathBuf::from)
                .unwrap_or_else(|| PathBuf::from(d))
        };
        Self {
            policy_dir: env("AUTHFORGE_POLICY_DIR", "/etc/authforge/policy.d"),
            pending_dir: env("AUTHFORGE_PENDING_DIR", "/var/lib/authforge/pending"),
            recovery_dir: env("AUTHFORGE_RECOVERY_DIR", "/var/lib/authforge/recovery"),
            #[cfg(feature = "totp")]
            totp_dir: env("AUTHFORGE_TOTP_DIR", "/etc/google-authenticator"),
            userdb_path: env("AUTHFORGE_USERDB", "/var/lib/authforge/users.db"),
            pam_profile_path: env("AUTHFORGE_PAMCONF_PATH", "/usr/share/pam-configs/authforge"),
            pam_auth_update: env("AUTHFORGE_PAM_AUTH_UPDATE", "/usr/sbin/pam-auth-update"),
        }
    }
}

pub struct AppState {
    policy: PolicyStore,
    pending: PendingStore,
    #[allow(dead_code)] // wired through D-Bus in Task 6.
    recovery: RecoveryStore,
    #[cfg(feature = "totp")]
    #[allow(dead_code)] // wired through D-Bus in Task 5.
    totp: TotpStore,
    userdb: Mutex<UserDb>,
    authn: Arc<dyn Authenticator>,
    applier: PolicyApplier,
}

impl AppState {
    pub fn open(cfg: StorageConfig, authn: Arc<dyn Authenticator>) -> Result<Self, StateError> {
        let policy = PolicyStore::new(cfg.policy_dir);
        let pending = PendingStore::new(cfg.pending_dir);
        let recovery = RecoveryStore::new(cfg.recovery_dir);
        #[cfg(feature = "totp")]
        let totp = TotpStore::new(cfg.totp_dir);
        let userdb = Mutex::new(UserDb::open(&cfg.userdb_path)?);
        let applier = PolicyApplier::new(cfg.pam_profile_path, cfg.pam_auth_update);
        Ok(Self {
            policy,
            pending,
            recovery,
            #[cfg(feature = "totp")]
            totp,
            userdb,
            authn,
            applier,
        })
    }

    fn creds_resolver(&self) -> Result<CredsPathResolver, StateError> {
        let pol = self.policy.load()?;
        Ok(CredsPathResolver::new(pol.storage))
    }

    pub async fn list_credentials(&self, user: &str) -> Vec<Credential> {
        let resolver = match self.creds_resolver() {
            Ok(r) => r,
            Err(_) => return vec![],
        };
        let path = match resolver.path_for(user) {
            Ok(p) => p,
            Err(_) => return vec![],
        };
        let store = CredentialsStore::new(path);
        let lines = store.list(user).unwrap_or_default();
        lines
            .iter()
            .map(|c| {
                let id = CredEntry::cred_id(c).to_string();
                Credential {
                    id: id.clone(),
                    nickname: id,
                    method: Method::Fido2,
                    transport: Transport::Unknown,
                    created_unix: 0,
                }
            })
            .collect()
    }

    /// Run the authenticator's `make_credential` and persist the resulting
    /// pam_u2f line + userdb entry. Returns the Credential the GUI displays.
    pub async fn enroll(&self, user: &str, nickname: &str) -> Result<Credential, StateError> {
        let resolver = self.creds_resolver()?;
        let path = resolver.path_for(user)?;
        let store = CredentialsStore::new(path);

        let pam_cred = self
            .authn
            .make_credential("pam://localhost", user, None)
            .map_err(|e| StateError::Authn(e.to_string()))?;
        let segment = pam_cred.to_pam_segment();
        let cred_id = hex::encode(&pam_cred.key_handle);

        store.add(user, &segment)?;
        let db = self.userdb.lock().await;
        db.record_enrollment(user, Method::Fido2)?;

        Ok(Credential {
            id: cred_id,
            nickname: nickname.to_string(),
            method: Method::Fido2,
            transport: Transport::Usb,
            created_unix: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.as_secs())
                .unwrap_or(0),
        })
    }

    pub async fn remove_credential(&self, user: &str, cred_id: &str) -> Result<bool, StateError> {
        let resolver = self.creds_resolver()?;
        let path = resolver.path_for(user)?;
        let store = CredentialsStore::new(path);
        let removed = store.remove(user, cred_id)?;
        if removed && store.list(user)?.is_empty() {
            self.userdb
                .lock()
                .await
                .drop_enrollment(user, Method::Fido2)?;
        }
        Ok(removed)
    }

    pub async fn get_policy(&self) -> Result<Policy, StateError> {
        Ok(self.policy.load()?)
    }

    /// Apply a new policy. Runs the lockout simulator first; if any users
    /// would be locked out and `force == false`, returns the violations in
    /// `PolicyApplyResult { applied: false, violations }` without writing or
    /// invoking pam-auth-update. With `force == true` (or no violations),
    /// persists the policy and runs pam-auth-update via PolicyApplier; on
    /// pam-auth-update failure the prior profile is restored.
    pub async fn set_policy(
        &self,
        p: Policy,
        force: bool,
    ) -> Result<PolicyApplyResult, StateError> {
        let registry = self.build_enrollment_registry().await?;
        let violations = simulate(&p, &registry);
        if !violations.is_empty() && !force {
            return Ok(PolicyApplyResult {
                applied: false,
                violations,
            });
        }
        self.policy.save(&p)?;
        // pam-auth-update is best-effort in dev; the AUTHFORGE_PAM_AUTH_UPDATE
        // env var lets tests / non-prod runs swap in a no-op shim. If the
        // resolved binary doesn't exist on the host, log + continue.
        if let Err(e) = self.applier.apply(&p) {
            tracing::warn!(error=%e, "pam-auth-update apply failed; profile rolled back");
            return Err(StateError::Apply(e));
        }
        Ok(PolicyApplyResult {
            applied: true,
            violations,
        })
    }

    async fn build_enrollment_registry(&self) -> Result<Vec<UserEnrollment>, StateError> {
        let db = self.userdb.lock().await;
        let fido2 = db.users_with(Method::Fido2)?;
        let totp = db.users_with(Method::Totp)?;
        drop(db);

        let mut by_user: std::collections::HashMap<String, HashSet<Method>> =
            std::collections::HashMap::new();
        for u in fido2 {
            by_user.entry(u).or_default().insert(Method::Fido2);
        }
        for u in totp {
            by_user.entry(u).or_default().insert(Method::Totp);
        }
        Ok(by_user
            .into_iter()
            .map(|(user, methods)| UserEnrollment { user, methods })
            .collect())
    }

    pub async fn set_pending(&self, user: &str, f: PendingFlag) -> Result<(), StateError> {
        self.pending.set(user, &f)?;
        Ok(())
    }

    pub async fn clear_pending(&self, user: &str) -> Result<bool, StateError> {
        Ok(self.pending.clear(user)?)
    }

    #[allow(dead_code)] // wired through D-Bus in Task 6.
    pub async fn issue_recovery(&self, user: &str) -> Result<String, StateError> {
        Ok(self
            .recovery
            .issue(user, crate::recovery::DEFAULT_TTL_SECS)?)
    }

    #[allow(dead_code)] // wired through D-Bus in Task 6.
    pub async fn list_recovery(&self) -> Result<Vec<RecoveryEntry>, StateError> {
        Ok(self.recovery.list()?)
    }

    #[allow(dead_code)] // wired through D-Bus in Task 6.
    pub async fn revoke_recovery(&self, user: &str) -> Result<bool, StateError> {
        Ok(self.recovery.revoke(user)?)
    }

    /// Tests assert pending flag round-trips through SetPendingFlag /
    /// ClearPendingFlag via this read accessor. Production readers go
    /// through `get_pending_status` (GUI) or read the on-disk file
    /// directly (PAM module).
    #[cfg(test)]
    pub async fn has_pending(&self, user: &str) -> Result<bool, StateError> {
        Ok(self.pending.get(user)?.is_some())
    }

    pub async fn get_pending_status(&self, user: &str) -> Result<PendingStatus, StateError> {
        match self.pending.get(user)? {
            Some(flag) => Ok(PendingStatus {
                present: true,
                flag,
            }),
            None => Ok(PendingStatus {
                present: false,
                flag: PendingFlag::default(),
            }),
        }
    }
}

#[cfg(feature = "totp")]
#[allow(dead_code)] // wired through D-Bus in Task 5.
impl AppState {
    pub async fn enroll_totp(
        &self,
        user: &str,
    ) -> Result<authforge_common::types::TotpEnrollment, StateError> {
        let inner = self
            .totp
            .enroll(user, "AuthForge")
            .map_err(|e| StateError::Storage(e.to_string()))?;
        Ok(authforge_common::types::TotpEnrollment {
            user: inner.user,
            secret_b32: inner.secret_b32,
            otpauth_uri: inner.otpauth_uri,
        })
    }

    pub async fn is_totp_enrolled(&self, user: &str) -> Result<bool, StateError> {
        self.totp
            .is_enrolled(user)
            .map_err(|e| StateError::Storage(e.to_string()))
    }

    pub async fn revoke_totp(&self, user: &str) -> Result<bool, StateError> {
        self.totp
            .revoke(user)
            .map_err(|e| StateError::Storage(e.to_string()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fido::mock::MockAuthenticator;
    use tempfile::tempdir;

    fn open_in(d: &std::path::Path) -> AppState {
        // Stage a no-op pam-auth-update shim inside the tempdir so SetPolicy
        // calls succeed without touching /usr/share/pam-configs.
        let shim = d.join("fake-pau.sh");
        std::fs::write(&shim, "#!/bin/sh\nexit 0\n").unwrap();
        let mut perms = std::fs::metadata(&shim).unwrap().permissions();
        std::os::unix::fs::PermissionsExt::set_mode(&mut perms, 0o755);
        std::fs::set_permissions(&shim, perms).unwrap();

        AppState::open(
            StorageConfig {
                policy_dir: d.join("policy.d"),
                pending_dir: d.join("pending"),
                recovery_dir: d.join("recovery"),
                #[cfg(feature = "totp")]
                totp_dir: d.join("totp"),
                userdb_path: d.join("users.db"),
                pam_profile_path: d.join("authforge-pamconf"),
                pam_auth_update: shim,
            },
            Arc::new(MockAuthenticator::with_one_yubikey()),
        )
        .unwrap()
    }

    #[tokio::test]
    async fn empty_state_has_nothing() {
        let d = tempdir().unwrap();
        let s = open_in(d.path());
        assert!(s.list_credentials("alice").await.is_empty());
        assert!(!s.has_pending("alice").await.unwrap());
        assert!(s.get_policy().await.unwrap().stacks.is_empty());
    }

    #[tokio::test]
    async fn pending_set_clear_roundtrip() {
        let d = tempdir().unwrap();
        let s = open_in(d.path());
        let f = PendingFlag {
            required_methods: vec![Method::Fido2],
            created_unix: 1,
            deadline_unix: 0,
            re_enroll: false,
        };
        s.set_pending("alice", f).await.unwrap();
        assert!(s.has_pending("alice").await.unwrap());
        assert!(s.clear_pending("alice").await.unwrap());
        assert!(!s.has_pending("alice").await.unwrap());
    }

    #[tokio::test]
    async fn enroll_writes_real_pam_u2f_line() {
        let d = tempdir().unwrap();
        // Configure central storage so the resolver doesn't try /home/alice.
        let policy_dir = d.path().join("policy.d");
        std::fs::create_dir_all(&policy_dir).unwrap();
        let central = d.path().join("u2f_keys");
        std::fs::write(
            policy_dir.join("00.conf"),
            format!(
                "[storage]\nbackend = \"central\"\ncentral_path = \"{}\"\n",
                central.display()
            ),
        )
        .unwrap();
        let shim = d.path().join("fake-pau.sh");
        std::fs::write(&shim, "#!/bin/sh\nexit 0\n").unwrap();
        let mut perms = std::fs::metadata(&shim).unwrap().permissions();
        std::os::unix::fs::PermissionsExt::set_mode(&mut perms, 0o755);
        std::fs::set_permissions(&shim, perms).unwrap();

        let s = AppState::open(
            StorageConfig {
                policy_dir,
                pending_dir: d.path().join("pending"),
                recovery_dir: d.path().join("recovery"),
                #[cfg(feature = "totp")]
                totp_dir: d.path().join("totp"),
                userdb_path: d.path().join("users.db"),
                pam_profile_path: d.path().join("authforge-pamconf"),
                pam_auth_update: shim,
            },
            Arc::new(MockAuthenticator::with_one_yubikey()),
        )
        .unwrap();

        let cred = s.enroll("alice", "Yellow").await.unwrap();
        // cred.id should be hex-encoded keyHandle.
        assert!(cred.id.chars().all(|c| c.is_ascii_hexdigit()));
        assert_eq!(cred.nickname, "Yellow");
        assert_eq!(cred.method, Method::Fido2);

        // The on-disk pam_u2f file holds a real line — username:kh,pk,es256,+presence
        let body = std::fs::read_to_string(&central).unwrap();
        assert!(body.starts_with("alice:"));
        assert!(body.contains(",es256,+presence"));
    }
}