//! Store keys and other information needed by SSH CA tool.

use crate::info::{Host, User, CA};

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// A secure store for keys.
///
/// The store will contain CA keys (both public and private) and host
/// and user public keys, as well as additional data associated with
/// the keys.
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct KeyStore {
    #[serde(skip)]
    dirty: bool,
    host_ca_keys: HashMap<String, CA>,
    user_ca_keys: HashMap<String, CA>,
    hosts: Vec<Host>,
    users: Vec<User>,
}

impl KeyStore {
    /// Create a new, unmodified store in memory.
    pub fn new() -> Self {
        Self::default()
    }

    /// Has store been modified?
    pub fn is_dirty(&self) -> bool {
        self.dirty
    }

    /// Mark store as modified.
    pub fn make_dirty(&mut self) {
        self.dirty = true;
    }

    /// Mark store as not modified.
    pub fn make_clean(&mut self) {
        self.dirty = false;
    }

    /// Serialize store as YAML.
    pub fn to_yaml(&self) -> String {
        serde_norway::to_string(self).unwrap()
    }

    /// Load store from YAML.
    pub fn from_yaml(yaml: &str) -> Result<Self, serde_norway::Error> {
        let x: Result<Self, serde_norway::Error> = serde_norway::from_str(yaml);
        match x {
            Err(err) => Err(err),
            Ok(mut store) => {
                store.make_clean();
                Ok(store)
            }
        }
    }

    /// Insert a host CA.
    pub fn insert_host_ca(&mut self, ca: CA) {
        self.make_dirty();
        self.host_ca_keys.insert(ca.name().to_string(), ca);
    }

    /// Remove a host CA.
    pub fn delete_host_ca(&mut self, name: &str) {
        self.make_dirty();
        self.host_ca_keys.remove(name);
    }

    /// List host CAs.
    pub fn host_cas(&self) -> Vec<CA> {
        self.host_ca_keys.values().cloned().collect()
    }

    /// Get host CA by name.
    pub fn get_host_ca(&self, name: &str) -> Option<&CA> {
        self.host_ca_keys.get(name)
    }

    /// Insert a user CA.
    pub fn insert_user_ca(&mut self, ca: CA) {
        self.make_dirty();
        self.user_ca_keys.insert(ca.name().to_string(), ca);
    }

    /// Remove a user CA.
    pub fn delete_user_ca(&mut self, name: &str) {
        self.make_dirty();
        self.user_ca_keys.remove(name);
    }

    /// Rename CA in store.
    pub fn rename_ca(&mut self, oldname: &str, newname: &str) -> Result<(), KeyStoreError> {
        if let Some(mut ca) = self.host_ca_keys.remove(oldname) {
            ca.rename(newname);
            self.insert_host_ca(ca);
        } else if let Some(mut ca) = self.user_ca_keys.remove(oldname) {
            ca.rename(newname);
            self.insert_host_ca(ca);
        } else {
            return Err(KeyStoreError::UnknownCa(oldname.to_string()));
        }
        self.make_dirty();
        Ok(())
    }

    /// List user CAs.
    pub fn user_cas(&self) -> Vec<CA> {
        self.user_ca_keys.values().cloned().collect()
    }

    /// Get user CA by name.
    pub fn get_user_ca(&self, name: &str) -> Option<&CA> {
        self.user_ca_keys.get(name)
    }

    /// List hosts in store.
    pub fn hosts(&self) -> &[Host] {
        &self.hosts
    }

    /// Insert a host to store.
    pub fn insert_host(&mut self, host: Host) -> Result<(), KeyStoreError> {
        if self.get_host(host.name()).is_some() {
            return Err(KeyStoreError::DuplicateHostName(host.name().into()));
        }
        self.hosts.push(host);
        self.make_dirty();
        Ok(())
    }

    /// Get host from store by name.
    pub fn get_host(&self, name: &str) -> Option<&Host> {
        self.hosts.iter().find(|&host| host.name() == name)
    }

    /// Get a mutable reference to a host in store, by name.
    pub fn get_host_mut(&mut self, name: &str) -> Option<&mut Host> {
        for h in self.hosts.iter_mut() {
            if h.name() == name {
                self.dirty = true;
                return Some(h);
            }
        }
        None
    }

    /// Rename host in store.
    pub fn rename_host(&mut self, oldname: &str, newname: &str) -> Result<(), KeyStoreError> {
        let mut renamed = false;
        for host in self.hosts.iter_mut() {
            if host.name() == oldname {
                host.rename(newname);
                renamed = true;
                break;
            }
        }
        if !renamed {
            return Err(KeyStoreError::UnknownHost(oldname.to_string()));
        }
        self.make_dirty();
        Ok(())
    }

    /// Remove host from store by name.
    pub fn remove_host(&mut self, name: &str) -> Result<(), KeyStoreError> {
        let mut i = 0;
        let mut removed = false;
        while i < self.hosts.len() {
            if self.hosts[i].name() == name {
                self.hosts.remove(i);
                removed = true;
            } else {
                i += 1;
            }
        }
        if !removed {
            return Err(KeyStoreError::UnknownHost(name.to_string()));
        }
        self.make_dirty();
        Ok(())
    }

    /// List users in store.
    pub fn users(&self) -> &[User] {
        &self.users
    }

    /// Insert user into store.
    pub fn insert_user(&mut self, user: User) -> Result<(), KeyStoreError> {
        if self.get_user(user.name()).is_some() {
            return Err(KeyStoreError::DuplicateUserName(user.name().into()));
        }
        self.users.push(user);
        self.make_dirty();
        Ok(())
    }

    /// Find user in store by name.
    pub fn get_user(&self, name: &str) -> Option<&User> {
        self.users.iter().find(|user| user.name() == name)
    }

    /// Get a mutable reference to a user in store, by name.
    pub fn get_user_mut(&mut self, name: &str) -> Option<&mut User> {
        for u in self.users.iter_mut() {
            if u.name() == name {
                self.dirty = true;
                return Some(u);
            }
        }
        None
    }

    /// Rename user in store.
    pub fn rename_user(&mut self, oldname: &str, newname: &str) -> Result<(), KeyStoreError> {
        let mut renamed = false;
        for user in self.users.iter_mut() {
            if user.name() == oldname {
                user.rename(newname);
                renamed = true;
                break;
            }
        }
        if !renamed {
            return Err(KeyStoreError::UnknownUser(oldname.to_string()));
        }
        self.make_dirty();
        Ok(())
    }

    /// Remove user from store.
    pub fn remove_user(&mut self, name: &str) -> Result<(), KeyStoreError> {
        let mut i = 0;
        let mut removed = false;
        while i < self.users.len() {
            if self.users[i].name() == name {
                self.users.remove(i);
                removed = true;
            } else {
                i += 1;
            }
        }
        if !removed {
            return Err(KeyStoreError::UnknownUser(name.to_string()));
        }
        self.make_dirty();
        Ok(())
    }
}

/// Errors related to [`KeyStore`].
#[derive(Debug, thiserror::Error)]
pub enum KeyStoreError {
    /// Duplicate host name.
    #[error("host key with host name {0} already exists in store")]
    DuplicateHostName(String),

    /// Duplicate user name.
    #[error("user {0} already exists in store")]
    DuplicateUserName(String),

    /// CA does not exist.
    #[error("CA {0} not found in the key store")]
    UnknownCa(String),

    /// Host does not exist.
    #[error("host {0} not found in the key store")]
    UnknownHost(String),

    /// User does not exist.
    #[error("user {0} not found in the key store")]
    UnknownUser(String),
}

#[cfg(test)]
mod test {
    use super::{Host, KeyStore, KeyStoreError, User, CA};
    use crate::key::{KeyPair, PrivateKey, PublicKey, SshKey};

    fn dummy_ca() -> CA {
        let pair = KeyPair::new(
            PublicKey::new("dummy-public".to_string()),
            PrivateKey::new("dummy-private".to_string()),
        );
        CA::new("CAv1", pair)
    }

    fn dummy_host() -> Host {
        Host::new(
            PublicKey::new("dummy-public".into()),
            Some(PrivateKey::new("dummy-private".into())),
            "dummy-host".into(),
        )
        .with_principals(vec!["dummy-host".into()])
    }

    fn dummy_user() -> User {
        User::new(
            PublicKey::new("dummy-user".into()),
            "dummy".into(),
            vec!["dummy".into()],
        )
    }

    fn host_ca_names(store: &KeyStore) -> Vec<String> {
        store.host_cas().iter().map(|ca| ca.name().into()).collect()
    }

    fn user_ca_names(store: &KeyStore) -> Vec<String> {
        store.user_cas().iter().map(|ca| ca.name().into()).collect()
    }

    fn hostnames(store: &KeyStore) -> Vec<String> {
        store.hosts().iter().map(|h| h.name().into()).collect()
    }

    #[test]
    fn store_is_clean_initially() {
        let store = KeyStore::new();
        assert!(!store.is_dirty());
    }

    #[test]
    fn store_is_empty_initially() {
        let store = KeyStore::new();
        assert_eq!(store.host_cas().len(), 0);
    }

    #[test]
    fn adds_updates_removes_host_ca() {
        let ca = dummy_ca();
        let mut store = KeyStore::new();
        store.insert_host_ca(ca.clone());
        assert!(store.is_dirty());
        assert_eq!(store.host_cas().len(), 1);

        let ca2 = store.get_host_ca(ca.name());
        assert_eq!(ca2, Some(&ca));

        store.delete_host_ca(ca.name());
        assert_eq!(store.host_cas().len(), 0);
        assert!(store.get_host_ca(ca.name()).is_none());
    }

    #[test]
    fn can_be_cleaned() {
        let mut store = KeyStore::new();
        store.insert_host_ca(dummy_ca());
        assert!(store.is_dirty());
        store.make_clean();
        assert!(!store.is_dirty());
    }

    #[test]
    fn returns_no_host_ca_given_nonexistent_name() {
        let store = KeyStore::new();
        assert!(store.get_host_ca("CAv1").is_none());
    }

    #[test]
    fn deletes_nonexistent_host_ca() {
        let mut store = KeyStore::new();
        store.insert_host_ca(dummy_ca());
        store.delete_host_ca("CAv2");
        assert_eq!(store.host_cas().len(), 1);
        assert_eq!(host_ca_names(&store), &["CAv1"]);
    }

    #[test]
    fn adds_updates_removes_user_ca() {
        let ca = dummy_ca();
        let mut store = KeyStore::new();
        store.insert_user_ca(ca.clone());
        assert!(store.is_dirty());
        assert_eq!(store.user_cas().len(), 1);

        let ca2 = store.get_user_ca(ca.name());
        assert_eq!(ca2, Some(&ca));

        store.delete_user_ca(ca.name());
        assert_eq!(store.user_cas().len(), 0);
        assert!(store.get_user_ca(ca.name()).is_none());
    }

    #[test]
    fn returns_no_user_ca_key_given_nonexistent_name() {
        let store = KeyStore::new();
        assert!(store.get_user_ca("CAv1").is_none());
    }

    #[test]
    fn deletes_nonexistent_user_ca_key() {
        let mut store = KeyStore::new();
        store.insert_user_ca(dummy_ca());
        store.delete_user_ca("CAv2");
        assert_eq!(store.user_cas().len(), 1);
        assert_eq!(user_ca_names(&store), &["CAv1"]);
    }

    #[test]
    fn round_trip_via_yaml() {
        let yaml = {
            let mut store = KeyStore::new();
            store.insert_host_ca(dummy_ca());
            store.to_yaml()
        };

        let store = KeyStore::from_yaml(&yaml).unwrap();
        assert!(!store.is_dirty());
        assert_eq!(host_ca_names(&store), &["CAv1"]);

        let info = store.get_host_ca("CAv1");
        assert!(info.is_some());
        let ca = info.unwrap();
        assert_eq!(
            format!("{}", ca.keypair().public().as_str()),
            "dummy-public"
        );
    }

    #[test]
    fn has_no_hosts_initially() {
        let store = KeyStore::new();
        assert!(store.hosts().is_empty());
    }

    #[test]
    fn adds_updates_removes_host() {
        let mut store = KeyStore::new();

        let host = dummy_host();
        store.insert_host(host.clone()).unwrap();
        assert!(store.is_dirty());
        assert_eq!(store.hosts().len(), 1);

        let host2 = store.get_host(host.name());
        assert_eq!(host2, Some(&host));

        store.remove_host(host.name()).unwrap();
        assert_eq!(store.hosts().len(), 0);
        assert!(store.get_host(host.name()).is_none());
    }

    #[test]
    fn refuses_host_name_collision() {
        let mut store = KeyStore::new();
        let host = dummy_host();
        assert!(store.insert_host(host.clone()).is_ok());
        assert!(store.insert_host(host).is_err());
        assert_eq!(hostnames(&store).len(), 1);
    }

    #[test]
    fn returns_error_if_cant_find_host() {
        let store = KeyStore::new();
        assert!(store.get_host("host").is_none());
    }

    #[test]
    fn fails_if_removing_unknown_host() {
        let mut store = KeyStore::new();
        let r = store.remove_host("foo.example.com");
        match r {
            Err(KeyStoreError::UnknownHost(_)) => (),
            _ => panic!("could not remove host from store"),
        }
    }

    #[test]
    fn has_no_users_initially() {
        let store = KeyStore::new();
        assert!(store.users().is_empty());
    }

    #[test]
    fn adds_updates_removes_user() {
        let mut store = KeyStore::new();

        let user = dummy_user();
        store.insert_user(user.clone()).unwrap();
        assert!(store.is_dirty());
        assert_eq!(store.users().len(), 1);

        let user2 = store.get_user(user.name());
        assert_eq!(user2, Some(&user));

        store.remove_user(user.name()).unwrap();
        assert_eq!(store.users().len(), 0);
        assert!(store.get_host(user.name()).is_none());
    }

    #[test]
    fn refuses_user_name_collision() {
        let mut store = KeyStore::new();
        let user = dummy_user();
        assert!(store.insert_user(user.clone()).is_ok());
        assert!(store.insert_user(user).is_err());
        assert_eq!(store.users().len(), 1);
    }

    #[test]
    fn returns_error_if_cant_find_user() {
        let store = KeyStore::new();
        assert!(store.get_user("tomjon").is_none());
    }

    #[test]
    fn finds_user() {
        let mut store = KeyStore::new();
        let user = dummy_user();
        store.insert_user(user.clone()).unwrap();
        assert_eq!(store.get_user(user.name()), Some(&user));
    }

    #[test]
    fn removes_user() {
        let mut store = KeyStore::new();
        let user = dummy_user();
        store.insert_user(user.clone()).unwrap();
        store.make_clean();
        assert!(store.remove_user(user.name()).is_ok());
        assert!(store.is_dirty());
        assert!(store.users().is_empty());
    }

    #[test]
    fn fails_if_removing_unknown_user() {
        let mut store = KeyStore::new();
        let r = store.remove_user("tomjon");
        match r {
            Err(KeyStoreError::UnknownUser(_)) => (),
            _ => panic!("could not remove user from store"),
        }
    }

    #[test]
    fn sets_user_principals() {
        let mut store = KeyStore::new();
        let user = dummy_user();
        store.insert_user(user.clone()).unwrap();
        if let Some(u) = store.get_user_mut(user.name()) {
            u.set_principals(&["alias".into()]);
        } else {
            panic!("failed to get mutable user");
        }

        assert!(store.is_dirty());
        let user2 = store.get_user(user.name()).unwrap();
        assert_eq!(user2.principals(), &["alias"]);
    }

    #[test]
    fn adds_to_user_principals() {
        let mut store = KeyStore::new();
        let user = dummy_user();
        store.insert_user(user.clone()).unwrap();
        if let Some(u) = store.get_user_mut(user.name()) {
            u.add_principals(&["alias".into()]);
        } else {
            panic!("failed to get mutable user");
        }

        assert!(store.is_dirty());
        let user2 = store.get_user(user.name()).unwrap();
        assert_eq!(user2.principals(), &["dummy", "alias"]);
    }

    #[test]
    fn removes_from_user_principals() {
        let mut store = KeyStore::new();
        let user = dummy_user();
        store.insert_user(user.clone()).unwrap();
        if let Some(u) = store.get_user_mut(user.name()) {
            u.add_principals(&["alias".into()]);
            u.remove_principals(&["dummy".into()]);
        } else {
            panic!("failed to get mutable user");
        }

        assert!(store.is_dirty());
        let user2 = store.get_user(user.name()).unwrap();
        assert_eq!(user2.principals(), &["alias"]);
    }
}
