use crate::identity::{AgentId, MachineId};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum TrustLevel {
Blocked,
#[default]
Unknown,
Known,
Trusted,
}
impl std::fmt::Display for TrustLevel {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Blocked => write!(f, "blocked"),
Self::Unknown => write!(f, "unknown"),
Self::Known => write!(f, "known"),
Self::Trusted => write!(f, "trusted"),
}
}
}
impl std::str::FromStr for TrustLevel {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"blocked" => Ok(Self::Blocked),
"unknown" => Ok(Self::Unknown),
"known" => Ok(Self::Known),
"trusted" => Ok(Self::Trusted),
_ => Err(format!(
"invalid trust level: {s} (valid values: blocked, unknown, known, trusted)"
)),
}
}
}
impl TrustLevel {
#[must_use]
pub fn rank(self) -> u8 {
match self {
Self::Blocked => 0,
Self::Unknown => 1,
Self::Known => 2,
Self::Trusted => 3,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum IdentityType {
#[default]
Anonymous,
Known,
Trusted,
Pinned,
}
impl std::fmt::Display for IdentityType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Anonymous => write!(f, "anonymous"),
Self::Known => write!(f, "known"),
Self::Trusted => write!(f, "trusted"),
Self::Pinned => write!(f, "pinned"),
}
}
}
impl std::str::FromStr for IdentityType {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"anonymous" => Ok(Self::Anonymous),
"known" => Ok(Self::Known),
"trusted" => Ok(Self::Trusted),
"pinned" => Ok(Self::Pinned),
_ => Err(format!("invalid identity type: {s}")),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MachineRecord {
pub machine_id: MachineId,
pub label: Option<String>,
pub first_seen: u64,
pub last_seen: u64,
pub pinned: bool,
}
impl MachineRecord {
#[must_use]
pub fn new(machine_id: MachineId, label: Option<String>) -> Self {
let now = now_secs();
Self {
machine_id,
label,
first_seen: now,
last_seen: now,
pinned: false,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Contact {
pub agent_id: AgentId,
pub trust_level: TrustLevel,
pub label: Option<String>,
pub added_at: u64,
pub last_seen: Option<u64>,
#[serde(default)]
pub identity_type: IdentityType,
#[serde(default)]
pub machines: Vec<MachineRecord>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RevocationRecord {
pub agent_id: AgentId,
pub reason: String,
pub timestamp: u64,
pub revoker_id: Option<AgentId>,
}
#[derive(Debug)]
pub struct ContactStore {
contacts: HashMap<[u8; 32], Contact>,
revoked_keys: HashSet<[u8; 32]>,
revocations: Vec<RevocationRecord>,
storage_path: PathBuf,
}
#[derive(Serialize, Deserialize)]
struct ContactsFile {
contacts: Vec<Contact>,
#[serde(default)]
revocations: Vec<RevocationRecord>,
}
fn now_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
impl ContactStore {
pub fn new(storage_path: PathBuf) -> Self {
let mut store = Self {
contacts: HashMap::new(),
revoked_keys: HashSet::new(),
revocations: Vec::new(),
storage_path,
};
let _ = store.load();
store
}
pub fn add(&mut self, mut contact: Contact) {
if self.revoked_keys.contains(&contact.agent_id.0) {
contact.trust_level = TrustLevel::Blocked;
}
if matches!(contact.trust_level, TrustLevel::Known | TrustLevel::Trusted)
&& contact.identity_type == IdentityType::Anonymous
{
contact.identity_type = IdentityType::Known;
}
self.contacts.insert(contact.agent_id.0, contact);
let _ = self.save();
}
pub fn remove(&mut self, agent_id: &AgentId) -> Option<Contact> {
let result = self.contacts.remove(&agent_id.0);
if result.is_some() {
let _ = self.save();
}
result
}
pub fn set_trust(&mut self, agent_id: &AgentId, trust_level: TrustLevel) {
let effective_trust = if self.revoked_keys.contains(&agent_id.0) {
TrustLevel::Blocked
} else {
trust_level
};
let entry = self.contacts.entry(agent_id.0).or_insert_with(|| Contact {
agent_id: *agent_id,
trust_level: effective_trust,
label: None,
added_at: now_secs(),
last_seen: None,
identity_type: IdentityType::default(),
machines: Vec::new(),
});
entry.trust_level = effective_trust;
if matches!(effective_trust, TrustLevel::Known | TrustLevel::Trusted)
&& entry.identity_type == IdentityType::Anonymous
{
entry.identity_type = IdentityType::Known;
}
let _ = self.save();
}
pub fn revoke(&mut self, agent_id: &AgentId, reason: &str) {
if self.revoked_keys.contains(&agent_id.0) {
return;
}
self.revoked_keys.insert(agent_id.0);
self.revocations.push(RevocationRecord {
agent_id: *agent_id,
reason: reason.to_string(),
timestamp: now_secs(),
revoker_id: None,
});
self.set_trust(agent_id, TrustLevel::Blocked);
}
pub fn revoke_with_revoker(&mut self, agent_id: &AgentId, reason: &str, revoker_id: &AgentId) {
if self.revoked_keys.contains(&agent_id.0) {
return;
}
self.revoked_keys.insert(agent_id.0);
self.revocations.push(RevocationRecord {
agent_id: *agent_id,
reason: reason.to_string(),
timestamp: now_secs(),
revoker_id: Some(*revoker_id),
});
self.set_trust(agent_id, TrustLevel::Blocked);
}
pub fn is_revoked(&self, agent_id: &AgentId) -> bool {
self.revoked_keys.contains(&agent_id.0)
}
pub fn revocations(&self) -> &[RevocationRecord] {
&self.revocations
}
pub fn get(&self, agent_id: &AgentId) -> Option<&Contact> {
self.contacts.get(&agent_id.0)
}
pub fn get_mut(&mut self, agent_id: &AgentId) -> Option<&mut Contact> {
self.contacts.get_mut(&agent_id.0)
}
pub fn list(&self) -> Vec<&Contact> {
self.contacts.values().collect()
}
pub fn is_trusted(&self, agent_id: &AgentId) -> bool {
self.contacts
.get(&agent_id.0)
.map(|c| c.trust_level == TrustLevel::Trusted)
.unwrap_or(false)
}
pub fn is_blocked(&self, agent_id: &AgentId) -> bool {
self.contacts
.get(&agent_id.0)
.map(|c| c.trust_level == TrustLevel::Blocked)
.unwrap_or(false)
}
pub fn trust_level(&self, agent_id: &AgentId) -> TrustLevel {
self.contacts
.get(&agent_id.0)
.map(|c| c.trust_level)
.unwrap_or(TrustLevel::Unknown)
}
pub fn touch(&mut self, agent_id: &AgentId) {
if let Some(contact) = self.contacts.get_mut(&agent_id.0) {
contact.last_seen = Some(now_secs());
let _ = self.save();
}
}
pub fn add_machine(&mut self, agent_id: &AgentId, record: MachineRecord) -> bool {
let contact = self.contacts.entry(agent_id.0).or_insert_with(|| Contact {
agent_id: *agent_id,
trust_level: TrustLevel::Unknown,
label: None,
added_at: now_secs(),
last_seen: None,
identity_type: IdentityType::default(),
machines: Vec::new(),
});
if let Some(existing) = contact
.machines
.iter_mut()
.find(|m| m.machine_id == record.machine_id)
{
existing.last_seen = now_secs();
if record.label.is_some() {
existing.label = record.label;
}
let _ = self.save();
false
} else {
contact.machines.push(record);
let _ = self.save();
true
}
}
pub fn remove_machine(&mut self, agent_id: &AgentId, machine_id: &MachineId) -> bool {
if let Some(contact) = self.contacts.get_mut(&agent_id.0) {
let before = contact.machines.len();
contact.machines.retain(|m| m.machine_id != *machine_id);
let removed = contact.machines.len() < before;
if removed {
let _ = self.save();
}
removed
} else {
false
}
}
pub fn machines(&self, agent_id: &AgentId) -> &[MachineRecord] {
self.contacts
.get(&agent_id.0)
.map(|c| c.machines.as_slice())
.unwrap_or(&[])
}
pub fn pin_machine(&mut self, agent_id: &AgentId, machine_id: &MachineId) -> bool {
if let Some(contact) = self.contacts.get_mut(&agent_id.0) {
if let Some(record) = contact
.machines
.iter_mut()
.find(|m| m.machine_id == *machine_id)
{
record.pinned = true;
contact.identity_type = IdentityType::Pinned;
let _ = self.save();
return true;
}
}
false
}
pub fn unpin_machine(&mut self, agent_id: &AgentId, machine_id: &MachineId) -> bool {
if let Some(contact) = self.contacts.get_mut(&agent_id.0) {
if let Some(record) = contact
.machines
.iter_mut()
.find(|m| m.machine_id == *machine_id)
{
record.pinned = false;
if !contact.machines.iter().any(|m| m.pinned) {
contact.identity_type = IdentityType::Known;
}
let _ = self.save();
return true;
}
}
false
}
pub fn set_identity_type(&mut self, agent_id: &AgentId, identity_type: IdentityType) {
let contact = self.contacts.entry(agent_id.0).or_insert_with(|| Contact {
agent_id: *agent_id,
trust_level: TrustLevel::Unknown,
label: None,
added_at: now_secs(),
last_seen: None,
identity_type: IdentityType::default(),
machines: Vec::new(),
});
contact.identity_type = identity_type;
let _ = self.save();
}
fn save(&self) -> std::io::Result<()> {
let file = ContactsFile {
contacts: self.contacts.values().cloned().collect(),
revocations: self.revocations.clone(),
};
let json = serde_json::to_string_pretty(&file)
.map_err(|e| std::io::Error::other(format!("serialize: {e}")))?;
if let Some(parent) = self.storage_path.parent() {
std::fs::create_dir_all(parent)?;
}
let tmp = self.storage_path.with_extension("tmp");
std::fs::write(&tmp, &json)?;
std::fs::rename(&tmp, &self.storage_path)?;
Ok(())
}
fn load(&mut self) -> std::io::Result<()> {
if !self.storage_path.exists() {
return Ok(());
}
let json = std::fs::read_to_string(&self.storage_path)?;
let file: ContactsFile = serde_json::from_str(&json)
.map_err(|e| std::io::Error::other(format!("deserialize: {e}")))?;
for contact in file.contacts {
self.contacts.insert(contact.agent_id.0, contact);
}
for record in &file.revocations {
self.revoked_keys.insert(record.agent_id.0);
}
self.revocations = file.revocations;
Ok(())
}
pub fn storage_path(&self) -> &Path {
&self.storage_path
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::identity::AgentKeypair;
fn test_agent_id() -> AgentId {
AgentKeypair::generate().expect("keygen").agent_id()
}
fn test_machine_id() -> MachineId {
let kp = crate::identity::MachineKeypair::generate().expect("keygen");
kp.machine_id()
}
#[test]
fn test_trust_level_display_and_parse() {
for level in [
TrustLevel::Blocked,
TrustLevel::Unknown,
TrustLevel::Known,
TrustLevel::Trusted,
] {
let s = level.to_string();
let parsed: TrustLevel = s.parse().expect("parse");
assert_eq!(parsed, level);
}
}
#[test]
fn test_trust_level_parse_invalid() {
assert!("invalid".parse::<TrustLevel>().is_err());
}
#[test]
fn test_identity_type_display_and_parse() {
for ty in [
IdentityType::Anonymous,
IdentityType::Known,
IdentityType::Trusted,
IdentityType::Pinned,
] {
let s = ty.to_string();
let parsed: IdentityType = s.parse().expect("parse");
assert_eq!(parsed, ty);
}
}
#[test]
fn test_identity_type_parse_invalid() {
assert!("invalid".parse::<IdentityType>().is_err());
}
#[test]
fn test_identity_type_default() {
assert_eq!(IdentityType::default(), IdentityType::Anonymous);
}
#[test]
fn test_contact_store_add_get_remove() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let id = test_agent_id();
store.add(Contact {
agent_id: id,
trust_level: TrustLevel::Trusted,
label: Some("Test".to_string()),
added_at: 1000,
last_seen: None,
identity_type: IdentityType::default(),
machines: Vec::new(),
});
assert!(store.get(&id).is_some());
assert!(store.is_trusted(&id));
assert!(!store.is_blocked(&id));
let removed = store.remove(&id);
assert!(removed.is_some());
assert!(store.get(&id).is_none());
}
#[test]
fn test_contact_store_set_trust() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let id = test_agent_id();
store.set_trust(&id, TrustLevel::Known);
assert_eq!(store.trust_level(&id), TrustLevel::Known);
store.set_trust(&id, TrustLevel::Blocked);
assert!(store.is_blocked(&id));
}
#[test]
fn test_contact_store_default_unknown() {
let dir = tempfile::tempdir().expect("tmpdir");
let store = ContactStore::new(dir.path().join("contacts.json"));
let id = test_agent_id();
assert_eq!(store.trust_level(&id), TrustLevel::Unknown);
}
#[test]
fn test_contact_store_list() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let id1 = test_agent_id();
let id2 = test_agent_id();
store.set_trust(&id1, TrustLevel::Trusted);
store.set_trust(&id2, TrustLevel::Known);
assert_eq!(store.list().len(), 2);
}
#[test]
fn test_contact_store_persistence() {
let dir = tempfile::tempdir().expect("tmpdir");
let path = dir.path().join("contacts.json");
let id = test_agent_id();
{
let mut store = ContactStore::new(path.clone());
store.add(Contact {
agent_id: id,
trust_level: TrustLevel::Trusted,
label: Some("Persistent".to_string()),
added_at: 2000,
last_seen: None,
identity_type: IdentityType::default(),
machines: Vec::new(),
});
}
let store = ContactStore::new(path);
let contact = store.get(&id).expect("should exist after reload");
assert_eq!(contact.trust_level, TrustLevel::Trusted);
assert_eq!(contact.label.as_deref(), Some("Persistent"));
}
#[test]
fn test_contact_store_touch() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let id = test_agent_id();
store.set_trust(&id, TrustLevel::Known);
assert!(store.get(&id).expect("exists").last_seen.is_none());
store.touch(&id);
assert!(store.get(&id).expect("exists").last_seen.is_some());
}
#[test]
fn test_trust_level_serde() {
let json = serde_json::to_string(&TrustLevel::Trusted).expect("ser");
assert_eq!(json, "\"trusted\"");
let parsed: TrustLevel = serde_json::from_str(&json).expect("de");
assert_eq!(parsed, TrustLevel::Trusted);
}
#[test]
fn test_machine_record_new() {
let mid = test_machine_id();
let rec = MachineRecord::new(mid, Some("laptop".to_string()));
assert_eq!(rec.machine_id, mid);
assert_eq!(rec.label.as_deref(), Some("laptop"));
assert!(!rec.pinned);
assert!(rec.first_seen > 0);
assert_eq!(rec.first_seen, rec.last_seen);
}
#[test]
fn test_add_machine_new() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let agent = test_agent_id();
let machine = test_machine_id();
let rec = MachineRecord::new(machine, None);
let is_new = store.add_machine(&agent, rec);
assert!(is_new);
assert_eq!(store.machines(&agent).len(), 1);
}
#[test]
fn test_add_machine_duplicate_updates_last_seen() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let agent = test_agent_id();
let machine = test_machine_id();
store.add_machine(&agent, MachineRecord::new(machine, None));
let is_new = store.add_machine(&agent, MachineRecord::new(machine, Some("new".into())));
assert!(!is_new);
assert_eq!(store.machines(&agent).len(), 1);
assert_eq!(store.machines(&agent)[0].label.as_deref(), Some("new"));
}
#[test]
fn test_remove_machine() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let agent = test_agent_id();
let machine = test_machine_id();
store.add_machine(&agent, MachineRecord::new(machine, None));
assert!(store.remove_machine(&agent, &machine));
assert_eq!(store.machines(&agent).len(), 0);
assert!(!store.remove_machine(&agent, &machine));
}
#[test]
fn test_pin_and_unpin_machine() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let agent = test_agent_id();
let machine = test_machine_id();
store.add_machine(&agent, MachineRecord::new(machine, None));
assert!(store.pin_machine(&agent, &machine));
assert_eq!(
store.get(&agent).expect("exists").identity_type,
IdentityType::Pinned
);
assert!(store.machines(&agent)[0].pinned);
let other = test_machine_id();
assert!(!store.pin_machine(&agent, &other));
assert!(store.unpin_machine(&agent, &machine));
assert!(!store.machines(&agent)[0].pinned);
assert_eq!(
store.get(&agent).expect("exists").identity_type,
IdentityType::Known
);
}
#[test]
fn test_set_identity_type() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let agent = test_agent_id();
store.set_identity_type(&agent, IdentityType::Trusted);
assert_eq!(
store.get(&agent).expect("exists").identity_type,
IdentityType::Trusted
);
}
#[test]
fn test_machine_record_persistence() {
let dir = tempfile::tempdir().expect("tmpdir");
let path = dir.path().join("contacts.json");
let agent = test_agent_id();
let machine = test_machine_id();
{
let mut store = ContactStore::new(path.clone());
store.add_machine(&agent, MachineRecord::new(machine, Some("desktop".into())));
store.pin_machine(&agent, &machine);
}
let store = ContactStore::new(path);
let machines = store.machines(&agent);
assert_eq!(machines.len(), 1);
assert_eq!(machines[0].machine_id, machine);
assert_eq!(machines[0].label.as_deref(), Some("desktop"));
assert!(machines[0].pinned);
assert_eq!(
store.get(&agent).expect("exists").identity_type,
IdentityType::Pinned
);
}
#[test]
fn test_backward_compat_no_machines_field() {
let dir = tempfile::tempdir().expect("tmpdir");
let path = dir.path().join("contacts.json");
let agent_id = test_agent_id();
{
let mut store = ContactStore::new(path.clone());
store.add(Contact {
agent_id,
trust_level: TrustLevel::Trusted,
label: None,
added_at: 1000,
last_seen: None,
identity_type: IdentityType::Anonymous,
machines: Vec::new(),
});
}
let json = std::fs::read_to_string(&path).expect("read");
let mut root: serde_json::Value =
serde_json::from_str(&json).expect("parse saved contacts");
if let Some(contacts) = root.get_mut("contacts").and_then(|v| v.as_array_mut()) {
for c in contacts.iter_mut() {
if let Some(obj) = c.as_object_mut() {
obj.remove("identity_type");
obj.remove("machines");
}
}
}
let stripped = serde_json::to_string_pretty(&root).expect("serialize");
std::fs::write(&path, &stripped).expect("write");
let store = ContactStore::new(path);
let contact = store.get(&agent_id).expect("should load");
assert_eq!(contact.trust_level, TrustLevel::Trusted);
assert_eq!(contact.identity_type, IdentityType::Anonymous);
assert!(contact.machines.is_empty());
}
#[test]
fn test_revoke_blocks_future_messages() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let id = test_agent_id();
store.set_trust(&id, TrustLevel::Trusted);
assert!(store.is_trusted(&id));
assert!(!store.is_revoked(&id));
store.revoke(&id, "key compromised");
assert!(store.is_revoked(&id));
assert!(store.is_blocked(&id));
assert!(!store.is_trusted(&id));
assert_eq!(store.trust_level(&id), TrustLevel::Blocked);
}
#[test]
fn test_revocations_persist_across_reload() {
let dir = tempfile::tempdir().expect("tmpdir");
let path = dir.path().join("contacts.json");
let id = test_agent_id();
{
let mut store = ContactStore::new(path.clone());
store.set_trust(&id, TrustLevel::Trusted);
store.revoke(&id, "stolen key");
}
let store = ContactStore::new(path);
assert!(store.is_revoked(&id));
assert!(store.is_blocked(&id));
assert_eq!(store.revocations().len(), 1);
assert_eq!(store.revocations()[0].reason, "stolen key");
}
#[test]
fn test_revoked_key_cannot_be_unrevoked_by_set_trust() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let id = test_agent_id();
store.set_trust(&id, TrustLevel::Trusted);
store.revoke(&id, "compromised");
store.set_trust(&id, TrustLevel::Trusted);
assert!(store.is_revoked(&id));
assert!(store.is_blocked(&id));
assert_eq!(store.trust_level(&id), TrustLevel::Blocked);
}
#[test]
fn test_revoked_key_stays_blocked_after_add() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let id = test_agent_id();
store.revoke(&id, "bad actor");
store.add(Contact {
agent_id: id,
trust_level: TrustLevel::Trusted,
label: Some("Sneaky".to_string()),
added_at: 3000,
last_seen: None,
identity_type: IdentityType::default(),
machines: Vec::new(),
});
assert!(store.is_revoked(&id));
assert!(store.is_blocked(&id));
}
#[test]
fn test_revoke_with_revoker() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let target = test_agent_id();
let revoker = test_agent_id();
store.revoke_with_revoker(&target, "audit finding", &revoker);
assert!(store.is_revoked(&target));
let record = &store.revocations()[0];
assert_eq!(record.revoker_id, Some(revoker));
assert_eq!(record.reason, "audit finding");
}
#[test]
fn test_duplicate_revocation_ignored() {
let dir = tempfile::tempdir().expect("tmpdir");
let mut store = ContactStore::new(dir.path().join("contacts.json"));
let id = test_agent_id();
store.revoke(&id, "first revocation");
store.revoke(&id, "second revocation");
assert_eq!(store.revocations().len(), 1);
assert_eq!(store.revocations()[0].reason, "first revocation");
}
}