use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum Scope {
#[default]
Session,
Persistent,
}
impl std::fmt::Display for Scope {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Scope::Session => write!(f, "Session"),
Scope::Persistent => write!(f, "Persistent"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Grant {
pub tool: String,
pub params_hash: Option<String>,
#[serde(default)]
pub scope: Scope,
pub created_at: DateTime<Utc>,
}
impl Grant {
pub fn tool(name: impl Into<String>) -> Self {
Self {
tool: name.into(),
params_hash: None,
scope: Scope::default(),
created_at: Utc::now(),
}
}
pub fn exact(name: impl Into<String>, params_hash: impl Into<String>) -> Self {
Self {
tool: name.into(),
params_hash: Some(params_hash.into()),
scope: Scope::default(),
created_at: Utc::now(),
}
}
pub fn with_scope(mut self, scope: Scope) -> Self {
self.scope = scope;
self
}
pub fn is_tool_wide(&self) -> bool {
self.params_hash.is_none()
}
pub fn matches(&self, params_hash: &str) -> bool {
match &self.params_hash {
None => true, Some(h) => h == params_hash,
}
}
}
impl PartialEq for Grant {
fn eq(&self, other: &Self) -> bool {
self.tool == other.tool
&& self.params_hash == other.params_hash
&& self.scope == other.scope
}
}
impl Eq for Grant {}
pub fn hash_params(params: &serde_json::Value) -> String {
use sha2::{Digest, Sha256};
let canonical = canonicalize_json(params);
let json = serde_json::to_string(&canonical).unwrap_or_default();
let hash = Sha256::digest(json.as_bytes());
format!("{:x}", hash)
}
fn canonicalize_json(value: &serde_json::Value) -> serde_json::Value {
use serde_json::Value;
use std::collections::BTreeMap;
match value {
Value::Object(map) => {
let sorted: BTreeMap<_, _> = map
.iter()
.map(|(k, v)| (k.clone(), canonicalize_json(v)))
.collect();
Value::Object(sorted.into_iter().collect())
}
Value::Array(arr) => Value::Array(arr.iter().map(canonicalize_json).collect()),
other => other.clone(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_grant_tool() {
let grant = Grant::tool("echo");
assert_eq!(grant.tool, "echo");
assert!(grant.params_hash.is_none());
assert!(grant.is_tool_wide());
assert_eq!(grant.scope, Scope::Session);
}
#[test]
fn test_grant_exact() {
let grant = Grant::exact("database", "abc123");
assert_eq!(grant.tool, "database");
assert_eq!(grant.params_hash, Some("abc123".to_string()));
assert!(!grant.is_tool_wide());
}
#[test]
fn test_grant_with_scope() {
let grant = Grant::tool("test").with_scope(Scope::Persistent);
assert_eq!(grant.scope, Scope::Persistent);
}
#[test]
fn test_grant_matches() {
let tool_grant = Grant::tool("test");
assert!(tool_grant.matches("any_hash"));
assert!(tool_grant.matches("other_hash"));
let exact_grant = Grant::exact("test", "specific_hash");
assert!(exact_grant.matches("specific_hash"));
assert!(!exact_grant.matches("other_hash"));
}
#[test]
fn test_hash_params() {
let params = serde_json::json!({"key": "value"});
let hash = hash_params(¶ms);
assert!(!hash.is_empty());
assert_eq!(hash.len(), 64);
let params2 = serde_json::json!({"key": "value"});
assert_eq!(hash_params(¶ms2), hash);
let params3 = serde_json::json!({"key": "other"});
assert_ne!(hash_params(¶ms3), hash);
}
#[test]
fn test_hash_params_canonical_order() {
let params1 = serde_json::json!({"a": 1, "b": 2, "c": 3});
let params2 = serde_json::json!({"c": 3, "b": 2, "a": 1});
assert_eq!(hash_params(¶ms1), hash_params(¶ms2));
let nested1 = serde_json::json!({"outer": {"z": 1, "a": 2}});
let nested2 = serde_json::json!({"outer": {"a": 2, "z": 1}});
assert_eq!(hash_params(&nested1), hash_params(&nested2));
}
#[test]
fn test_scope_display() {
assert_eq!(Scope::Session.to_string(), "Session");
assert_eq!(Scope::Persistent.to_string(), "Persistent");
}
#[test]
fn test_grant_equality() {
let g1 = Grant::tool("test");
let g2 = Grant::tool("test");
assert_eq!(g1, g2);
let g3 = Grant::exact("test", "hash");
assert_ne!(g1, g3); }
#[test]
fn test_grant_serialization() {
let grant = Grant::exact("tool", "hash123").with_scope(Scope::Persistent);
let json = serde_json::to_string(&grant).unwrap();
let parsed: Grant = serde_json::from_str(&json).unwrap();
assert_eq!(grant.tool, parsed.tool);
assert_eq!(grant.params_hash, parsed.params_hash);
assert_eq!(grant.scope, parsed.scope);
}
}