1 files changed, 118 insertions, 0 deletions

diff --git a/src/client/database.rs b/src/client/database.rs
new file mode 100644
index 0000000..e1fa42d
--- /dev/null
+++ b/src/client/database.rs
@@ -0,0 +1,118 @@
+use rusqlite::Connection;
+
+pub fn init_db(conn: &Connection) {
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS hosts (
+            host TEXT PRIMARY KEY,
+            last_used DATETIME DEFAULT CURRENT_TIMESTAMP,
+            fingerprint TEXT
+        )",
+        [],
+    )
+    .unwrap();
+    // Migration for databases created before TOFU trust.
+    let _ = conn.execute("ALTER TABLE hosts ADD COLUMN fingerprint TEXT", []);
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS credentials (
+            host TEXT PRIMARY KEY,
+            username TEXT NOT NULL,
+            password TEXT NOT NULL
+        )",
+        [],
+    )
+    .unwrap();
+}
+
+pub fn get_last_host(conn: &Connection) -> Option<String> {
+    let mut stmt = conn
+        .prepare("SELECT host FROM hosts ORDER BY last_used DESC LIMIT 1")
+        .unwrap();
+    let mut rows = stmt.query([]).unwrap();
+    if let Ok(Some(row)) = rows.next() {
+        row.get(0).ok()
+    } else {
+        None
+    }
+}
+
+pub fn save_host(conn: &Connection, host: &str) {
+    conn.execute(
+        "INSERT INTO hosts (host) VALUES (?1)
+         ON CONFLICT(host) DO UPDATE SET last_used = CURRENT_TIMESTAMP",
+        [host],
+    )
+    .unwrap();
+}
+
+pub fn save_credentials(
+    conn: &Connection,
+    host: &str,
+    username: &str,
+    password: &str,
+) -> Result<(), Box<dyn std::error::Error>> {
+    conn.execute(
+        "INSERT OR REPLACE INTO credentials (host, username, password) VALUES (?1, ?2, ?3)",
+        (host, username, password),
+    )?;
+    Ok(())
+}
+
+pub fn get_credentials(conn: &Connection, host: &str) -> Option<(String, String)> {
+    let mut stmt = conn
+        .prepare("SELECT username, password FROM credentials WHERE host = ?1")
+        .unwrap();
+    let mut rows = stmt.query([host]).unwrap();
+    if let Ok(Some(row)) = rows.next() {
+        let username: String = row.get(0).unwrap();
+        let password: String = row.get(1).unwrap();
+        Some((username, password))
+    } else {
+        None
+    }
+}
+
+/// SHA-256 SPKI fingerprint of a host's certificate, if trusted.
+pub fn get_fingerprint(conn: &Connection, host: &str) -> Option<String> {
+    let mut stmt = conn
+        .prepare("SELECT fingerprint FROM hosts WHERE host = ?1")
+        .ok()?;
+    let mut rows = stmt.query([host]).ok()?;
+    match rows.next() {
+        Ok(Some(row)) => row.get(0).ok(),
+        _ => None,
+    }
+}
+
+pub fn save_fingerprint(
+    conn: &Connection,
+    host: &str,
+    fingerprint: &str,
+) -> Result<(), rusqlite::Error> {
+    // Hosts may not exist yet (fingerprint can be saved before first
+    // successful connect completes save_host).
+    conn.execute("INSERT OR IGNORE INTO hosts (host) VALUES (?1)", [host])?;
+    conn.execute(
+        "UPDATE hosts SET fingerprint = ?2 WHERE host = ?1",
+        (host, fingerprint),
+    )?;
+    Ok(())
+}
+
+/// Remove all trust and credentials for a host ("forget this host").
+pub fn forget_host(conn: &Connection, host: &str) -> Result<(), rusqlite::Error> {
+    conn.execute("DELETE FROM hosts WHERE host = ?1", [host])?;
+    conn.execute("DELETE FROM credentials WHERE host = ?1", [host])?;
+    Ok(())
+}
+
+/// List all remembered hosts with their fingerprints.
+pub fn get_all_hosts(conn: &Connection) -> Vec<(String, Option<String>)> {
+    let mut stmt = conn
+        .prepare("SELECT host, fingerprint FROM hosts ORDER BY last_used DESC")
+        .unwrap();
+    let rows = stmt.query_map([], |row| {
+        Ok((row.get(0)?, row.get(1)?))
+    })
+    .unwrap();
+    rows.filter_map(|r| r.ok()).collect()
+}