#[macro_use]
extern crate uint;
pub mod chacha20;
pub mod csprng;
pub(crate) mod kdf;
pub mod poly1305;
pub mod salsa20;
construct_uint! {
    pub struct U256(4);
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum CipherType {
    Salsa20,
    Chacha20,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct SecretBox {
    key: [u8; 32],
    cipher: CipherType,
}

impl SecretBox {
    /// Creates a new SecretBox instance
    ///
    /// Returns None if the passed key is shorter than 32 bytes
    pub fn new<R>(key: R, cipher: CipherType) -> Option<Self>
    where
        R: AsRef<[u8]>,
    {
        let k = key.as_ref();
        if k.len() < 32 {
            return None;
        }
        Some(Self {
            key: [
                k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7], k[8], k[9], k[10], k[11], k[12],
                k[13], k[14], k[15], k[16], k[17], k[18], k[19], k[20], k[21], k[22], k[23], k[24],
                k[25], k[26], k[27], k[28], k[29], k[30], k[31],
            ],
            cipher,
        })
    }
    /// Creates a new SecretBox instance by doing an ECDH key exchange using curve25519
    ///
    /// Returns None if the passed public key is not 32 bytes long.
    #[cfg(feature = "curve25519")]
    pub fn from_ecdh<T, R>(
        peer_public_key: T,
        rng: &mut R,
        cipher: CiperType,
    ) -> Option<(Self, [u8; 32])>
    where
        T: AsRef<[u8]>,
        R: rand::Rng + rand::CryptoRng,
    {
        let k = peer_public_key.as_ref();
        if k.len() != 32 {
            return None;
        }
        let peer_pubkey: x25519_dalek::PublicKey = From::from([
            k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7], k[8], k[9], k[10], k[11], k[12], k[13],
            k[14], k[15], k[16], k[17], k[18], k[19], k[20], k[21], k[22], k[23], k[24], k[25],
            k[26], k[27], k[28], k[29], k[30], k[31],
        ]);
        let privkey = x25519_dalek::EphermalSecret::new(rng);
        let pubkey: x25519_dalek::PublicĶey = From::from(&privkey);
        let shared_secret = privkey.diffie_hellman(&peer_pubkey);
        Some((
            Self {
                key: *shared_secret.as_bytes(),
                cipher,
            },
            *pubkey.as_bytes(),
        ))
    }
    /// Creates a new SecretBox instance with a generated key
    pub fn from_random_key<R>(rng: &mut R, cipher: CipherType) -> (Self, [u8; 32])
    where
        R: rand::Rng + rand::CryptoRng,
    {
        let mut buf = [0u8; 32];
        rng.fill_bytes(&mut buf);
        (Self { key: buf, cipher }, buf)
    }
    /// This function returns an encrypted and authenticated copy of the message. The key and nonce
    /// pair must be unique for every message.
    pub fn seal(&self, message: &[u8], nonce: [u8; 24]) -> Vec<u8> {
        let mut out = Vec::with_capacity(message.len() + 16);
        out.extend_from_slice(&[0u8; 16]);
        let (sub_key, nonce) = kdf::generate_subkey(nonce, self.key);
        // The poly1305 key is generated by encrypting 32 bytes of zeroes.
        // The chosen encryption method works with 64-byte blocks
        let initial_block = match self.cipher {
            CipherType::Chacha20 => chacha20::XChacha20::new(sub_key, nonce, 0).generate_block(),
            CipherType::Salsa20 => salsa20::XSalsa20::new(sub_key, nonce, 0).generate_block(),
        }
        .unwrap();
        let poly1305_r = u128::from_le_bytes([
            initial_block[0],
            initial_block[1],
            initial_block[2],
            initial_block[3],
            initial_block[4],
            initial_block[5],
            initial_block[6],
            initial_block[7],
            initial_block[8],
            initial_block[9],
            initial_block[10],
            initial_block[11],
            initial_block[12],
            initial_block[13],
            initial_block[14],
            initial_block[15],
        ]);
        let poly1305_s = u128::from_le_bytes([
            initial_block[16],
            initial_block[17],
            initial_block[18],
            initial_block[19],
            initial_block[20],
            initial_block[21],
            initial_block[22],
            initial_block[23],
            initial_block[24],
            initial_block[25],
            initial_block[26],
            initial_block[27],
            initial_block[28],
            initial_block[29],
            initial_block[30],
            initial_block[31],
        ]);
        let mut mac = poly1305::Poly1305::new(poly1305_r, poly1305_s);
        // Encrypt 32 bytes
        for i in 0..std::cmp::min(32, message.len()) {
            out.push(initial_block[32 + i] ^ message[i]);
        }
        if message.len() > 32 {
            out.extend_from_slice(&message[32..]);
            match self.cipher {
                CipherType::Chacha20 => {
                    chacha20::XChacha20::new(sub_key, nonce, 1).crypt(&mut out[48..])
                }
                CipherType::Salsa20 => {
                    salsa20::XSalsa20::new(sub_key, nonce, 1).crypt(&mut out[48..])
                }
            };
        }
        let mac = mac.hash(&out[16..]);
        for (i, c) in mac.to_le_bytes().into_iter().enumerate() {
            out[i] = *c;
        }
        out
    }

    /// This function works like the above, except that it automatically generates a unique nonce.
    pub fn easy_seal(&self, message: &[u8]) -> Vec<u8> {
        use rand::RngCore;
        let mut rng = rand::rngs::OsRng;
        let mut nonce = [0u8; 24];
        rng.fill_bytes(&mut nonce);
        let mut v = Vec::with_capacity(message.len() + 16 + 24);
        v.extend_from_slice(&nonce);
        v.extend_from_slice(&self.seal(message, nonce));
        v
    }

    /// This function works like unseal, except that it finds the nonce automatically.
    pub fn easy_unseal(&self, data: &[u8]) -> Option<Vec<u8>> {
        let nonce = [
            data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7], data[8],
            data[9], data[10], data[11], data[12], data[13], data[14], data[15], data[16],
            data[17], data[18], data[19], data[20], data[21], data[22], data[23],
        ];
        self.unseal(&data[24..], nonce)
    }

    /// This function tries to authenticate and decrypt a box
    pub fn unseal(&self, data: &[u8], nonce: [u8; 24]) -> Option<Vec<u8>> {
        let mut out = Vec::with_capacity(data.len() - 16);
        let (sub_key, nonce) = kdf::generate_subkey(nonce, self.key);
        // The poly1305 key is generated by encrypting 32 bytes of zeroes.
        // The chosen encryption method works with 64-byte blocks
        let initial_block = match self.cipher {
            CipherType::Chacha20 => chacha20::XChacha20::new(sub_key, nonce, 0).generate_block(),
            CipherType::Salsa20 => salsa20::XSalsa20::new(sub_key, nonce, 0).generate_block(),
        }
        .unwrap();
        let poly1305_r = u128::from_le_bytes([
            initial_block[0],
            initial_block[1],
            initial_block[2],
            initial_block[3],
            initial_block[4],
            initial_block[5],
            initial_block[6],
            initial_block[7],
            initial_block[8],
            initial_block[9],
            initial_block[10],
            initial_block[11],
            initial_block[12],
            initial_block[13],
            initial_block[14],
            initial_block[15],
        ]);
        let poly1305_s = u128::from_le_bytes([
            initial_block[16],
            initial_block[17],
            initial_block[18],
            initial_block[19],
            initial_block[20],
            initial_block[21],
            initial_block[22],
            initial_block[23],
            initial_block[24],
            initial_block[25],
            initial_block[26],
            initial_block[27],
            initial_block[28],
            initial_block[29],
            initial_block[30],
            initial_block[31],
        ]);
        let expected_mac = u128::from_le_bytes([
            data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7], data[8],
            data[9], data[10], data[11], data[12], data[13], data[14], data[15],
        ]);
        let mut mac = poly1305::Poly1305::new(poly1305_r, poly1305_s);
        if !mac.verify(&data[16..], expected_mac) {
            return None;
        }
        // Encrypt 32 bytes
        for i in 0..std::cmp::min(32, data.len() - 16) {
            out.push(initial_block[32 + i] ^ data[16 + i]);
        }
        if data.len() > 32 {
            out.extend_from_slice(&data[48..]);
            match self.cipher {
                CipherType::Chacha20 => {
                    chacha20::XChacha20::new(sub_key, nonce, 1).crypt(&mut out[32..])
                }
                CipherType::Salsa20 => {
                    salsa20::XSalsa20::new(sub_key, nonce, 1).crypt(&mut out[32..])
                }
            };
        }
        Some(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn pynacl_vector() {
        let key = b"\x1b\x27\x55\x64\x73\xe9\x85\xd4\x62\xcd\x51\x19\x7a\x9a\x46\xc7\x60\x09\x54\x9e\xac\x64\x74\xf2\x06\xc4\xee\x08\x44\xf6\x83\x89".clone();
        let nonce = b"\x69\x69\x6e\xe9\x55\xb6\x2b\x73\xcd\x62\xbd\xa8\x75\xfc\x73\xd6\x82\x19\xe0\x03\x6b\x7a\x0b\x37".clone();
        let plain = b"\xbe\x07\x5f\xc5\x3c\x81\xf2\xd5\xcf\x14\x13\x16\xeb\xeb\x0c\x7b\x52\x28\xc5\x2a\x4c\x62\xcb\xd4\x4b\x66\x84\x9b\x64\x24\x4f\xfc\xe5\xec\xba\xaf\x33\xbd\x75\x1a\x1a\xc7\x28\xd4\x5e\x6c\x61\x29\x6c\xdc\x3c\x01\x23\x35\x61\xf4\x1d\xb6\x6c\xce\x31\x4a\xdb\x31\x0e\x3b\xe8\x25\x0c\x46\xf0\x6d\xce\xea\x3a\x7f\xa1\x34\x80\x57\xe2\xf6\x55\x6a\xd6\xb1\x31\x8a\x02\x4a\x83\x8f\x21\xaf\x1f\xde\x04\x89\x77\xeb\x48\xf5\x9f\xfd\x49\x24\xca\x1c\x60\x90\x2e\x52\xf0\xa0\x89\xbc\x76\x89\x70\x40\xe0\x82\xf9\x37\x76\x38\x48\x64\x5e\x07\x05".clone();
        let ciphertext = b"\xf3\xff\xc7\x70\x3f\x94\x00\xe5\x2a\x7d\xfb\x4b\x3d\x33\x05\xd9\x8e\x99\x3b\x9f\x48\x68\x12\x73\xc2\x96\x50\xba\x32\xfc\x76\xce\x48\x33\x2e\xa7\x16\x4d\x96\xa4\x47\x6f\xb8\xc5\x31\xa1\x18\x6a\xc0\xdf\xc1\x7c\x98\xdc\xe8\x7b\x4d\xa7\xf0\x11\xec\x48\xc9\x72\x71\xd2\xc2\x0f\x9b\x92\x8f\xe2\x27\x0d\x6f\xb8\x63\xd5\x17\x38\xb4\x8e\xee\xe3\x14\xa7\xcc\x8a\xb9\x32\x16\x45\x48\xe5\x26\xae\x90\x22\x43\x68\x51\x7a\xcf\xea\xbd\x6b\xb3\x73\x2b\xc0\xe9\xda\x99\x83\x2b\x61\xca\x01\xb6\xde\x56\x24\x4a\x9e\x88\xd5\xf9\xb3\x79\x73\xf6\x22\xa4\x3d\x14\xa6\x59\x9b\x1f\x65\x4c\xb4\x5a\x74\xe3\x55\xa5".clone();
        let s = SecretBox::new(key, CipherType::Salsa20).unwrap();
        let output = s.seal(&plain[..], nonce);
        println!("output: {}, ciphertext: {}", output.len(), ciphertext.len());
        assert_eq!(&output[..], &ciphertext[..]);
        let output2 = s.unseal(&output[..], nonce).unwrap();
        assert_eq!(&output2[..], &plain[..]);
    }

    #[test]
    fn easy_seal_unseal() {
        use rand::RngCore;
        let mut rng = rand::rngs::OsRng;
        let mut key = [0u8; 32];
        rng.fill_bytes(&mut key);
        let plain = b"hello world".to_vec();
        let s = SecretBox::new(key, CipherType::Salsa20).unwrap();
        let sealed = s.easy_seal(&plain);
        let unsealed = s.easy_unseal(&sealed).unwrap();
        assert_eq!(&unsealed[..], &plain[..]);
    }
}