foil/src/db.rs

use super::crypt;
use crate::crypt::EncryptedValue;
use rusqlite::{params, Connection};
use rustc_serialize::base64;
use rustc_serialize::base64::{FromBase64, ToBase64};
use std::error::Error;
use std::fmt;
use std::path::PathBuf;

static DB_INIT_QUERY: &str = include_str!("init.sql");

pub struct DbHandle {
    conn: Connection,
}

#[derive(Debug, Clone)]
struct DbProperty {
    name: String,
    value: Option<String>,
}

pub struct DbNamespace {
    pub id: i64,
    pub name: EncryptedValue,
}

pub struct Namespace {
    pub id: i64,
    pub name: String,
}

pub struct DbNote {
    pub id: i64,
    pub namespace: DbNamespace,
    pub category: String,
    pub title: EncryptedValue,
    pub value: EncryptedValue,
}

#[derive(Debug)]
pub struct Note {
    pub id: i64,
    pub namespace: String,
    pub category: String,
    pub title: String,
    pub value: String,
}

impl DbHandle {
    pub fn new(db_path: &Option<String>) -> DbHandle {
        let path: PathBuf = db_path
            .as_ref()
            .map(PathBuf::from)
            .unwrap_or_else(|| dirs::home_dir().unwrap().join(".foil").to_path_buf());
        let mut conn: Connection = Connection::open(path).unwrap();
        let tx = conn.transaction().unwrap();
        tx.execute_batch(DB_INIT_QUERY).unwrap();
        tx.commit().unwrap();
        DbHandle { conn }
    }

    pub fn get_namespace_id(
        &mut self,
        name: &str,
        master_key: [u8; 32],
    ) -> Result<i64, Box<dyn Error>> {
        {
            let mut stmt = self
                .conn
                .prepare("SELECT id, name FROM namespaces")
                .unwrap();
            let rows = stmt.query_map(params![], |row| {
                Ok(DbNamespace {
                    id: row.get(0)?,
                    name: row.get(1)?,
                })
            })?;

            for row_result in rows {
                let row: DbNamespace = row_result?;
                let row_name: String = row.name.decrypt_to_string(master_key)?;
                if name == row_name {
                    return Ok(row.id);
                }
            }
        }

        let new_namespace = crypt::encrypt_value(name, master_key);
        let tx = self.conn.transaction().unwrap();
        tx.execute(
            "INSERT INTO namespaces (name) VALUES ($1)",
            &[&new_namespace],
        )
        .unwrap();
        let rowid: i64 = tx.last_insert_rowid();
        tx.commit().unwrap();
        Ok(rowid)
    }

    pub fn write_note(&mut self, master_key: [u8; 32], note: Note) {
        let existing_notes = self.read_notes(master_key).unwrap();
        let namespace_id = self.get_namespace_id(&note.namespace, master_key).unwrap();
        let tx = self.conn.transaction().unwrap();

        for existing_note in existing_notes {
            if existing_note.namespace == note.namespace
                && existing_note.category == note.category
                && existing_note.title == note.title
            {
                tx.execute("DELETE FROM notes WHERE id=$1;", params![existing_note.id])
                    .unwrap();
            }
        }

        let encrypted_title = crypt::encrypt_value(&note.title, master_key);
        let encrypted_value = crypt::encrypt_value(&note.value, master_key);
        tx.execute(
            "INSERT INTO notes (namespace, category, title, value) VALUES ($1, $2, $3, $4);",
            params![
                namespace_id,
                note.category,
                encrypted_title,
                encrypted_value
            ],
        )
        .unwrap();

        tx.commit().unwrap();
    }

    pub fn read_notes(&mut self, master_key: [u8; 32]) -> rusqlite::Result<Vec<Note>> {
        let mut stmt = self.conn.prepare("SELECT notes.id, notes.category, notes.title, notes.value, namespaces.id, namespaces.name FROM notes JOIN namespaces ON notes.namespace=namespaces.id").unwrap();
        let rows = stmt.query_map(params![], |row| {
            Ok(DbNote {
                id: row.get(0)?,
                category: row.get(1)?,
                title: row.get(2)?,
                value: row.get(3)?,
                namespace: DbNamespace {
                    id: row.get(4)?,
                    name: row.get(5)?,
                },
            })
        })?;

        let notes: Result<Vec<Note>, _> = rows
            .map(|note_result| match note_result {
                Ok(note) => Ok(Note {
                    id: note.id,
                    namespace: note.namespace.name.decrypt_to_string(master_key).unwrap(),
                    category: note.category,
                    title: note.title.decrypt_to_string(master_key).unwrap(),
                    value: note.value.decrypt_to_string(master_key).unwrap(),
                }),
                Err(e) => Err(e),
            })
            .collect();

        notes
    }

    pub fn list_namespaces(&mut self, master_key: [u8; 32]) -> rusqlite::Result<Vec<Namespace>> {
        let mut stmt = self
            .conn
            .prepare("SELECT id, name FROM namespaces")
            .unwrap();
        let rows = stmt.query_map(params![], |row| {
            Ok(DbNamespace {
                id: row.get(0)?,
                name: row.get(1)?,
            })
        })?;

        let namespaces: Result<Vec<Namespace>, _> = rows
            .map(|namespace_result| match namespace_result {
                Ok(namespace) => Ok(Namespace {
                    id: namespace.id,
                    name: namespace.name.decrypt_to_string(master_key).unwrap(),
                }),
                Err(e) => Err(e),
            })
            .collect();
        namespaces
    }

    pub fn get_db_property(&self, name: &str) -> Result<Option<String>, Box<dyn Error>> {
        let mut stmt = self
            .conn
            .prepare("SELECT name, value FROM props WHERE name=$1")
            .unwrap();
        let mut props = stmt.query_map(&[&name], |row| {
            Ok(DbProperty {
                name: row.get(0)?,
                value: row.get(1)?,
            })
        })?;

        match props.next() {
            Some(prop) => Ok(Some(prop.unwrap().value.unwrap())),
            None => Ok(None),
        }
    }

    pub fn set_db_property(&self, name: &str, value: &str) {
        self.conn
            .execute(
                "INSERT OR REPLACE INTO props (name, value) VALUES ($1, $2)",
                &[&name, &value],
            )
            .unwrap();
    }

    pub fn get_db_property_bytes(&self, name: &str) -> Result<Option<Vec<u8>>, Box<dyn Error>> {
        self.get_db_property(name)
            .map(|option| option.map(|prop| prop.from_base64().unwrap()))
    }

    pub fn set_db_property_bytes(&self, name: &str, value: &[u8]) {
        let b64value: String = value.to_base64(base64::STANDARD);
        self.set_db_property(name, &b64value);
    }
}

impl fmt::Display for Namespace {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name)
    }
}