use anyhow::Result;
use chrono::Utc;
use rusqlite::{params, Connection, OptionalExtension};
use serde_json::Value as JsonValue;
use dragoon_proto::pubkey::{fingerprint_pubkey_blob, parse_pubkey_blob};
use crate::auth;
#[derive(Debug, Clone)]
pub struct UserRow {
pub id: i64,
pub username: String,
pub password_hash: String,
pub totp_secret: String,
pub recovery_codes_hash: Vec<String>,
pub revoked_at: Option<String>,
}
fn iso_now() -> String {
Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true)
}
pub fn create_user(
conn: &Connection,
username: &str,
password: &str,
) -> Result<(i64, String, Vec<String>)> {
let pw_hash = auth::hash_password(password)?;
let totp_secret = auth::generate_totp_secret();
let (plain, hashes) = auth::generate_recovery_codes(10);
let codes_json = serde_json::to_string(&hashes)?;
conn.execute(
"INSERT INTO users
(username, password_hash, totp_secret_enc, recovery_codes_hash, created_at)
VALUES (?,?,?,?,?)",
params![username, pw_hash, totp_secret, codes_json, iso_now()],
)?;
Ok((conn.last_insert_rowid(), totp_secret, plain))
}
pub fn get_user(conn: &Connection, username: &str) -> Result<Option<UserRow>> {
let row = conn
.query_row(
"SELECT id, username, password_hash, totp_secret_enc, recovery_codes_hash, revoked_at
FROM users WHERE username=?",
[username],
|r| {
Ok((
r.get::<_, i64>(0)?,
r.get::<_, String>(1)?,
r.get::<_, String>(2)?,
r.get::<_, String>(3)?,
r.get::<_, Option<String>>(4)?,
r.get::<_, Option<String>>(5)?,
))
},
)
.optional()?;
let Some((id, username, password_hash, totp_secret, codes_json, revoked_at)) = row else {
return Ok(None);
};
let recovery_codes_hash: Vec<String> =
serde_json::from_str(codes_json.as_deref().unwrap_or("[]"))?;
Ok(Some(UserRow {
id,
username,
password_hash,
totp_secret,
recovery_codes_hash,
revoked_at,
}))
}
pub fn revoke_user(conn: &Connection, username: &str) -> Result<()> {
conn.execute(
"UPDATE users SET revoked_at=? WHERE username=?",
params![iso_now(), username],
)?;
Ok(())
}
pub fn set_password(conn: &Connection, username: &str, new_password: &str) -> Result<()> {
let h = auth::hash_password(new_password)?;
conn.execute(
"UPDATE users SET password_hash=? WHERE username=?",
params![h, username],
)?;
Ok(())
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct PubkeyRow {
pub fingerprint: String,
pub alg: String,
pub label: Option<String>,
pub added_at: String,
pub revoked_at: Option<String>,
}
pub fn add_pubkey(
conn: &Connection,
user_id: i64,
pubkey_blob: &[u8],
label: Option<&str>,
) -> Result<String> {
let parsed = parse_pubkey_blob(pubkey_blob)?;
let fp = fingerprint_pubkey_blob(pubkey_blob);
conn.execute(
"INSERT INTO user_pubkeys (user_id, fingerprint, alg, pubkey_blob, label, added_at)
VALUES (?,?,?,?,?,?)",
params![user_id, fp, parsed.alg(), pubkey_blob, label, iso_now()],
)?;
Ok(fp)
}
pub fn list_pubkeys(conn: &Connection, user_id: i64) -> Result<Vec<PubkeyRow>> {
let mut stmt = conn.prepare(
"SELECT fingerprint, alg, label, added_at, revoked_at FROM user_pubkeys
WHERE user_id=? ORDER BY added_at ASC",
)?;
let rows = stmt
.query_map([user_id], |r| {
Ok(PubkeyRow {
fingerprint: r.get(0)?,
alg: r.get(1)?,
label: r.get(2)?,
added_at: r.get(3)?,
revoked_at: r.get(4)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
pub fn revoke_pubkey(conn: &Connection, user_id: i64, fingerprint: &str) -> Result<bool> {
let n = conn.execute(
"UPDATE user_pubkeys SET revoked_at=? WHERE user_id=? AND fingerprint=?",
params![iso_now(), user_id, fingerprint],
)?;
Ok(n > 0)
}
pub fn lookup_user_id_by_name(conn: &Connection, username: &str) -> Result<Option<i64>> {
let id: Option<i64> = conn
.query_row(
"SELECT id FROM users WHERE username=? AND revoked_at IS NULL",
[username],
|r| r.get(0),
)
.optional()?;
Ok(id)
}
#[allow(dead_code)]
fn _silence() -> JsonValue {
JsonValue::Null
}
#[cfg(test)]
mod tests {
use super::*;
use ed25519_dalek::SigningKey;
fn fresh() -> Connection {
let c = crate::db::connect_in_memory().unwrap();
crate::db::bootstrap(&c).unwrap();
c
}
fn ed25519_pub_blob(seed: u8) -> Vec<u8> {
let sk = SigningKey::from_bytes(&[seed; 32]);
dragoon_proto::task_sig::ed25519_public_blob(&sk)
}
#[test]
fn create_user_returns_totp_and_codes() {
let c = fresh();
let (uid, secret, codes) = create_user(&c, "alice", "hunter2").unwrap();
assert!(uid > 0);
assert!(!secret.is_empty());
assert_eq!(codes.len(), 10);
}
#[test]
fn get_user_round_trip() {
let c = fresh();
let (uid, _, _) = create_user(&c, "alice", "x").unwrap();
let row = get_user(&c, "alice").unwrap().unwrap();
assert_eq!(row.id, uid);
assert!(auth::verify_password("x", &row.password_hash));
assert_eq!(row.recovery_codes_hash.len(), 10);
}
#[test]
fn add_list_revoke_pubkey() {
let c = fresh();
let (uid, _, _) = create_user(&c, "alice", "pw").unwrap();
let blob = ed25519_pub_blob(11);
let fp = add_pubkey(&c, uid, &blob, Some("laptop")).unwrap();
assert!(fp.starts_with("SHA256:"));
let keys = list_pubkeys(&c, uid).unwrap();
assert_eq!(keys.len(), 1);
assert_eq!(keys[0].fingerprint, fp);
assert_eq!(keys[0].label.as_deref(), Some("laptop"));
assert!(revoke_pubkey(&c, uid, &fp).unwrap());
let keys = list_pubkeys(&c, uid).unwrap();
assert!(keys[0].revoked_at.is_some());
}
#[test]
fn lookup_user_id_helper() {
let c = fresh();
let (uid, _, _) = create_user(&c, "bob", "x").unwrap();
assert_eq!(lookup_user_id_by_name(&c, "bob").unwrap(), Some(uid));
assert_eq!(lookup_user_id_by_name(&c, "ghost").unwrap(), None);
}
}