loa_core/

identity.rs

//! Agent Identity Management
//!
//! Handles Ed25519 keypairs for signing and X25519 keypairs for encryption.
//!
//! # Key Usage
//! - **Ed25519**: Currently used for heartbeat authentication (signing messages)
//! - **X25519**: Reserved for future encrypted agent-to-agent communication
//!   (encrypted metrics, encrypted commands, etc.). The API stores both public
//!   keys but only validates Ed25519 signatures currently.

use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use crate::{Error, Result};
use libp2p_identity::{ed25519, Keypair, PeerId};
use once_cell::sync::OnceCell;
use rand::rngs::OsRng;
use serde::{Deserialize, Serialize};
use std::{fs, path::Path, path::PathBuf, sync::Arc};
use x25519_dalek::{PublicKey as X25519PublicKey, StaticSecret as X25519SecretKey};

/// Filename for consolidated agent data (new format)
const AGENT_TOML_FILE: &str = "agent.toml";
/// Filename for legacy key storage (old format)
const LEGACY_KEY_FILE: &str = "agent_id.key";
/// Filename for legacy claim token (old format)
const LEGACY_CLAIM_TOKEN_FILE: &str = "claim_token";
/// Filename for legacy agent name (old format)
const LEGACY_AGENT_NAME_FILE: &str = "agent_name";

/// Lightweight agent information that can be read without starting the runtime.
///
/// Use this for CLI info commands, status checks, or any scenario where you need
/// agent metadata without spawning the full actor system.
#[derive(Debug, Clone)]
pub struct AgentInfo {
    /// Unique agent identifier derived from Ed25519 public key
    pub peer_id: PeerId,
    /// Human-readable 3-word name (e.g., "conscious-jade-mongoose")
    /// None if agent hasn't registered with the server yet
    pub name: Option<String>,
    /// Path to agent storage directory
    pub storage_path: PathBuf,
    /// Dashboard URL for this agent
    pub dashboard_url: String,
    /// Ed25519 public key in hex format
    pub ed25519_public_key_hex: String,
    /// X25519 public key in hex format (for future encrypted communication)
    pub x25519_public_key_hex: String,
    /// Workspace claim token if agent is claimed
    pub claim_token: Option<String>,
}

impl AgentInfo {
    /// Read agent info from storage without starting the runtime.
    ///
    /// This is a lightweight operation that only reads files from disk.
    /// No actors, databases, or network connections are started.
    ///
    /// Tries TOML format first (`agent.toml`), then falls back to legacy
    /// separate files (`agent_id.key`, `claim_token`, `agent_name`).
    ///
    /// # Arguments
    ///
    /// * `storage_path` - Path to the agent's storage directory
    ///
    /// # Returns
    ///
    /// Returns `AgentInfo` if the identity file exists, or an error if the
    /// agent has not been initialized yet.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use elo::AgentInfo;
    /// use std::path::Path;
    ///
    /// let info = AgentInfo::read(Path::new("/var/lib/loa"))?;
    /// println!("Agent: {} ({})", info.name, info.peer_id);
    /// println!("Dashboard: {}", info.dashboard_url);
    /// # Ok::<(), elo::Error>(())
    /// ```
    pub fn read(storage_path: &Path) -> Result<Self> {
        let toml_path = storage_path.join(AGENT_TOML_FILE);
        let legacy_path = storage_path.join(LEGACY_KEY_FILE);

        if !toml_path.exists() && !legacy_path.exists() {
            return Err(Error::Config(format!(
                "Agent not initialized. No identity found at {} or {}",
                toml_path.display(),
                legacy_path.display()
            )));
        }

        // Try to read registration data from TOML first
        let (name, claim_token) = if toml_path.exists() {
            Self::read_registration_from_toml(&toml_path)?
        } else {
            // Fall back to legacy separate files
            let name = AgentIdentity::read_legacy_file(storage_path, LEGACY_AGENT_NAME_FILE);
            let claim_token = AgentIdentity::read_legacy_file(storage_path, LEGACY_CLAIM_TOKEN_FILE);
            (name, claim_token)
        };

        // Load identity (handles TOML vs legacy automatically)
        let identity = AgentIdentity::from_storage_or_generate_new(storage_path)?;
        let peer_id = identity.peer_id().clone();

        // Build dashboard URL
        let api_url = crate::constants::api_url();
        let dashboard_url = if api_url.contains("localhost") {
            format!("http://localhost:3000/agents/{}", peer_id)
        } else if api_url.contains("workers.dev") {
            // Dev environment
            format!("https://loa-web.pages.dev/agents/{}", peer_id)
        } else {
            // Production
            format!("https://loa.sh/agents/{}", peer_id)
        };

        Ok(Self {
            peer_id,
            name,
            storage_path: storage_path.to_path_buf(),
            dashboard_url,
            ed25519_public_key_hex: identity.ed25519_public_key_hex().to_string(),
            x25519_public_key_hex: identity.x25519_public_key_hex().to_string(),
            claim_token,
        })
    }

    /// Read registration section from TOML file
    fn read_registration_from_toml(toml_path: &Path) -> Result<(Option<String>, Option<String>)> {
        let toml_content = fs::read_to_string(toml_path)
            .map_err(|e| Error::Io(std::io::Error::new(e.kind(), format!("read agent.toml: {}", e))))?;
        let agent_toml: AgentToml = toml::from_str(&toml_content)
            .map_err(|e| Error::Config(format!("parse agent.toml: {}", e)))?;

        let (name, claim_token) = match agent_toml.registration {
            Some(reg) => (reg.name, reg.claim_token),
            None => (None, None),
        };

        Ok((name, claim_token))
    }

    /// Check if the agent identity exists at the given storage path.
    ///
    /// This is a quick check that doesn't load the identity file.
    /// Checks for both TOML and legacy formats.
    pub fn exists(storage_path: &Path) -> bool {
        storage_path.join(AGENT_TOML_FILE).exists()
            || storage_path.join(LEGACY_KEY_FILE).exists()
    }
}

/// Global agent identity singleton
///
/// Initialized once during agent startup and accessible from all actors.
static AGENT_IDENTITY: OnceCell<Arc<AgentIdentity>> = OnceCell::new();

/// Initialize the global agent identity
///
/// This should be called once during agent startup, typically in the builder.
///
/// # Panics
///
/// Panics if called more than once.
pub fn init_global_identity(identity: Arc<AgentIdentity>) {
    AGENT_IDENTITY
        .set(identity)
        .expect("Global identity already initialized");
}

/// Get a clone of the global agent identity
///
/// Returns an Arc clone for efficient sharing across actors.
///
/// # Panics
///
/// Panics if called before `init_global_identity()`.
pub fn get_global_identity() -> Arc<AgentIdentity> {
    AGENT_IDENTITY
        .get()
        .expect("Global identity not initialized - call init_global_identity() first")
        .clone()
}

#[derive(Serialize, Deserialize)]
struct PersistedKeys {
    ed25519_bytes: Vec<u8>,
    x25519_bytes: Vec<u8>,
}

// ============================================================================
// TOML FORMAT (new consolidated format)
// ============================================================================

/// TOML structure for consolidated agent persistence
#[derive(Serialize, Deserialize)]
struct AgentToml {
    identity: IdentitySection,
    #[serde(default)]
    registration: Option<RegistrationSection>,
}

#[derive(Serialize, Deserialize)]
struct IdentitySection {
    /// Base64-encoded Ed25519 secret key (64 bytes)
    ed25519_secret: String,
    /// Base64-encoded X25519 secret key (32 bytes)
    x25519_secret: String,
}

#[derive(Serialize, Deserialize, Default, Clone)]
struct RegistrationSection {
    /// Human-readable 3-word name (e.g., "conscious-jade-mongoose")
    #[serde(skip_serializing_if = "Option::is_none")]
    name: Option<String>,
    /// Workspace claim token
    #[serde(skip_serializing_if = "Option::is_none")]
    claim_token: Option<String>,
}

pub struct AgentIdentity {
    ed25519_keypair: Keypair,
    x25519_secret: X25519SecretKey,
    /// X25519 public key - reserved for future encrypted communication.
    /// Currently only sent to API during registration; not used locally yet.
    #[allow(dead_code)]
    x25519_public: X25519PublicKey,
    peer_id: PeerId,
    ed25519_public_key_hex: String,
    x25519_public_key_hex: String,
}

// Manual Debug impl that skips secret fields for security
impl std::fmt::Debug for AgentIdentity {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AgentIdentity")
            .field("peer_id", &self.peer_id)
            .field("ed25519_public_key_hex", &self.ed25519_public_key_hex)
            .field("x25519_public_key_hex", &self.x25519_public_key_hex)
            .finish_non_exhaustive()
    }
}

impl AgentIdentity {
    /// Load agent identity from storage, migrating or generating as needed.
    ///
    /// Priority:
    /// 1. Load from `agent.toml` (new consolidated format)
    /// 2. Migrate from `agent_id.key` + separate files (legacy format)
    /// 3. Generate new identity and save as `agent.toml`
    pub fn from_storage_or_generate_new(storage_path: &Path) -> Result<Self> {
        let toml_path = storage_path.join(AGENT_TOML_FILE);
        let legacy_path = storage_path.join(LEGACY_KEY_FILE);

        let (ed25519_keypair, x25519_secret, x25519_public) = if toml_path.exists() {
            Self::load_keypairs_from_toml(&toml_path)?
        } else if legacy_path.exists() {
            Self::migrate_legacy_to_toml(storage_path)?
        } else {
            Self::generate_and_save_new_toml(storage_path)?
        };

        let peer_id = PeerId::from_public_key(&ed25519_keypair.public());
        let ed25519_public_key_hex = hex::encode(ed25519_keypair.public().encode_protobuf());
        let x25519_public_key_hex = hex::encode(x25519_public.as_bytes());

        Ok(Self {
            ed25519_keypair,
            x25519_secret,
            x25519_public,
            peer_id,
            ed25519_public_key_hex,
            x25519_public_key_hex,
        })
    }

    /// Legacy method for backwards compatibility - delegates to from_storage_or_generate_new
    #[deprecated(note = "Use from_storage_or_generate_new instead")]
    pub fn from_file_or_generate_new(path: &Path) -> Result<Self> {
        // Assume path is to agent_id.key, get parent directory
        let storage_path = path.parent().ok_or_else(|| {
            Error::Config("Invalid identity path - no parent directory".to_string())
        })?;
        Self::from_storage_or_generate_new(storage_path)
    }

    pub fn peer_id(&self) -> &PeerId {
        &self.peer_id
    }

    pub fn ed25519_public_key_hex(&self) -> &str {
        &self.ed25519_public_key_hex
    }

    pub fn x25519_public_key_hex(&self) -> &str {
        &self.x25519_public_key_hex
    }

    pub fn sign_message_as_hex(&self, message: &[u8]) -> Result<String> {
        let signature = self
            .ed25519_keypair
            .sign(message)
            .map_err(|e| Self::err("sign", e))?;
        Ok(hex::encode(signature))
    }

    /// Decrypt data using x25519 + XChaCha20-Poly1305
    ///
    /// **Note:** This method is reserved for future encrypted communication features
    /// (e.g., encrypted config updates, encrypted metrics). Currently unused but
    /// maintained for forward compatibility with the encryption package.
    ///
    /// Expected format of encrypted_data:
    /// - First 32 bytes: ephemeral x25519 public key
    /// - Remaining bytes: XChaCha20-Poly1305 ciphertext (includes auth tag)
    #[allow(dead_code)]
    pub fn decrypt_x25519(
        &self,
        encrypted_data: &[u8],
        nonce: &[u8],
    ) -> std::result::Result<Vec<u8>, String> {
        use chacha20poly1305::{
            aead::{Aead, KeyInit},
            XChaCha20Poly1305, XNonce,
        };

        // Extract ephemeral public key (first 32 bytes)
        if encrypted_data.len() < 32 {
            return Err("encrypted data too short".to_string());
        }

        let (ephemeral_public_bytes, ciphertext) = encrypted_data.split_at(32);
        let ephemeral_public_array: [u8; 32] = ephemeral_public_bytes
            .try_into()
            .map_err(|_| "invalid ephemeral public key")?;
        let ephemeral_public = X25519PublicKey::from(ephemeral_public_array);

        // Compute shared secret using Diffie-Hellman
        let shared_secret = self.x25519_secret.diffie_hellman(&ephemeral_public);

        // Use shared secret as XChaCha20-Poly1305 key (24-byte nonce)
        let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into());

        // Decrypt with provided 24-byte nonce
        let nonce_array: XNonce = nonce
            .try_into()
            .map_err(|_| "Invalid nonce length".to_string())?;
        let plaintext = cipher
            .decrypt(&nonce_array, ciphertext)
            .map_err(|e| format!("decryption failed: {}", e))?;

        Ok(plaintext)
    }

    /// Load keypairs from TOML format (new consolidated format)
    fn load_keypairs_from_toml(
        toml_path: &Path,
    ) -> Result<(Keypair, X25519SecretKey, X25519PublicKey)> {
        tracing::info!("Loading agent identity from {}", toml_path.display());

        let toml_content =
            fs::read_to_string(toml_path).map_err(|e| Self::err("read agent.toml", e))?;
        let agent_toml: AgentToml =
            toml::from_str(&toml_content).map_err(|e| Self::err("parse agent.toml", e))?;

        // Decode Ed25519 secret from base64
        let ed25519_bytes = BASE64
            .decode(&agent_toml.identity.ed25519_secret)
            .map_err(|e| Self::err("decode Ed25519 base64", e))?;
        let mut ed25519_bytes_array: [u8; 64] = ed25519_bytes
            .try_into()
            .map_err(|_| Self::err("invalid Ed25519 key length", "expected 64 bytes"))?;
        let ed25519_kp = ed25519::Keypair::try_from_bytes(&mut ed25519_bytes_array)
            .map_err(|e| Self::err("decode Ed25519 keypair", e))?;

        // Decode X25519 secret from base64
        let x25519_bytes = BASE64
            .decode(&agent_toml.identity.x25519_secret)
            .map_err(|e| Self::err("decode X25519 base64", e))?;
        let x25519_bytes_array: [u8; 32] = x25519_bytes
            .try_into()
            .map_err(|_| Self::err("invalid X25519 key length", "expected 32 bytes"))?;
        let x25519_secret = X25519SecretKey::from(x25519_bytes_array);
        let x25519_public = X25519PublicKey::from(&x25519_secret);

        Ok((ed25519_kp.into(), x25519_secret, x25519_public))
    }

    /// Migrate from legacy format (agent_id.key + separate files) to TOML
    fn migrate_legacy_to_toml(
        storage_path: &Path,
    ) -> Result<(Keypair, X25519SecretKey, X25519PublicKey)> {
        let legacy_path = storage_path.join(LEGACY_KEY_FILE);
        tracing::info!(
            "Migrating legacy agent identity from {} to agent.toml",
            legacy_path.display()
        );

        // Load keys from legacy format (JSON or binary)
        let (ed25519_kp, x25519_secret, x25519_public) =
            Self::load_legacy_keypairs(&legacy_path)?;

        // Read optional registration data from separate files
        let claim_token = Self::read_legacy_file(storage_path, LEGACY_CLAIM_TOKEN_FILE);
        let name = Self::read_legacy_file(storage_path, LEGACY_AGENT_NAME_FILE);

        // Build and save TOML
        let agent_toml = AgentToml {
            identity: IdentitySection {
                ed25519_secret: BASE64.encode(ed25519_kp.to_bytes()),
                x25519_secret: BASE64.encode(x25519_secret.to_bytes()),
            },
            registration: if name.is_some() || claim_token.is_some() {
                Some(RegistrationSection { name, claim_token })
            } else {
                None
            },
        };

        let toml_path = storage_path.join(AGENT_TOML_FILE);
        let toml_content =
            toml::to_string_pretty(&agent_toml).map_err(|e| Self::err("serialize agent.toml", e))?;
        fs::write(&toml_path, &toml_content).map_err(|e| Self::err("write agent.toml", e))?;

        tracing::info!(
            "Migration complete - agent data consolidated in {}",
            toml_path.display()
        );

        Ok((ed25519_kp.into(), x25519_secret, x25519_public))
    }

    /// Load keypairs from legacy JSON format (agent_id.key)
    fn load_legacy_keypairs(
        path: &Path,
    ) -> Result<(ed25519::Keypair, X25519SecretKey, X25519PublicKey)> {
        // Try loading as JSON first
        if let Ok(json) = fs::read_to_string(path) {
            if let Ok(persisted) = serde_json::from_str::<PersistedKeys>(&json) {
                let mut ed25519_bytes = persisted.ed25519_bytes;
                let ed25519_kp = ed25519::Keypair::try_from_bytes(&mut ed25519_bytes)
                    .map_err(|e| Self::err("decode Ed25519 key", e))?;

                let x25519_bytes: [u8; 32] = persisted
                    .x25519_bytes
                    .try_into()
                    .map_err(|_| Self::err("invalid X25519 key length", "expected 32 bytes"))?;
                let x25519_secret = X25519SecretKey::from(x25519_bytes);
                let x25519_public = X25519PublicKey::from(&x25519_secret);

                return Ok((ed25519_kp, x25519_secret, x25519_public));
            }
        }

        // Fall back to old binary format (ed25519 only, needs X25519 generation)
        tracing::info!("Loading old binary identity format");
        let mut bytes = fs::read(path).map_err(|e| Self::err("read identity file", e))?;
        let ed25519_kp = ed25519::Keypair::try_from_bytes(&mut bytes)
            .map_err(|e| Self::err("decode Ed25519 key", e))?;

        // Generate new X25519 keypair for migration
        let x25519_secret = X25519SecretKey::random_from_rng(OsRng);
        let x25519_public = X25519PublicKey::from(&x25519_secret);

        Ok((ed25519_kp, x25519_secret, x25519_public))
    }

    /// Read a legacy file (claim_token or agent_name), returning None if missing/empty
    fn read_legacy_file(storage_path: &Path, filename: &str) -> Option<String> {
        let path = storage_path.join(filename);
        if path.exists() {
            fs::read_to_string(&path)
                .ok()
                .map(|s| s.trim().to_string())
                .filter(|s| !s.is_empty())
        } else {
            None
        }
    }

    /// Generate new identity and save as TOML
    fn generate_and_save_new_toml(
        storage_path: &Path,
    ) -> Result<(Keypair, X25519SecretKey, X25519PublicKey)> {
        let toml_path = storage_path.join(AGENT_TOML_FILE);
        tracing::info!(
            "Generating new agent identity (saved to {})",
            toml_path.display()
        );

        // Generate Ed25519 keypair
        let mut bytes = [0u8; 32];
        use rand::RngCore;
        OsRng.fill_bytes(&mut bytes);
        let ed25519_kp = ed25519::Keypair::from(
            ed25519::SecretKey::try_from_bytes(&mut bytes)
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))?,
        );

        // Generate X25519 keypair
        let x25519_secret = X25519SecretKey::random_from_rng(OsRng);
        let x25519_public = X25519PublicKey::from(&x25519_secret);

        // Build TOML (no registration data for new identity)
        let agent_toml = AgentToml {
            identity: IdentitySection {
                ed25519_secret: BASE64.encode(ed25519_kp.to_bytes()),
                x25519_secret: BASE64.encode(x25519_secret.to_bytes()),
            },
            registration: None,
        };

        // Ensure storage directory exists
        fs::create_dir_all(storage_path).map_err(|e| Self::err("create storage directory", e))?;

        let toml_content =
            toml::to_string_pretty(&agent_toml).map_err(|e| Self::err("serialize agent.toml", e))?;
        fs::write(&toml_path, &toml_content).map_err(|e| Self::err("write agent.toml", e))?;

        Ok((ed25519_kp.into(), x25519_secret, x25519_public))
    }

    fn err(context: &str, error: impl std::fmt::Display) -> Error {
        Error::Io(std::io::Error::new(
            std::io::ErrorKind::Other,
            format!("{}: {}", context, error),
        ))
    }
}

// ============================================================================
// PUBLIC HELPERS
// ============================================================================

/// Update the registration section of agent.toml
///
/// This is used by the AlertsWorker to persist name and claim_token after
/// successful registration with the server.
///
/// If agent.toml doesn't exist, this will return an error - the agent must
/// be initialized first via AgentIdentity::from_storage_or_generate_new().
pub fn update_agent_toml_registration(
    storage_path: &Path,
    name: Option<&str>,
    claim_token: Option<&str>,
) -> Result<()> {
    let toml_path = storage_path.join(AGENT_TOML_FILE);

    if !toml_path.exists() {
        return Err(Error::Config(format!(
            "Cannot update registration: {} does not exist",
            toml_path.display()
        )));
    }

    // Read existing TOML
    let toml_content =
        fs::read_to_string(&toml_path).map_err(|e| Error::Io(std::io::Error::new(
            e.kind(),
            format!("read agent.toml: {}", e),
        )))?;
    let mut agent_toml: AgentToml =
        toml::from_str(&toml_content).map_err(|e| Error::Config(format!("parse agent.toml: {}", e)))?;

    // Update registration section
    let mut registration = agent_toml.registration.unwrap_or_default();

    if let Some(n) = name {
        registration.name = Some(n.to_string());
    }
    if let Some(ct) = claim_token {
        registration.claim_token = Some(ct.to_string());
    }

    // Only set registration if there's something to store
    agent_toml.registration = if registration.name.is_some() || registration.claim_token.is_some() {
        Some(registration)
    } else {
        None
    };

    // Write back
    let toml_content = toml::to_string_pretty(&agent_toml)
        .map_err(|e| Error::Config(format!("serialize agent.toml: {}", e)))?;
    fs::write(&toml_path, &toml_content).map_err(|e| Error::Io(std::io::Error::new(
        e.kind(),
        format!("write agent.toml: {}", e),
    )))?;

    tracing::debug!("Updated registration in {}", toml_path.display());
    Ok(())
}

/// Read claim token from storage (TOML or legacy file)
///
/// Used by AlertsWorker during startup to load existing claim token.
pub fn read_claim_token(storage_path: &Path) -> Option<String> {
    let toml_path = storage_path.join(AGENT_TOML_FILE);

    // Try TOML first
    if toml_path.exists() {
        if let Ok(toml_content) = fs::read_to_string(&toml_path) {
            if let Ok(agent_toml) = toml::from_str::<AgentToml>(&toml_content) {
                if let Some(reg) = agent_toml.registration {
                    if reg.claim_token.is_some() {
                        return reg.claim_token;
                    }
                }
            }
        }
    }

    // Fall back to legacy file
    AgentIdentity::read_legacy_file(storage_path, LEGACY_CLAIM_TOKEN_FILE)
}