1use anyhow::Result;
4use chrono::Utc;
5use rusqlite::{params, Connection, OptionalExtension};
6use serde_json::Value as JsonValue;
7
8use dragoon_proto::pubkey::{fingerprint_pubkey_blob, parse_pubkey_blob};
9
10use crate::auth;
11
12#[derive(Debug, Clone)]
13pub struct UserRow {
14 pub id: i64,
15 pub username: String,
16 pub password_hash: String,
17 pub totp_secret: String,
18 pub recovery_codes_hash: Vec<String>,
19 pub revoked_at: Option<String>,
20}
21
22fn iso_now() -> String {
23 Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true)
24}
25
26pub fn create_user(
28 conn: &Connection,
29 username: &str,
30 password: &str,
31) -> Result<(i64, String, Vec<String>)> {
32 let pw_hash = auth::hash_password(password)?;
33 let totp_secret = auth::generate_totp_secret();
34 let (plain, hashes) = auth::generate_recovery_codes(10);
35 let codes_json = serde_json::to_string(&hashes)?;
36 conn.execute(
37 "INSERT INTO users
38 (username, password_hash, totp_secret_enc, recovery_codes_hash, created_at)
39 VALUES (?,?,?,?,?)",
40 params![username, pw_hash, totp_secret, codes_json, iso_now()],
41 )?;
42 Ok((conn.last_insert_rowid(), totp_secret, plain))
43}
44
45pub fn get_user(conn: &Connection, username: &str) -> Result<Option<UserRow>> {
46 let row = conn
47 .query_row(
48 "SELECT id, username, password_hash, totp_secret_enc, recovery_codes_hash, revoked_at
49 FROM users WHERE username=?",
50 [username],
51 |r| {
52 Ok((
53 r.get::<_, i64>(0)?,
54 r.get::<_, String>(1)?,
55 r.get::<_, String>(2)?,
56 r.get::<_, String>(3)?,
57 r.get::<_, Option<String>>(4)?,
58 r.get::<_, Option<String>>(5)?,
59 ))
60 },
61 )
62 .optional()?;
63 let Some((id, username, password_hash, totp_secret, codes_json, revoked_at)) = row else {
64 return Ok(None);
65 };
66 let recovery_codes_hash: Vec<String> =
67 serde_json::from_str(codes_json.as_deref().unwrap_or("[]"))?;
68 Ok(Some(UserRow {
69 id,
70 username,
71 password_hash,
72 totp_secret,
73 recovery_codes_hash,
74 revoked_at,
75 }))
76}
77
78pub fn revoke_user(conn: &Connection, username: &str) -> Result<()> {
79 conn.execute(
80 "UPDATE users SET revoked_at=? WHERE username=?",
81 params![iso_now(), username],
82 )?;
83 Ok(())
84}
85
86pub fn set_password(conn: &Connection, username: &str, new_password: &str) -> Result<()> {
87 let h = auth::hash_password(new_password)?;
88 conn.execute(
89 "UPDATE users SET password_hash=? WHERE username=?",
90 params![h, username],
91 )?;
92 Ok(())
93}
94
95#[derive(Debug, Clone, serde::Serialize)]
96pub struct PubkeyRow {
97 pub fingerprint: String,
98 pub alg: String,
99 pub label: Option<String>,
100 pub added_at: String,
101 pub revoked_at: Option<String>,
102}
103
104pub fn add_pubkey(
105 conn: &Connection,
106 user_id: i64,
107 pubkey_blob: &[u8],
108 label: Option<&str>,
109) -> Result<String> {
110 let parsed = parse_pubkey_blob(pubkey_blob)?;
111 let fp = fingerprint_pubkey_blob(pubkey_blob);
112 conn.execute(
113 "INSERT INTO user_pubkeys (user_id, fingerprint, alg, pubkey_blob, label, added_at)
114 VALUES (?,?,?,?,?,?)",
115 params![user_id, fp, parsed.alg(), pubkey_blob, label, iso_now()],
116 )?;
117 Ok(fp)
118}
119
120pub fn list_pubkeys(conn: &Connection, user_id: i64) -> Result<Vec<PubkeyRow>> {
121 let mut stmt = conn.prepare(
122 "SELECT fingerprint, alg, label, added_at, revoked_at FROM user_pubkeys
123 WHERE user_id=? ORDER BY added_at ASC",
124 )?;
125 let rows = stmt
126 .query_map([user_id], |r| {
127 Ok(PubkeyRow {
128 fingerprint: r.get(0)?,
129 alg: r.get(1)?,
130 label: r.get(2)?,
131 added_at: r.get(3)?,
132 revoked_at: r.get(4)?,
133 })
134 })?
135 .collect::<rusqlite::Result<Vec<_>>>()?;
136 Ok(rows)
137}
138
139pub fn revoke_pubkey(conn: &Connection, user_id: i64, fingerprint: &str) -> Result<bool> {
140 let n = conn.execute(
141 "UPDATE user_pubkeys SET revoked_at=? WHERE user_id=? AND fingerprint=?",
142 params![iso_now(), user_id, fingerprint],
143 )?;
144 Ok(n > 0)
145}
146
147pub fn lookup_user_id_by_name(conn: &Connection, username: &str) -> Result<Option<i64>> {
149 let id: Option<i64> = conn
150 .query_row(
151 "SELECT id FROM users WHERE username=? AND revoked_at IS NULL",
152 [username],
153 |r| r.get(0),
154 )
155 .optional()?;
156 Ok(id)
157}
158
159#[allow(dead_code)]
161fn _silence() -> JsonValue {
162 JsonValue::Null
163}
164
165#[cfg(test)]
166mod tests {
167 use super::*;
168 use ed25519_dalek::SigningKey;
169
170 fn fresh() -> Connection {
171 let c = crate::db::connect_in_memory().unwrap();
172 crate::db::bootstrap(&c).unwrap();
173 c
174 }
175
176 fn ed25519_pub_blob(seed: u8) -> Vec<u8> {
177 let sk = SigningKey::from_bytes(&[seed; 32]);
178 dragoon_proto::task_sig::ed25519_public_blob(&sk)
179 }
180
181 #[test]
182 fn create_user_returns_totp_and_codes() {
183 let c = fresh();
184 let (uid, secret, codes) = create_user(&c, "alice", "hunter2").unwrap();
185 assert!(uid > 0);
186 assert!(!secret.is_empty());
187 assert_eq!(codes.len(), 10);
188 }
189
190 #[test]
191 fn get_user_round_trip() {
192 let c = fresh();
193 let (uid, _, _) = create_user(&c, "alice", "x").unwrap();
194 let row = get_user(&c, "alice").unwrap().unwrap();
195 assert_eq!(row.id, uid);
196 assert!(auth::verify_password("x", &row.password_hash));
197 assert_eq!(row.recovery_codes_hash.len(), 10);
198 }
199
200 #[test]
201 fn add_list_revoke_pubkey() {
202 let c = fresh();
203 let (uid, _, _) = create_user(&c, "alice", "pw").unwrap();
204 let blob = ed25519_pub_blob(11);
205 let fp = add_pubkey(&c, uid, &blob, Some("laptop")).unwrap();
206 assert!(fp.starts_with("SHA256:"));
207 let keys = list_pubkeys(&c, uid).unwrap();
208 assert_eq!(keys.len(), 1);
209 assert_eq!(keys[0].fingerprint, fp);
210 assert_eq!(keys[0].label.as_deref(), Some("laptop"));
211 assert!(revoke_pubkey(&c, uid, &fp).unwrap());
212 let keys = list_pubkeys(&c, uid).unwrap();
213 assert!(keys[0].revoked_at.is_some());
214 }
215
216 #[test]
217 fn lookup_user_id_helper() {
218 let c = fresh();
219 let (uid, _, _) = create_user(&c, "bob", "x").unwrap();
220 assert_eq!(lookup_user_id_by_name(&c, "bob").unwrap(), Some(uid));
221 assert_eq!(lookup_user_id_by_name(&c, "ghost").unwrap(), None);
222 }
223}