use anyhow::{Result, anyhow};
use ark_bn254::Fr;
use ark_ff::PrimeField;
use nostr_sdk::prelude::*;
use serde::{Serialize, Deserialize};
use crate::key_manager::ManagedKey;
use crate::membership::bn254_leaf_commitment;
use crate::nostr::derive_nsec_from_leaf_secret;
use super::storage::{
SecureStorage,
KEY_DID_PRIVATE, KEY_LEAF_SECRET,
};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentIdentityState {
pub did: String,
pub pubkey_hex: String,
pub agent_npub: String,
pub agent_npub_hex: String,
pub leaf_commitment: String,
}
pub struct AgentIdentity<S: SecureStorage> {
storage: S,
}
impl<S: SecureStorage> AgentIdentity<S> {
pub fn new(storage: S) -> Self {
Self { storage }
}
pub fn is_initialized(&self) -> bool {
self.storage.exists(KEY_DID_PRIVATE) && self.storage.exists(KEY_LEAF_SECRET)
}
pub fn initialize(&self) -> Result<AgentIdentityState> {
if self.is_initialized() {
return Err(anyhow!("Agent identity already initialized"));
}
let did_key = ManagedKey::generate()?;
self.storage.store(KEY_DID_PRIVATE, &did_key.secret_key.secret_bytes())
.map_err(|e| anyhow!("Failed to store DID private key: {}", e))?;
let leaf_secret = derive_leaf_secret_from_did(&did_key)?;
let leaf_secret_bytes = serialize_leaf_secret(&leaf_secret);
self.storage.store(KEY_LEAF_SECRET, &leaf_secret_bytes)
.map_err(|e| anyhow!("Failed to store leaf_secret: {}", e))?;
let leaf_commitment = bn254_leaf_commitment(&leaf_secret);
let leaf_commitment_hex = fr_to_hex(&leaf_commitment);
let agent_keys = derive_nsec_from_leaf_secret(&leaf_secret)?;
let agent_npub = Keys::new(agent_keys.clone()).public_key();
Ok(AgentIdentityState {
did: did_key.to_did(),
pubkey_hex: did_key.pubkey_hex(),
agent_npub: agent_npub.to_bech32()?,
agent_npub_hex: agent_npub.to_hex(),
leaf_commitment: leaf_commitment_hex,
})
}
pub fn load(&self) -> Result<AgentIdentityState> {
if !self.is_initialized() {
return Err(anyhow!("Agent identity not initialized"));
}
let did_bytes = self.storage.retrieve(KEY_DID_PRIVATE)
.map_err(|e| anyhow!("Failed to load DID: {}", e))?;
let did_key = ManagedKey::from_bytes(&did_bytes)?;
let leaf_secret_bytes = self.storage.retrieve(KEY_LEAF_SECRET)
.map_err(|e| anyhow!("Failed to load leaf_secret: {}", e))?;
let leaf_secret = deserialize_leaf_secret(&leaf_secret_bytes)?;
let leaf_commitment = bn254_leaf_commitment(&leaf_secret);
let leaf_commitment_hex = fr_to_hex(&leaf_commitment);
let agent_keys = derive_nsec_from_leaf_secret(&leaf_secret)?;
let agent_npub = Keys::new(agent_keys).public_key();
Ok(AgentIdentityState {
did: did_key.to_did(),
pubkey_hex: did_key.pubkey_hex(),
agent_npub: agent_npub.to_bech32()?,
agent_npub_hex: agent_npub.to_hex(),
leaf_commitment: leaf_commitment_hex,
})
}
pub fn get_agent_keys(&self) -> Result<Keys> {
let leaf_secret_bytes = self.storage.retrieve(KEY_LEAF_SECRET)
.map_err(|e| anyhow!("Failed to load leaf_secret: {}", e))?;
let leaf_secret = deserialize_leaf_secret(&leaf_secret_bytes)?;
let nsec = derive_nsec_from_leaf_secret(&leaf_secret)?;
Ok(Keys::new(nsec))
}
pub fn get_leaf_secret(&self) -> Result<[Fr; 5]> {
let leaf_secret_bytes = self.storage.retrieve(KEY_LEAF_SECRET)
.map_err(|e| anyhow!("Failed to load leaf_secret: {}", e))?;
deserialize_leaf_secret(&leaf_secret_bytes)
}
}
fn derive_leaf_secret_from_did(did_key: &ManagedKey) -> Result<[Fr; 5]> {
use sha2::{Sha256, Digest};
let mut leaf_secret = [Fr::from(0u64); 5];
for i in 0..5 {
let mut hasher = Sha256::new();
hasher.update(b"signedby_leaf_secret_v1:");
hasher.update(&did_key.secret_key.secret_bytes());
hasher.update(&[i as u8]);
let hash = hasher.finalize();
leaf_secret[i] = Fr::from_be_bytes_mod_order(&hash);
}
Ok(leaf_secret)
}
fn serialize_leaf_secret(leaf_secret: &[Fr; 5]) -> Vec<u8> {
let mut bytes = Vec::with_capacity(5 * 32);
for elem in leaf_secret {
let bigint = elem.into_bigint();
for limb in bigint.0.iter() {
bytes.extend_from_slice(&limb.to_le_bytes());
}
}
bytes
}
fn deserialize_leaf_secret(bytes: &[u8]) -> Result<[Fr; 5]> {
if bytes.len() != 5 * 32 {
return Err(anyhow!("Invalid leaf_secret length: expected {}, got {}", 5 * 32, bytes.len()));
}
let mut leaf_secret = [Fr::from(0u64); 5];
for i in 0..5 {
let offset = i * 32;
let elem_bytes: [u8; 32] = bytes[offset..offset + 32].try_into()?;
leaf_secret[i] = Fr::from_le_bytes_mod_order(&elem_bytes);
}
Ok(leaf_secret)
}
fn fr_to_hex(fr: &Fr) -> String {
let bigint = fr.into_bigint();
let mut bytes = Vec::new();
for limb in bigint.0.iter().rev() {
bytes.extend_from_slice(&limb.to_be_bytes());
}
while bytes.len() > 32 && bytes[0] == 0 {
bytes.remove(0);
}
while bytes.len() < 32 {
bytes.insert(0, 0);
}
hex::encode(bytes)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sdk::storage::EncryptedFileStorage;
use tempfile::tempdir;
#[test]
fn test_initialize_identity() {
let dir = tempdir().unwrap();
let storage = EncryptedFileStorage::new(dir.path().to_path_buf()).unwrap();
let identity = AgentIdentity::new(storage);
assert!(!identity.is_initialized());
let state = identity.initialize().unwrap();
assert!(identity.is_initialized());
assert!(state.did.starts_with("did:btcr:"));
}
#[test]
fn test_load_identity() {
let dir = tempdir().unwrap();
let storage = EncryptedFileStorage::new(dir.path().to_path_buf()).unwrap();
let identity = AgentIdentity::new(storage);
let state1 = identity.initialize().unwrap();
let state2 = identity.load().unwrap();
assert_eq!(state1.did, state2.did);
assert_eq!(state1.agent_npub, state2.agent_npub);
assert_eq!(state1.leaf_commitment, state2.leaf_commitment);
}
}