use rusqlite::{params, Connection}; use rust_util::{opt_result, simple_error, XResult}; pub const DEFAULT_MASTER_KEY_VERIFICATION_KEY: &str = "__master_verification_key"; pub struct Key { pub name: String, pub encrypted_key: String, pub comment: Option, } pub struct Keys { pub count: usize, pub keys: Vec, } pub fn make_value_key_name(name: &str) -> String { format!("value:{}", name) } pub fn make_data_key_name(name: &str) -> String { format!("data_key:{}", name) } pub fn open_db(db: &str) -> XResult { let con = opt_result!(Connection::open(db), "Open sqlite db: {}, failed: {}", db); log::debug!("Db auto commit: {}", con.is_autocommit()); Ok(con) } pub fn init_db(conn: &Connection) -> XResult { if let Ok(false) = check_table_keys(conn) { repair_table_keys(conn)?; } Ok(true) } fn repair_table_keys(conn: &Connection) -> XResult<()> { let field_names = list_table_fields(conn, "keys")?; let field_names = field_names.iter().map(|n| n.as_str()).collect::>(); if !field_names.contains(&"comment") { log::info!("Repair table keys, add column comment"); let _ = conn.execute("ALTER TABLE keys ADD COLUMN comment TEXT", ())?; } Ok(()) } fn list_table_fields(conn: &Connection, table: &str) -> XResult> { let mut stmt_query_fields = conn.prepare(&format!("pragma table_info({})", table))?; let mut rows_query_field = stmt_query_fields.query(())?; let mut field_names = vec![]; let mut next_query_field_opt = rows_query_field.next()?; while let Some(next_query_field) = next_query_field_opt { // cid|name|type|notnull|dflt_value|pk // ^ ^ ^ ^ ^ ^ // | | | | | [5] - Is column PK // | | | | [4] Column default value // | | | [3] Is column not null // | | [2] Column type // | [1] Column name // [0] Column index let field_name: String = next_query_field.get(1)?; field_names.push(field_name.to_lowercase()); next_query_field_opt = rows_query_field.next()?; } log::trace!("Table {} fields: {:?}", table, field_names); Ok(field_names) } fn check_table_keys(conn: &Connection) -> XResult { let mut stmt = conn.prepare("SELECT name FROM sqlite_master WHERE type='table' AND name='keys'")?; let mut rows = stmt.query(())?; if rows.next()?.is_none() { log::info!("Create table keys"); create_table_keys(conn)?; Ok(true) } else { Ok(false) } } fn create_table_keys(conn: &Connection) -> XResult<()> { let _ = conn.execute( r##" CREATE TABLE keys ( id INTEGER PRIMARY KEY, name TEXT NOT NULL, value TEXT, comment TEXT )"##, (), )?; Ok(()) } pub fn insert_key(conn: &Connection, key: &Key) -> XResult<()> { let default_comment = "".to_string(); log::debug!("insert key name={}", &key.name); let _ = conn.execute( "INSERT INTO keys (name, value, comment) VALUES (?1, ?2, ?3)", ( &key.name, &key.encrypted_key, key.comment.as_ref().unwrap_or(&default_comment), ), )?; Ok(()) } pub fn update_key(conn: &Connection, key: &Key) -> XResult<()> { log::debug!("update key name={}", &key.name); if let Some(comment) = &key.comment { let _ = conn.execute( "UPDATE keys SET value = ?1, comment = ?2 WHERE name = ?3", (&key.encrypted_key, comment, &key.name), )?; } else { let _ = conn.execute( "UPDATE keys SET value = ?1 WHERE name = ?2", (&key.encrypted_key, &key.name), )?; } Ok(()) } pub fn list_keys(conn: &Connection, ty: &str, search: &str, limit: usize) -> XResult { let name = format!("{}:%{}%", ty, search); let mut count_stmt = conn.prepare("SELECT count(*) FROM keys WHERE name like ?1")?; let mut count_iter = count_stmt.query_map(params![name], |row| { let count: usize = row.get(0)?; Ok(count) })?; let count = match count_iter.next() { None => 0, Some(Ok(count)) => count, Some(Err(e)) => return simple_error!("List keys failed: {}", e), }; log::debug!("found {} keys via: {}, limit: {}", count, name, limit); let mut keys = vec![]; if count > 0 { let mut list_stmt = conn.prepare("SELECT id, name, value, comment FROM keys WHERE name like ?1 LIMIT ?2")?; let mut list_iter = list_stmt.query_map(params![name, limit], |row| { Ok(Key { name: row.get(1)?, encrypted_key: row.get(2)?, comment: row.get(3)?, }) })?; loop { match list_iter.next() { None => { break; } Some(Ok(r)) => { log::debug!("found key name={}", r.name); keys.push(r); } Some(Err(e)) => return simple_error!("List keys failed: {}", e), } } } Ok(Keys { count, keys }) } pub fn find_key(conn: &Connection, name: &str) -> XResult> { let mut stmt = conn.prepare("SELECT id, name, value, comment FROM keys WHERE name = ?1")?; let mut key_iter = stmt.query_map(params![name], |row| { Ok(Key { name: row.get(1)?, encrypted_key: row.get(2)?, comment: row.get(3)?, }) })?; match key_iter.next() { None => { log::debug!("key name={} not exists", name); Ok(None) } Some(Ok(r)) => { log::debug!("found key name={}", name); Ok(Some(r)) } Some(Err(e)) => simple_error!("Find key failed: {}", e), } }