bitbottle 0.9.1

use aes::Aes256Ctr;
use bcrypt_pbkdf::bcrypt_pbkdf;
use cipher::{NewCipher, StreamCipher};
use digest::generic_array::GenericArray;
use dryoc::classic::crypto_sign_ed25519::{crypto_sign_ed25519_pk_to_curve25519, crypto_sign_ed25519_sk_to_curve25519};
use dryoc::constants::CRYPTO_SIGN_PUBLICKEYBYTES;
use dryoc::dryocbox;
use dryoc::sign;
use num_enum::TryFromPrimitive;
use std::fs::File;
use std::fmt;
use std::io::Read;
use crate::bottle_error::{BottleError, BottleResult};
use crate::header::Header;


const EXPECTED_SSH_SK_START: &str = "-----BEGIN OPENSSH PRIVATE KEY-----";
const EXPECTED_SSH_SK_END: &str = "-----END OPENSSH PRIVATE KEY-----";

// helper function for extracting fields from "ssh encoding", which is
// [be32 length + data]. returns (field, remainder).
fn get_ssh_field(data: &[u8]) -> BottleResult<(&[u8], &[u8])> {
    let (len, data) = get_ssh_u32(data)?;
    if data.len() < len as usize { return Err(BottleError::SshEncodingError); }
    Ok(data.split_at(len as usize))
}

fn get_ssh_u32(data: &[u8]) -> BottleResult<(u32, &[u8])> {
    if data.len() < 4 { return Err(BottleError::SshEncodingError); }
    let rv = u32::from_be_bytes(data[0..4].try_into().unwrap());
    Ok((rv, &data[4..]))
}

fn skip_ssh_fields(data: &[u8], count: usize) -> BottleResult<&[u8]> {
    let mut rv = data;
    for _ in 0..count {
        let (_, data) = get_ssh_field(data)?;
        rv = data;
    }
    Ok(rv)
}


const FIELD_RECIPIENT_NAME_STRING: u8 = 0;
const FIELD_PUBLIC_KEY_BYTES: u8 = 0;
const FIELD_ENCRYPTED_KEY_BYTES: u8 = 1;

pub trait BottlePublicKey {
    fn algorithm(&self) -> PublicKeyAlgorithm;
    fn name(&self) -> &str;
    fn as_bytes(&self) -> &[u8];
    fn encrypt(&self, data: &[u8]) -> BottleResult<Vec<u8>>;
    fn box_clone(&self) -> Box<dyn BottlePublicKey>;

    fn to_header(&self, key: &[u8]) -> BottleResult<Header> {
        let mut h = Header::new();
        h.add_string(FIELD_RECIPIENT_NAME_STRING, self.name())?;
        h.add_bytes(FIELD_PUBLIC_KEY_BYTES, self.as_bytes())?;
        h.add_bytes(FIELD_ENCRYPTED_KEY_BYTES, &self.encrypt(key)?)?;
        Ok(h)
    }
}

impl fmt::Debug for dyn BottlePublicKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Ed25519PublicKey({}, {})", self.name(), hex::encode(self.as_bytes()))
    }
}

pub trait BottleSecretKey {
    fn algorithm(&self) -> PublicKeyAlgorithm;
    fn name(&self) -> &str;
    fn matches_public_key(&self, public_key: &dyn BottlePublicKey) -> bool;
    fn decrypt(&self, data: &[u8]) -> BottleResult<Vec<u8>>;
}


#[allow(non_camel_case_types)]
#[derive(Clone, Copy, Debug, PartialEq, TryFromPrimitive)]
#[repr(u8)]
pub enum PublicKeyAlgorithm {
    ED25519_NACL_SEALED = 0,
}

fn public_key_from_bytes(
    algorithm: PublicKeyAlgorithm,
    data: &[u8],
    name: String
) -> Option<Box<dyn BottlePublicKey>> {
    match algorithm {
        PublicKeyAlgorithm::ED25519_NACL_SEALED => Some(Box::new(Ed25519PublicKey::from_bytes(data, name))),
        // _ => None
    }
}


#[derive(Debug)]
pub struct EncryptedKeyHeader {
    pub public_key: Box<dyn BottlePublicKey>,
    pub encrypted_key: Vec<u8>,
    pub header: Header,
}

/// Parse out the public key from a Header which is used to encode a key
/// for an encrypted bottle, and if it's valid, return the BottlePublicKey
/// representing it, as well as the encrypted key it can decrypt.
pub(crate) fn public_key_from_header(
    algorithm: PublicKeyAlgorithm,
    header: Header
) -> Option<EncryptedKeyHeader> {
    let encrypted_key = header.get_bytes(FIELD_ENCRYPTED_KEY_BYTES).unwrap_or(&[]).to_vec();
    let key_bytes = header.get_bytes(FIELD_PUBLIC_KEY_BYTES).unwrap_or(&[]);
    let name = header.get_string(FIELD_RECIPIENT_NAME_STRING).unwrap_or("");
    if encrypted_key.is_empty() || key_bytes.is_empty() || name.is_empty() { return None; }
    public_key_from_bytes(algorithm, key_bytes, String::from(name)).map(|public_key| {
        EncryptedKeyHeader { public_key, encrypted_key, header }
    })
}

impl Clone for EncryptedKeyHeader {
    fn clone(&self) -> Self {
        EncryptedKeyHeader {
            public_key: self.public_key.box_clone(),
            encrypted_key: self.encrypted_key.clone(),
            header: self.header.clone(),
        }
    }
}


#[derive(Clone)]
pub struct Ed25519PublicKey {
    pk: dryocbox::PublicKey,
    name: String,
}

impl Ed25519PublicKey {
    pub fn from_bytes(data: &[u8], name: String) -> Ed25519PublicKey {
        let mut pk = dryocbox::PublicKey::new();
        pk.copy_from_slice(data);
        Ed25519PublicKey { pk, name }
    }

    /// Convert an SSH key into the "curve" form used by nacl (dryoc).
    pub fn from_ssh_bytes(data: &[u8], name: String) -> BottleResult<Ed25519PublicKey> {
        if data.len() != CRYPTO_SIGN_PUBLICKEYBYTES { return Err(BottleError::NotAnEd25519Key); }
        let vpk = sign::PublicKey::from(&data.try_into().unwrap());
        let mut pk = dryocbox::PublicKey::new();
        crypto_sign_ed25519_pk_to_curve25519(pk.as_mut(), vpk.as_ref()).map_err(|_| BottleError::NotAnEd25519Key)?;
        Ok(Ed25519PublicKey { pk, name })
    }

    pub fn from_ssh_contents(filename: &str, file_contents: String) -> BottleResult<Ed25519PublicKey> {
        let first_line = file_contents.split('\n').next().ok_or_else(|| {
            BottleError::InvalidSshFile(filename.to_string())
        })?;
        let segments: Vec<&str> = first_line.split(' ').collect();
        if segments.len() < 3 || segments[0] != "ssh-ed25519" {
            return Err(BottleError::InvalidSshFile(filename.to_string()));
        }
        let name = segments[2];
        let data = base64::decode(segments[1]).map_err(|_| BottleError::InvalidSshFile(filename.to_string()))?;
        if data.len() < 8 { return Err(BottleError::InvalidSshFile(filename.to_string())); }

        // decode the base64: algorithm-name (again), and key data
        let (alg_name, data) = get_ssh_field(&data)?;
        if alg_name != b"ssh-ed25519" { return Err(BottleError::InvalidSshFile(filename.to_string())); }
        let (key_data, _) = get_ssh_field(data)?;
        if key_data.len() != CRYPTO_SIGN_PUBLICKEYBYTES {
            return Err(BottleError::InvalidSshFile(filename.to_string()));
        }
        Ed25519PublicKey::from_ssh_bytes(key_data, String::from(name))
    }

    pub fn from_ssh_file(filename: &str) -> BottleResult<Ed25519PublicKey> {
        let mut contents = String::new();
        let mut file = File::open(filename).map_err(|e| BottleError::FileError(filename.to_string(), e.kind()))?;
        file.read_to_string(&mut contents).map_err(|e| BottleError::FileError(filename.to_string(), e.kind()))?;
        Ed25519PublicKey::from_ssh_contents(filename, contents)
    }
}

impl BottlePublicKey for Ed25519PublicKey {
    fn algorithm(&self) -> PublicKeyAlgorithm {
        PublicKeyAlgorithm::ED25519_NACL_SEALED
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn as_bytes(&self) -> &[u8] {
        self.pk.as_ref()
    }

    fn encrypt(&self, data: &[u8]) -> BottleResult<Vec<u8>> {
        let sealed = dryocbox::DryocBox::seal_to_vecbox(data, &self.pk).map_err(|_| BottleError::DryocBoxError)?;
        Ok(sealed.to_vec())
    }

    fn box_clone(&self) -> Box<dyn BottlePublicKey> {
        Box::new(self.clone())
    }
}


pub struct Ed25519SecretKey {
    key_pair: dryocbox::KeyPair,
    pub name: String,
}

impl Ed25519SecretKey {
    /// Convert an ssh secret key (which is actually the seed) into a nacl
    /// (dryoc) key pair.
    pub fn from_ssh_bytes(data: &[u8], name: String) -> BottleResult<Ed25519SecretKey> {
        let key_pair: sign::SigningKeyPair<sign::PublicKey, sign::SecretKey> = sign::SigningKeyPair::from_seed(&data);
        let mut sk = dryocbox::SecretKey::new();
        crypto_sign_ed25519_sk_to_curve25519(sk.as_mut(), key_pair.secret_key.as_ref());
        let mut pk = dryocbox::PublicKey::new();
        crypto_sign_ed25519_pk_to_curve25519(pk.as_mut(), key_pair.public_key.as_ref()).map_err(|_| {
            BottleError::NotAnEd25519Key
        })?;
        let key_pair = dryocbox::KeyPair::from_slices(&pk, &sk).map_err(|_| BottleError::NotAnEd25519Key)?;
        Ok(Ed25519SecretKey { key_pair, name })
    }

    /// There's no formal spec for an SSH private key file, but the internet
    /// has worked it out by scanning the OpenSSH source code (and apparently
    /// some trial and error). Between the "start" and "end" lines is base-64
    /// data in a custom format. This format encodes numbers as 32-bit
    /// big-endian ("u32be") and strings as UTF-8.
    ///
    /// I've tried to document this better in [docs/openssh-keys.md](docs/openssh-keys.md).
    pub fn from_ssh_contents(
        filename: &str,
        file_contents: String,
        password: Option<&str>,
    ) -> BottleResult<Ed25519SecretKey> {
        // decode base64 between marker lines
        let mut lines: Vec<&str> = file_contents.split('\n').collect();
        while !lines.is_empty() && lines[0] != EXPECTED_SSH_SK_START {
            lines.remove(0);
        }
        lines.remove(0);
        while !lines.is_empty() && lines.last() != Some(&EXPECTED_SSH_SK_END) {
            lines.pop();
        }
        lines.pop();
        let data = base64::decode(lines.join("")).map_err(|_| BottleError::InvalidSshFile(filename.to_string()))?;

        // skip identifier string
        let i = data.iter().position(|b| *b == 0).unwrap_or(0);
        if &data[0..i] != b"openssh-key-v1" { return Err(BottleError::InvalidSshFile(filename.to_string())); }

        // is it encrypted? we only support "none" and "bcrypt/aes256-ctr"
        let (cipher_name, data) = get_ssh_field(&data[i + 1 ..])?;
        let (kdf_name, data) = get_ssh_field(data)?;
        let (kdf_options, data) = get_ssh_field(data)?;
        let is_encrypted = match (&kdf_name, &cipher_name) {
            (&b"none", &b"none") => false,
            (&b"bcrypt", &b"aes256-ctr") => true,
            _ => {
                // make a decent effort to tell them what's wrong with the key file.
                let kind = format!(
                    "{}/{}",
                    String::from_utf8(kdf_name.to_vec()).unwrap_or_else(|_| "?".to_string()),
                    String::from_utf8(cipher_name.to_vec()).unwrap_or_else(|_| "?".to_string()),
                );
                return Err(BottleError::UnsupportedSshFileEncryption { filename: filename.to_string(), kind });
            }
        };

        // public key ("keys") for this secret key (redundant, since we'll see it again in a minute)
        let (public_key_count, data) = get_ssh_u32(data)?;
        let data = skip_ssh_fields(data, public_key_count as usize)?;
        let (data, _) = get_ssh_field(data)?;

        // ok, final lap now. this section may be encrypted.
        let mut data = data.to_owned();
        if is_encrypted {
            if password.is_none() {
                return Err(BottleError::SshPasswordRequired);
            }
            let (salt, kdf_options) = get_ssh_field(kdf_options)?;
            let (rounds, _) = get_ssh_u32(kdf_options)?;

            // generate AES-256 key (32 bytes) + counter IV (16 bytes)
            let mut generated = [0u8; 48];
            bcrypt_pbkdf(password.unwrap(), salt, rounds, &mut generated).unwrap();
            let (key, iv) = generated.split_at(32);

            let mut cipher = Aes256Ctr::new(GenericArray::from_slice(key), GenericArray::from_slice(iv));
            cipher.apply_keystream(&mut data);
        }

        // skip 8-byte salt.
        let (alg_name, data) = get_ssh_field(&data[8..])?;
        if alg_name != b"ssh-ed25519" { return Err(BottleError::InvalidSshFile(filename.to_string())); }
        let (_public_key, data) = get_ssh_field(data)?;
        let (secret_key_data, data) = get_ssh_field(data)?;
        let (secret_key, _public_key) = secret_key_data.split_at(CRYPTO_SIGN_PUBLICKEYBYTES);
        let (name, _data) = get_ssh_field(data)?;
        let name = String::from_utf8(name.to_vec()).map_err(|_| BottleError::InvalidSshFile(filename.to_string()))?;

        Ed25519SecretKey::from_ssh_bytes(secret_key, name)
    }

    pub fn from_ssh_file(filename: &str, password: Option<&str>) -> BottleResult<Ed25519SecretKey> {
        let mut contents = String::new();
        let mut file = File::open(filename).map_err(|e| BottleError::FileError(filename.to_string(), e.kind()))?;
        file.read_to_string(&mut contents).map_err(|e| BottleError::FileError(filename.to_string(), e.kind()))?;
        Ed25519SecretKey::from_ssh_contents(filename, contents, password)
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.key_pair.secret_key.as_ref()
    }
}

impl fmt::Debug for Ed25519SecretKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Ed25519SecretKey({}, {})", self.name, hex::encode(&self.key_pair.secret_key))
    }
}

impl BottleSecretKey for Ed25519SecretKey {
    fn algorithm(&self) -> PublicKeyAlgorithm {
        PublicKeyAlgorithm::ED25519_NACL_SEALED
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn matches_public_key(&self, public_key: &dyn BottlePublicKey) -> bool {
        public_key.as_bytes() == &self.key_pair.public_key[..]
    }

    fn decrypt(&self, data: &[u8]) -> BottleResult<Vec<u8>> {
        let sealed = dryocbox::DryocBox::from_sealed_bytes(data).map_err(|_| BottleError::DryocBoxError)?;
        sealed.unseal_to_vec(&self.key_pair).map_err(|_| BottleError::DryocBoxError)
    }
}


#[cfg(test)]
mod test {
    use bcrypt_pbkdf::bcrypt_pbkdf;
    use hex;
    use std::fs;
    use super::{BottlePublicKey, BottleSecretKey, Ed25519PublicKey, Ed25519SecretKey};

    const PK_HEX: &str = "112905df7c79c726b6250ec6e87aaa35d9127e2f48736a65c7352e8768d77865";
    const SK_HEX: &str = "bc29877fa28b3e03ea6b58d35818e65a1fb4870dae35c4a452de880205108efb";

    const PK_NACL_HEX: &str = "34fd22aae3c59072fd6f48147309eb302ea30f6ae5fc6376f683df3e74485a7c";
    const SK_NACL_HEX: &str = "c026f16ec964ffc5314fc295d983828cf00f65a53cb9aee1375e7bee2c547e61";

    const MESSAGE: &str = "One person's rando is another person's homey.";

    #[test]
    fn from_ssh_bytes() {
        let pk = Ed25519PublicKey::from_ssh_bytes(&hex::decode(PK_HEX).unwrap(), String::new()).unwrap();
        assert_eq!(hex::encode(pk.as_bytes()), PK_NACL_HEX);
        let sk = Ed25519SecretKey::from_ssh_bytes(&hex::decode(SK_HEX).unwrap(), String::new()).unwrap();
        assert_eq!(hex::encode(sk.as_bytes()), SK_NACL_HEX);
    }

    #[test]
    fn encrypt() {
        let pk = Ed25519PublicKey::from_ssh_bytes(&hex::decode(PK_HEX).unwrap(), String::new()).unwrap();
        let sk = Ed25519SecretKey::from_ssh_bytes(&hex::decode(SK_HEX).unwrap(), String::new()).unwrap();
        let encrypted = pk.encrypt(MESSAGE.as_bytes()).unwrap();
        let decrypted = sk.decrypt(&encrypted).unwrap();
        assert_eq!(String::from_utf8(decrypted).unwrap(), MESSAGE);
    }

    #[test]
    fn read_ssh_pk() {
        let data = fs::read_to_string("./tests/data/test-key.pub").unwrap();
        let pk = Ed25519PublicKey::from_ssh_contents("test-key.pub", data).unwrap();
        assert_eq!(hex::encode(pk.as_bytes()), PK_NACL_HEX);
        assert_eq!(
            format!("{:?}", &pk as &dyn BottlePublicKey),
            "Ed25519PublicKey(robey@togusa, 34fd22aae3c59072fd6f48147309eb302ea30f6ae5fc6376f683df3e74485a7c)"
        );
    }

    #[test]
    fn read_ssh_sk() {
        let data = fs::read_to_string("./tests/data/test-key").unwrap();
        let sk = Ed25519SecretKey::from_ssh_contents("test-key", data, None).unwrap();
        assert_eq!(hex::encode(sk.as_bytes()), SK_NACL_HEX);
        assert_eq!(
            format!("{:?}", &sk),
            "Ed25519SecretKey(robey@togusa, c026f16ec964ffc5314fc295d983828cf00f65a53cb9aee1375e7bee2c547e61)"
        );
    }

    #[test]
    fn read_ssh_sk_encrypted() {
        // first, verify that bcrypt_pbkdf works:
        let mut generated = vec![0u8; 48];
        let salt = hex::decode("240d8b1cd1367fd79758b9569b991d8f").unwrap();
        bcrypt_pbkdf("password", &salt, 16, &mut generated).unwrap();
        assert_eq!(
            hex::encode(&generated), "016129df9f55f9a962d0f24bd09c2c9df2081e47663b1835090cfbab32a42b4b2651e759166186e66cd422cd381b5e58"
        );

        let data = fs::read_to_string("./tests/data/test-key-pw").unwrap();
        let bad_sk = Ed25519SecretKey::from_ssh_contents("test-key-pw", data.clone(), None);
        assert!(bad_sk.is_err());
        assert_eq!(format!("{:?}", bad_sk.unwrap_err()), "SshPasswordRequired");

        let sk = Ed25519SecretKey::from_ssh_contents("test-key-pw", data, Some("password")).unwrap();
        assert_eq!(
            format!("{:?}", &sk),
            "Ed25519SecretKey(robey@togusa, e8e78aa6f852348a9ba107beba008979011efce06115376bd7dc72c086acdf67)"
        );
    }

    #[test]
    fn encrypt_from_ssh() {
        let pk_data = fs::read_to_string("./tests/data/test-key.pub").unwrap();
        let pk = Ed25519PublicKey::from_ssh_contents("test-key.pub", pk_data).unwrap();
        let sk_data = fs::read_to_string("./tests/data/test-key").unwrap();
        let sk = Ed25519SecretKey::from_ssh_contents("test-key", sk_data, None).unwrap();

        let encrypted = pk.encrypt(MESSAGE.as_bytes()).unwrap();
        let decrypted = sk.decrypt(&encrypted).unwrap();
        assert_eq!(String::from_utf8(decrypted).unwrap(), MESSAGE);
    }
}