snow 0.0.1-preview.2

use constants::*;
use crypto_types::*;
use handshakestate::*;
use wrappers::rand_wrapper::*;
use wrappers::crypto_wrapper::*;
use cipherstate::*;
use session::*;
use utils::*;
use params::*;
use error::NoiseError;

pub trait CryptoResolver {
    fn resolve_rng(&self) -> Option<Box<Random>>;
    fn resolve_dh(&self, choice: &DHChoice) -> Option<Box<Dh>>;
    fn resolve_hash(&self, choice: &HashChoice) -> Option<Box<Hash>>;
    fn resolve_cipher(&self, choice: &CipherChoice) -> Option<Box<CipherStateType>>;
}

pub struct DefaultResolver;
impl CryptoResolver for DefaultResolver {
    fn resolve_rng(&self) -> Option<Box<Random>> {
        Some(Box::new(RandomOs::default()))
    }

    fn resolve_dh(&self, choice: &DHChoice) -> Option<Box<Dh>> {
        match *choice {
            DHChoice::Curve25519 => Some(Box::new(Dh25519::default())),
            _                    => None,
        }
    }

    fn resolve_hash(&self, choice: &HashChoice) -> Option<Box<Hash>> {
        match *choice {
            HashChoice::SHA256  => Some(Box::new(HashSHA256::default())),
            HashChoice::SHA512  => Some(Box::new(HashSHA512::default())),
            HashChoice::Blake2s => Some(Box::new(HashBLAKE2s::default())),
            HashChoice::Blake2b => Some(Box::new(HashBLAKE2b::default())),
        }
    }

    fn resolve_cipher(&self, choice: &CipherChoice) -> Option<Box<CipherStateType>> {
        match *choice {
            CipherChoice::ChaChaPoly => Some(Box::new(CipherState::<CipherChaChaPoly>::default())),
            CipherChoice::AESGCM     => Some(Box::new(CipherState::<CipherAESGCM>::default())),
        }
    }
}

/// Generates a `NoiseSession` and also validate that all the prerequisites for
/// the given parameters are satisfied.
///
/// # Examples
///
/// ```
/// # use snow::NoiseBuilder;
/// # let my_long_term_key = [0u8; 32];
/// # let their_pub_key = [0u8; 32];
/// let noise = NoiseBuilder::new("Noise_XX_25519_ChaChaPoly_BLAKE2s".parse().unwrap())
///                          .local_private_key(&my_long_term_key)
///                          .remote_public_key(&their_pub_key)
///                          .prologue("noise is just swell".as_bytes())
///                          .build_initiator()
///                          .unwrap();
/// ```
pub struct NoiseBuilder<'a> {
    params: NoiseParams,
    resolver: Box<CryptoResolver>, // The mapper from protocol choices to crypto implementations
    s: Option<&'a [u8]>,
    e_fixed: Option<&'a [u8]>,
    rs: Option<Vec<u8>>,
    re: Option<Vec<u8>>,
    psk: Option<Vec<u8>>,
    plog: Option<Vec<u8>>,
}

impl<'a> NoiseBuilder<'a> {
    /// Create a NoiseBuilder with the default crypto resolver.
    pub fn new(params: NoiseParams) -> Self {
        Self::with_resolver(params, Box::new(DefaultResolver{}))
    }

    /// Create a NoiseBuilder with a custom crypto resolver.
    pub fn with_resolver(params: NoiseParams, resolver: Box<CryptoResolver>) -> Self
    {
        NoiseBuilder {
            params: params,
            resolver: resolver,
            s: None,
            e_fixed: None,
            rs: None,
            re: None,
            plog: None,
            psk: None,
        }
    }

    /// Specify a PSK (only used with `NoisePSK` base parameter)
    pub fn preshared_key(mut self, key: &[u8]) -> Self {
        self.psk = Some(key.to_vec());
        self
    }

    /// Your static private key, can be generated by [`generate_private_key()`](#method.generate_private_key).
    pub fn local_private_key(mut self, key: &'a [u8]) -> Self {
        self.s = Some(key);
        self
    }

    #[doc(hidden)]
    pub fn fixed_ephemeral_key_for_testing_only(mut self, key: &'a [u8]) -> Self {
        self.e_fixed = Some(key);
        self
    }

    /// Arbitrary data to be hashed in to the handshake hash value.
    pub fn prologue(mut self, key: &[u8]) -> Self {
        self.plog = Some(key.to_vec());
        self
    }

    /// The responder's static public key.
    pub fn remote_public_key(mut self, pub_key: &[u8]) -> Self {
        self.rs = Some(pub_key.to_vec());
        self
    }

    /// Generate a new private key. It's up to the user of this library how to store this.
    pub fn generate_private_key(&self) -> Result<Vec<u8>, NoiseError> {
        let mut rng = self.resolver.resolve_rng()
            .ok_or(NoiseError::InitError("no suitable RNG"))?;
        let mut dh = self.resolver.resolve_dh(&self.params.dh)
            .ok_or(NoiseError::InitError("no suitable DH implementation"))?;
        let mut private = vec![0u8; dh.priv_len()];
        dh.generate(&mut *rng);
        private[..dh.priv_len()].copy_from_slice(dh.privkey());
        Ok(private)
    }

    /// Build a NoiseSession for the side who will initiate the handshake (send the first message)
    pub fn build_initiator(self) -> Result<NoiseSession<HandshakeState>, NoiseError> {
        self.build(true)
    }

    /// Build a NoiseSession for the side who will be responder (receive the first message)
    pub fn build_responder(self) -> Result<NoiseSession<HandshakeState>, NoiseError> {
        self.build(false)
    }

    fn build(self, initiator: bool) -> Result<NoiseSession<HandshakeState>, NoiseError> {
        if !self.s.is_some() && self.params.handshake.needs_local_static_key(initiator) {
            return Err(NoiseError::InitError("local key needed for chosen handshake pattern"));
        }

        if !self.rs.is_some() && self.params.handshake.need_known_remote_pubkey(initiator) {
            return Err(NoiseError::InitError("remote key needed for chosen handshake pattern"));
        }

        let rng = self.resolver.resolve_rng().ok_or(NoiseError::InitError("no suitable RNG"))?;
        let cipher = self.resolver.resolve_cipher(&self.params.cipher).ok_or(NoiseError::InitError("no cipher implementation"))?;
        let hash = self.resolver.resolve_hash(&self.params.hash).ok_or(NoiseError::InitError("no hash implementation"))?;
        let mut s_dh = self.resolver.resolve_dh(&self.params.dh).ok_or(NoiseError::InitError("no DH implementation"))?;
        let mut e_dh = self.resolver.resolve_dh(&self.params.dh).ok_or(NoiseError::InitError("no DH implementation"))?;
        let cipherstate1 = self.resolver.resolve_cipher(&self.params.cipher).ok_or(NoiseError::InitError("no cipher implementation"))?;
        let cipherstate2 = self.resolver.resolve_cipher(&self.params.cipher).ok_or(NoiseError::InitError("no cipher implementation"))?;

        let s = match self.s {
            Some(k) => {
                (&mut *s_dh).set(k);
                Toggle::on(s_dh)
            },
            None => {
                Toggle::off(s_dh)
            }
        };

        if let Some(fixed_k) = self.e_fixed {
            (&mut *e_dh).set(fixed_k);
        }
        let e = Toggle::off(e_dh);

        let e_fixed = match self.e_fixed {
            Some(k) => {
                let mut e_fixed_dh = self.resolver.resolve_dh(&self.params.dh).ok_or(NoiseError::InitError("no DH implementation"))?;
                (&mut *e_fixed_dh).set(k);
                Some(e_fixed_dh)
            },
            None => None
        };

        let mut rs_buf = [0u8; MAXDHLEN];
        let rs = match self.rs {
            Some(v) => {
                rs_buf[..v.len()].copy_from_slice(&v[..]);
                Toggle::on(rs_buf)
            },
            None => Toggle::off(rs_buf),
        };

        let mut re_buf = [0u8; MAXDHLEN];
        let re = match self.re {
            Some(v) => {
                re_buf[..v.len()].copy_from_slice(&v[..]);
                Toggle::on(re_buf)
            },
            None => Toggle::off(re_buf),
        };

        let hs = HandshakeState::new(rng, cipher, hash,
                                     s, e, self.e_fixed.is_some(), rs, re,
                                     initiator,
                                     self.params.handshake,
                                     &self.plog.unwrap_or_else(|| vec![0u8; 0]),
                                     self.psk,
                                     CipherStates::new(cipherstate1, cipherstate2)?)?;
        Ok(hs.into())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_builder() {
        let _noise = NoiseBuilder::new("Noise_NN_25519_ChaChaPoly_SHA256".parse().unwrap())
            .preshared_key(&[1,1,1,1,1,1,1])
            .prologue(&[2,2,2,2,2,2,2,2])
            .local_private_key(&[0u8; 32])
            .build_initiator().unwrap();
    }

    #[test]
    fn test_builder_keygen() {
        let builder = NoiseBuilder::new("Noise_NN_25519_ChaChaPoly_SHA256".parse().unwrap());
        let key1 = builder.generate_private_key();
        let key2 = builder.generate_private_key();
        assert!(key1.unwrap() != key2.unwrap());
    }

    #[test]
    fn test_builder_bad_spec() {
        let params: Result<NoiseParams, _> = "Noise_NK_25519_ChaChaPoly_BLAH256".parse();

        if let Ok(_) = params {
            panic!("NoiseParams should have failed");
        }
    }

    #[test]
    fn test_builder_missing_prereqs() {
        let noise = NoiseBuilder::new("Noise_NK_25519_ChaChaPoly_SHA256".parse().unwrap())
            .preshared_key(&[1,1,1,1,1,1,1])
            .prologue(&[2,2,2,2,2,2,2,2])
            .local_private_key(&[0u8; 32])
            .build_initiator(); // missing remote key, should result in Err

        if let Ok(_) = noise {
            panic!("builder should have failed on build");
        }
    }
}