use crate::errors::LitError;
use serde::{Deserialize, Serialize};
use sha3::{Digest, Sha3_256};
use std::fs;
use std::path::Path;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct Capability {
pub resource: String,
pub action: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub caveats: Option<serde_json::Value>,
}
impl std::fmt::Display for Capability {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}:{}", self.resource, self.action)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UcanToken {
pub version: String,
pub issuer: String,
pub audience: String,
pub capabilities: Vec<Capability>,
pub expiration: i64,
#[serde(skip_serializing_if = "Option::is_none")]
pub not_before: Option<i64>,
pub nonce: String,
#[serde(default)]
pub proof: Vec<String>,
pub signature: String,
}
impl UcanToken {
pub fn new(
issuer: String,
audience: String,
capabilities: Vec<Capability>,
duration_secs: i64,
) -> Self {
let now = chrono::Utc::now().timestamp();
let nonce = format!("{:016x}", now ^ std::process::id() as i64);
UcanToken {
version: "0.10.0".to_string(),
issuer,
audience,
capabilities,
expiration: now + duration_secs,
not_before: Some(now),
nonce,
proof: Vec::new(),
signature: String::new(),
}
}
pub fn sign(&mut self, private_key_hex: &str) -> Result<(), LitError> {
let payload = self.payload_bytes()?;
let private_bytes = hex::decode(private_key_hex)
.map_err(|e| LitError::general(format!("Invalid key: {}", e)))?;
let mut hasher = Sha3_256::new();
hasher.update(&private_bytes);
hasher.update(&payload);
let sig = hasher.finalize();
self.signature = hex::encode(sig);
Ok(())
}
pub fn verify(&self, issuer_public_key_hex: &str) -> Result<bool, LitError> {
let payload = self.payload_bytes()?;
let public_bytes = hex::decode(issuer_public_key_hex)
.map_err(|e| LitError::general(format!("Invalid key: {}", e)))?;
let mut hasher = Sha3_256::new();
hasher.update(&public_bytes);
hasher.update(&payload);
let expected = hasher.finalize();
let sig_bytes = hex::decode(&self.signature)
.map_err(|e| LitError::general(format!("Invalid signature: {}", e)))?;
Ok(subtle::ConstantTimeEq::ct_eq(sig_bytes.as_slice(), expected.as_slice()).into())
}
pub fn is_valid(&self) -> bool {
let now = chrono::Utc::now().timestamp();
if now > self.expiration {
return false;
}
if let Some(nb) = self.not_before {
if now < nb {
return false;
}
}
true
}
pub fn has_capability(&self, resource: &str, action: &str) -> bool {
self.capabilities.iter().any(|cap| {
let resource_match = cap.resource == "*"
|| cap.resource == resource
|| cap
.resource
.strip_suffix('*')
.map(|prefix| resource.starts_with(prefix))
.unwrap_or(false)
|| resource.starts_with(&cap.resource);
resource_match && (cap.action == "*" || cap.action == action)
})
}
pub fn cid(&self) -> Result<String, LitError> {
let bytes = self.payload_bytes()?;
let hash = Sha3_256::digest(&bytes);
Ok(hex::encode(hash))
}
pub fn delegate(
&self,
new_audience: String,
capabilities: Vec<Capability>,
duration_secs: i64,
) -> Result<UcanToken, LitError> {
for cap in &capabilities {
if !self.has_capability(&cap.resource, &cap.action) {
return Err(LitError::general(format!(
"Cannot delegate capability '{}:{}' — not held by parent token",
cap.resource, cap.action
)));
}
}
let parent_cid = self.cid()?;
let mut child = UcanToken::new(
self.audience.clone(), new_audience,
capabilities,
duration_secs,
);
child.proof.push(parent_cid);
for p in &self.proof {
child.proof.push(p.clone());
}
Ok(child)
}
fn payload_bytes(&self) -> Result<Vec<u8>, LitError> {
let payload = serde_json::json!({
"version": self.version,
"issuer": self.issuer,
"audience": self.audience,
"capabilities": self.capabilities,
"expiration": self.expiration,
"not_before": self.not_before,
"nonce": self.nonce,
"proof": self.proof,
});
serde_json::to_vec(&payload)
.map_err(|e| LitError::general(format!("Failed to serialize UCAN: {}", e)))
}
}
pub fn ucan_dir(repo_root: &Path) -> std::path::PathBuf {
repo_root.join(".lit").join("ucan")
}
pub fn save_token(repo_root: &Path, token: &UcanToken) -> Result<String, LitError> {
let dir = ucan_dir(repo_root);
fs::create_dir_all(&dir)
.map_err(|e| LitError::io(format!("Failed to create UCAN dir: {}", e)))?;
let cid = token.cid()?;
let path = dir.join(format!("{}.json", &cid[..16]));
let json = serde_json::to_string_pretty(token)
.map_err(|e| LitError::general(format!("Failed to serialize token: {}", e)))?;
fs::write(&path, json).map_err(|e| LitError::io(format!("Failed to write token: {}", e)))?;
Ok(cid)
}
pub fn load_tokens_for(repo_root: &Path, audience_did: &str) -> Result<Vec<UcanToken>, LitError> {
let dir = ucan_dir(repo_root);
if !dir.exists() {
return Ok(Vec::new());
}
let mut tokens = Vec::new();
for entry in fs::read_dir(&dir).map_err(|e| LitError::io(format!("IO error: {}", e)))? {
let entry = entry.map_err(|e| LitError::io(format!("IO error: {}", e)))?;
if entry.path().extension().is_some_and(|e| e == "json") {
if let Ok(json) = fs::read_to_string(entry.path()) {
if let Ok(token) = serde_json::from_str::<UcanToken>(&json) {
if token.audience == audience_did && token.is_valid() {
tokens.push(token);
}
}
}
}
}
Ok(tokens)
}
pub fn revoke_token(repo_root: &Path, cid_prefix: &str) -> Result<(), LitError> {
let dir = ucan_dir(repo_root);
let path = dir.join(format!("{}.json", cid_prefix));
if path.exists() {
fs::remove_file(&path)
.map_err(|e| LitError::io(format!("Failed to revoke token: {}", e)))?;
Ok(())
} else {
Err(LitError::general(format!(
"Token not found: {}",
cid_prefix
)))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ucan_creation() {
let token = UcanToken::new(
"did:lit:issuer123".to_string(),
"did:lit:audience456".to_string(),
vec![Capability {
resource: "branch:main".to_string(),
action: "push".to_string(),
caveats: None,
}],
3600,
);
assert!(token.is_valid());
assert!(token.has_capability("branch:main", "push"));
assert!(!token.has_capability("branch:main", "delete"));
}
#[test]
fn test_ucan_expiration() {
let mut token = UcanToken::new(
"did:lit:a".to_string(),
"did:lit:b".to_string(),
vec![],
-1, );
token.not_before = None;
assert!(!token.is_valid());
}
#[test]
fn test_ucan_delegation() {
let parent = UcanToken::new(
"did:lit:root".to_string(),
"did:lit:agent1".to_string(),
vec![Capability {
resource: "branch:*".to_string(),
action: "*".to_string(),
caveats: None,
}],
3600,
);
let child = parent
.delegate(
"did:lit:agent2".to_string(),
vec![Capability {
resource: "branch:feature".to_string(),
action: "push".to_string(),
caveats: None,
}],
1800,
)
.unwrap();
assert_eq!(child.issuer, "did:lit:agent1");
assert_eq!(child.audience, "did:lit:agent2");
assert!(!child.proof.is_empty());
}
}