foil/src/db.rs

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

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

pub struct DbHandle {
    conn: Connection,
}

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

#[derive(Debug)]
pub struct Account {
    pub id: i64,
    pub host: String,
    pub user: String,
    pub password: String,
}

impl DbHandle {
    pub fn new(db_path: &Option<String>) -> DbHandle {
        let path: PathBuf = db_path
            .as_ref()
            .map(|path: &String| PathBuf::from(path))
            .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: conn }
    }

    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: &Vec<u8>) {
        let b64value: String = value.to_base64(base64::STANDARD);
        self.set_db_property(name, &b64value);
    }

    pub fn list_accounts(&mut self, master_key: [u8; 32]) -> Vec<Account> {
        let mut stmt = self.conn
            .prepare("SELECT h.iv, h.ciphertext, h.mac, u.iv, u.ciphertext, u.mac, p.iv, p.ciphertext, p.mac, a.id FROM accounts a \
                      LEFT JOIN encrypted_values h ON a.server=h.id \
                      LEFT JOIN encrypted_values u ON a.user=u.id \
                      LEFT JOIN encrypted_values p ON a.password=p.id")
            .unwrap();
        let result: Vec<Account> = {
            let props = stmt
                .query_map(NO_PARAMS, |row| {
                    let host: String = String::from_utf8(decrypt_base64(
                        row.get(0).unwrap(),
                        row.get(1).unwrap(),
                        row.get(2).unwrap(),
                        master_key,
                    ))
                    .unwrap();
                    let user: String = String::from_utf8(decrypt_base64(
                        row.get(3).unwrap(),
                        row.get(4).unwrap(),
                        row.get(5).unwrap(),
                        master_key,
                    ))
                    .unwrap();
                    let password: String = String::from_utf8(decrypt_base64(
                        row.get(6).unwrap(),
                        row.get(7).unwrap(),
                        row.get(8).unwrap(),
                        master_key,
                    ))
                    .unwrap();
                    Ok(Account {
                        id: row.get(9).unwrap(),
                        host: host,
                        user: user,
                        password: password,
                    })
                })
                .unwrap();

            props.map(move |x| x.unwrap()).collect()
        };
        result
    }

    pub fn write_encrypted_value(&mut self, val: EncryptedValue) -> i64 {
        let b64ciphertext: String = val.ciphertext.to_base64(base64::STANDARD);
        let b64iv: String = val.iv.to_base64(base64::STANDARD);
        let b64mac: String = val.mac.code().to_base64(base64::STANDARD);
        let tx = self.conn.transaction().unwrap();
        tx.execute(
            "INSERT INTO encrypted_values (iv, ciphertext, mac) VALUES ($1, $2, $3)",
            &[&b64iv, &b64ciphertext, &b64mac],
        )
        .unwrap();
        let rowid: i64 = tx.last_insert_rowid();
        let _ = tx.commit().unwrap();
        rowid
    }

    pub fn write_account(
        &mut self,
        host: EncryptedValue,
        username: EncryptedValue,
        password: EncryptedValue,
    ) -> i64 {
        // TODO: This should be a transaction
        let host_id = self.write_encrypted_value(host);
        let user_id = self.write_encrypted_value(username);
        let password_id = self.write_encrypted_value(password);
        self.conn
            .execute(
                "INSERT INTO accounts (server, user, password) VALUES ($1, $2, $3)",
                &[&host_id, &user_id, &password_id],
            )
            .unwrap();
        self.conn.last_insert_rowid()
    }

    pub fn delete_account_with_host(&mut self, master_key: [u8; 32], host: &str) {
        let accounts: Vec<Account> = self.list_accounts(master_key);
        let tx = self.conn.transaction().unwrap();

        for account in accounts {
            if account.host == host {
                tx.execute(
                    "DELETE FROM encrypted_values WHERE exists \
                     (SELECT 1 FROM accounts WHERE \
                     (accounts.server = encrypted_values.id OR \
                     accounts.user = encrypted_values.id OR \
                     accounts.password = encrypted_values.id) AND \
                     accounts.id=$1);",
                    &[&account.id],
                )
                .unwrap();
                tx.execute("DELETE FROM accounts WHERE id=$1;", &[&account.id])
                    .unwrap();
            }
        }

        let _ = tx.commit().unwrap();
    }
}

fn decrypt_base64(iv: String, ciphertext: String, mac: String, master_key: [u8; 32]) -> Vec<u8> {
    let mut iv_bytes: [u8; 32] = [0; 32];
    let mut mac_bytes: [u8; 32] = [0; 32];

    iv_bytes.clone_from_slice(&iv.from_base64().unwrap());
    mac_bytes.clone_from_slice(&mac.from_base64().unwrap());

    let decrypted: Vec<u8> = crypt::decrypt_value(
        ciphertext.from_base64().unwrap(),
        master_key,
        iv_bytes,
        mac_bytes,
    );
    decrypted
}