#![allow(dead_code)]
use argon2::{
password_hash::{rand_core::OsRng, PasswordHash, PasswordHasher, PasswordVerifier, SaltString},
Argon2,
};
use cozo::{Db, SqliteStorage};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use uuid::Uuid;
pub type KeysDb = Db<SqliteStorage>;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ApiKey {
pub id: String,
pub name: String,
pub key_hash: String,
pub created_at: String,
pub last_used_at: Option<String>,
pub revoked_at: Option<String>,
}
pub struct ApiKeyStore {
db_path: std::path::PathBuf,
}
impl ApiKeyStore {
pub fn new() -> Result<Self, Box<dyn std::error::Error>> {
let home = dirs::home_dir().ok_or("Cannot find home directory")?;
let keys_dir = home.join(".leankg");
std::fs::create_dir_all(&keys_dir)?;
let db_path = keys_dir.join("keys.db");
Ok(Self { db_path })
}
pub fn init_db(&self) -> Result<KeysDb, Box<dyn std::error::Error>> {
let path_str = self.db_path.to_string_lossy().to_string();
let db = cozo::new_cozo_sqlite(path_str)?;
let check_relations = r#"::relations"#;
let relations_result = db.run_script(check_relations, Default::default())?;
let existing_relations: std::collections::HashSet<String> = relations_result
.rows
.iter()
.filter_map(|row| row.first().and_then(|v| v.as_str().map(String::from)))
.collect();
if !existing_relations.contains("api_keys") {
let create_table = r#":create api_keys {id: String, name: String, key_hash: String, created_at: String, last_used_at: String?, revoked_at: String?}"#;
db.run_script(create_table, Default::default())?;
}
Ok(db)
}
pub fn create_key(&self, name: &str) -> Result<(String, ApiKey), Box<dyn std::error::Error>> {
let db = self.init_db()?;
let key = generate_api_key();
let key_id = Uuid::new_v4().to_string();
let key_hash = hash_api_key(&key)?;
let created_at = chrono_timestamp();
let mut params = BTreeMap::new();
params.insert("id".to_string(), serde_json::json!(key_id));
params.insert("name".to_string(), serde_json::json!(name));
params.insert("key_hash".to_string(), serde_json::json!(key_hash));
params.insert("created_at".to_string(), serde_json::json!(created_at));
params.insert(
"last_used_at".to_string(),
serde_json::json!(serde_json::Value::Null),
);
params.insert(
"revoked_at".to_string(),
serde_json::json!(serde_json::Value::Null),
);
let insert = r#"
?[id, name, key_hash, created_at, last_used_at, revoked_at] <- [[$id, $name, $key_hash, $created_at, $last_used_at, $revoked_at]]
:put api_keys { id, name, key_hash, created_at, last_used_at, revoked_at }
"#;
db.run_script(insert, params)?;
let api_key = ApiKey {
id: key_id,
name: name.to_string(),
key_hash,
created_at,
last_used_at: None,
revoked_at: None,
};
Ok((key, api_key))
}
pub fn list_keys(&self) -> Result<Vec<ApiKey>, Box<dyn std::error::Error>> {
let db = self.init_db()?;
let query = r#"
?[id, name, key_hash, created_at, last_used_at, revoked_at] := *api_keys[id, name, key_hash, created_at, last_used_at, revoked_at]
"#;
let result = db.run_script(query, BTreeMap::new())?;
let mut keys: std::collections::HashMap<String, ApiKey> = std::collections::HashMap::new();
for row in result.rows {
let id = row[0].as_str().unwrap_or("").to_string();
let name = row[1].as_str().unwrap_or("").to_string();
let key_hash = String::new();
let created_at = row[3].as_str().unwrap_or("").to_string();
let last_used_at = row[4].as_str().map(String::from);
let revoked_at: Option<String> = row[5].as_str().map(String::from);
if revoked_at.is_some() {
keys.remove(&id);
continue;
}
if !keys.contains_key(&id) {
keys.insert(
id.clone(),
ApiKey {
id,
name,
key_hash,
created_at,
last_used_at,
revoked_at,
},
);
}
}
Ok(keys.into_values().collect())
}
pub fn revoke_key(&self, id: &str) -> Result<bool, Box<dyn std::error::Error>> {
let db = self.init_db()?;
let query = r#"
?[id, name, key_hash, created_at, last_used_at, revoked_at] := *api_keys[id, name, key_hash, created_at, last_used_at, revoked_at], id = $id
"#;
let mut params = BTreeMap::new();
params.insert("id".to_string(), serde_json::json!(id));
let result = db.run_script(query, params)?;
if result.rows.is_empty() {
return Ok(false);
}
let row = &result.rows[0];
let revoked_at_val = row[5].as_str();
if revoked_at_val.is_some() {
return Ok(false);
}
let revoked_at = chrono_timestamp();
let update = r#"
?[id, name, key_hash, created_at, last_used_at, revoked_at] <- [[$id, $name, $key_hash, $created_at, $last_used_at, $revoked_at]]
:put api_keys { id, name, key_hash, created_at, last_used_at, revoked_at }
"#;
let mut params = BTreeMap::new();
params.insert(
"id".to_string(),
serde_json::json!(row[0].as_str().unwrap_or("")),
);
params.insert(
"name".to_string(),
serde_json::json!(row[1].as_str().unwrap_or("")),
);
params.insert(
"key_hash".to_string(),
serde_json::json!(row[2].as_str().unwrap_or("")),
);
params.insert(
"created_at".to_string(),
serde_json::json!(row[3].as_str().unwrap_or("")),
);
params.insert(
"last_used_at".to_string(),
serde_json::json!(row[4].as_str().map(String::from)),
);
params.insert("revoked_at".to_string(), serde_json::json!(revoked_at));
db.run_script(update, params)?;
Ok(true)
}
pub fn validate_key(&self, key: &str) -> Result<Option<String>, Box<dyn std::error::Error>> {
let db = self.init_db()?;
let query = r#"
?[id, key_hash] := *api_keys[id, key_hash], revoked_at = null
"#;
let result = db.run_script(query, BTreeMap::new())?;
for row in result.rows {
let key_id = row[0].as_str().unwrap_or("").to_string();
let stored_hash = row[1].as_str().unwrap_or("");
if verify_api_key(key, stored_hash) {
let last_used = chrono_timestamp();
let delete = format!(r#":delete api_keys where id = "{}""#, key_id);
let _ = db.run_script(&delete, BTreeMap::new());
let insert = r#"
?[id, name, key_hash, created_at, last_used_at, revoked_at] <- [[$id, $name, $key_hash, $created_at, $last_used_at, $revoked_at]]
:put api_keys { id, name, key_hash, created_at, last_used_at, revoked_at }
"#;
let mut params = BTreeMap::new();
params.insert("id".to_string(), serde_json::json!(key_id.clone()));
params.insert("name".to_string(), serde_json::json!(""));
params.insert("key_hash".to_string(), serde_json::json!(stored_hash));
params.insert("created_at".to_string(), serde_json::json!(""));
params.insert("last_used_at".to_string(), serde_json::json!(last_used));
params.insert(
"revoked_at".to_string(),
serde_json::json!(serde_json::Value::Null),
);
let _ = db.run_script(insert, params);
return Ok(Some(key_id));
}
}
Ok(None)
}
}
fn generate_api_key() -> String {
let salt = SaltString::generate(&mut OsRng);
let key_part = Uuid::new_v4().to_string().replace("-", "");
format!("lkkg_{}_{}", key_part, &salt.as_str()[..8])
}
fn hash_api_key(key: &str) -> Result<String, String> {
let salt = SaltString::generate(&mut OsRng);
let argon2 = Argon2::default();
let hash = argon2
.hash_password(key.as_bytes(), &salt)
.map_err(|e| e.to_string())?
.to_string();
Ok(hash)
}
fn verify_api_key(key: &str, hash: &str) -> bool {
let parsed_hash = match PasswordHash::new(hash) {
Ok(h) => h,
Err(_) => return false,
};
Argon2::default()
.verify_password(key.as_bytes(), &parsed_hash)
.is_ok()
}
fn chrono_timestamp() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let duration = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default();
format!("{}", duration.as_secs())
}
impl Default for ApiKeyStore {
fn default() -> Self {
Self::new().expect("Failed to create API key store")
}
}