// Bitcoin secp256k1 bindings
// Written in 2014 by
//   Dawid Ciężarkiewicz
//   Andrew Poelstra
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! Public/Private keys

use std::intrinsics::copy_nonoverlapping_memory;
use std::cmp;
use std::default::Default;
use std::fmt;
use std::ptr::zero_memory;
use std::rand::Rng;
use serialize::{Decoder, Decodable, Encoder, Encodable};

use secretdata::SecretData;
use crypto::digest::Digest;
use crypto::sha2::Sha512;
use crypto::hmac::Hmac;
use crypto::mac::Mac;

use super::init;
use super::{Result, InvalidNonce, InvalidPublicKey, InvalidSecretKey, Unknown};
use constants;
use ffi;

/// Secret 256-bit nonce used as `k` in an ECDSA signature
pub struct Nonce([u8, ..constants::NONCE_SIZE]);
impl_array_newtype!(Nonce, u8, constants::NONCE_SIZE)

/// Secret 256-bit key used as `x` in an ECDSA signature
pub struct SecretKey<'a>(SecretData<'a, SecretKeyData>);

/// Secret 256-bit key used as `x` in an ECDSA signature
struct SecretKeyData([u8, ..constants::SECRET_KEY_SIZE]);
impl_array_newtype!(SecretKeyData, u8, constants::SECRET_KEY_SIZE)

impl Default for SecretKeyData {
    fn default() -> SecretKeyData {
        SecretKeyData([0, ..constants::SECRET_KEY_SIZE])
    }
}

/// Public key
#[deriving(Clone, PartialEq, Eq, Show)]
pub struct PublicKey(PublicKeyData);

enum PublicKeyData {
    Compressed([u8, ..constants::COMPRESSED_PUBLIC_KEY_SIZE]),
    Uncompressed([u8, ..constants::UNCOMPRESSED_PUBLIC_KEY_SIZE]),
}

fn random_32_bytes<R:Rng>(rng: &mut R) -> [u8, ..32] {
    let mut ret = [0u8, ..32];
    rng.fill_bytes(ret);
    ret
}

/// As described in RFC 6979
fn bits2octets(data: &[u8]) -> [u8, ..32] {
    let mut ret = [0, ..32];
    unsafe {
        copy_nonoverlapping_memory(ret.as_mut_ptr(),
                                   data.as_ptr(),
                                   cmp::min(data.len(), 32));
    }
    ret
}

impl Nonce {
    /// Creates a new random nonce
    #[inline]
    pub fn new<R:Rng>(rng: &mut R) -> Nonce {
        Nonce(random_32_bytes(rng))
    }

    /// Converts a `NONCE_SIZE`-byte slice to a nonce
    #[inline]
    pub fn from_slice(data: &[u8]) -> Result<Nonce> {
        match data.len() {
            constants::NONCE_SIZE => {
                let mut ret = [0, ..constants::NONCE_SIZE];
                unsafe {
                    copy_nonoverlapping_memory(ret.as_mut_ptr(),
                                               data.as_ptr(),
                                               data.len());
                }
                Ok(Nonce(ret))
            }
            _ => Err(InvalidNonce)
        }
    }

    /// Generates a deterministic nonce by RFC6979 with HMAC-SHA512
    #[inline]
    #[allow(non_snake_case)] // so we can match the names in the RFC
    pub fn deterministic<'a>(msg: &[u8], key: &SecretKey<'a>) -> Nonce {
        static HMAC_SIZE: uint = 64;

        macro_rules! hmac(
            ($res:expr <- key $key:expr, data $($data:expr),+) => ({
                let mut hmacker = Hmac::new(Sha512::new(), $key.as_slice());
                $(hmacker.input($data.as_slice());)+
                hmacker.raw_result($res.as_mut_slice());
            })
        )

        // Section 3.2a
        // Goofy block just to avoid marking `msg_hash` as mutable
        let mut hasher = Sha512::new();
        hasher.input(msg);
        let mut x = [0, ..HMAC_SIZE];
        hasher.result(x.as_mut_slice());
        let msg_hash = bits2octets(x.as_slice());

        // Section 3.2b
        let mut V = [0x01u8, ..HMAC_SIZE];
        // Section 3.2c
        let mut K = [0x00u8, ..HMAC_SIZE];

        // Section 3.2d
        hmac!(K <- key K, data V, [0x00], key, msg_hash)

        // Section 3.2e
        hmac!(V <- key K, data V)

        // Section 3.2f
        hmac!(K <- key K, data V, [0x01], key, msg_hash)

        // Section 3.2g
        hmac!(V <- key K, data V)

        // Section 3.2
        let mut k = Err(InvalidSecretKey);
        while k.is_err() {
            // Try to generate the nonce
            let mut T = [0x00u8, ..HMAC_SIZE];
            hmac!(T <- key K, data V)

            k = Nonce::from_slice(T.slice_to(constants::NONCE_SIZE));

            // Replace K, V
            if k.is_err() {
                hmac!(K <- key K, data V, [0x00])
                hmac!(V <- key K, data V)
            }
        }

        k.unwrap()
    }
}

impl<'a> SecretKey<'a> {
    /// Creates a new zeroed-out secret key
    #[inline]
    pub fn new() -> SecretKey<'a> {
        SecretKey(SecretData::new())
    }

    /// Creates a new random secret key
    #[inline]
    pub fn init_rng<R:Rng>(&'a mut self, rng: &mut R) {
        init();
        let mut data = random_32_bytes(rng);
        unsafe {
            while ffi::secp256k1_ecdsa_seckey_verify(data.as_ptr()) == 0 {
                data = random_32_bytes(rng);
            }
        }
        let &SecretKey(ref mut selfdata) = self;
        selfdata.move(&mut SecretKeyData(data))
    }

    /// Converts a `SECRET_KEY_SIZE`-byte slice to a secret key,
    /// zeroing out the original data
    #[inline]
    pub fn init_slice(&'a mut self, data: &mut [u8]) -> Result<()> {
        init();
        match data.len() {
            constants::SECRET_KEY_SIZE => {
                let &SecretKey(ref mut selfdata) = self;
                unsafe {
                    if ffi::secp256k1_ecdsa_seckey_verify(data.as_ptr()) == 0 {
                        return Err(InvalidSecretKey);
                    }
                    copy_nonoverlapping_memory(selfdata.data_mut().as_mut_ptr(),
                                               data.as_ptr(),
                                               data.len());
                    zero_memory(data.as_mut_ptr(), data.len());
                }
                Ok(())
            }
            _ => Err(InvalidSecretKey)
        }
    }

    /// Copies the data from one key to another without zeroing anyth out
    #[inline]
    pub fn clone_from<'b>(&'a mut self, other: &SecretKey<'b>) {
        let &SecretKey(ref mut selfdata) = self;
        let &SecretKey(ref otherdata) = other;
        selfdata.clone_from(otherdata);
    }

    #[inline]
    /// Adds one secret key to another, modulo the curve order
    /// Marked `unsafe` since you must
    /// call `init()` (or construct a `Secp256k1`, which does this for you) before
    /// using this function
    pub fn add_assign<'b>(&mut self, other: &SecretKey<'b>) -> Result<()> {
        init();
        unsafe {
            if ffi::secp256k1_ecdsa_privkey_tweak_add(self.as_mut_ptr(), other.as_ptr()) != 1 {
                Err(Unknown)
            } else {
                Ok(())
            }
        }
    }

    #[inline]
    /// Returns an immutable view of the data as a byteslice
    pub fn as_slice<'b>(&'b self) -> &'b [u8] {
        let &SecretKey(ref selfdata) = self;
        selfdata.data().as_slice()
    }

    #[inline]
    /// Returns a raw pointer to the underlying secret key data
    pub fn as_ptr(&self) -> *const u8 {
        let &SecretKey(ref selfdata) = self;
        selfdata.data().as_ptr()
    }

    #[inline]
    /// Returns a mutable raw pointer to the underlying secret key data
    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        let &SecretKey(ref mut selfdata) = self;
        selfdata.data_mut().as_mut_ptr()
    }
}

impl PublicKey {
    /// Creates a new zeroed out public key
    #[inline]
    pub fn new(compressed: bool) -> PublicKey {
        PublicKey(
            if compressed { Compressed([0, ..constants::COMPRESSED_PUBLIC_KEY_SIZE]) }
            else { Uncompressed([0, ..constants::UNCOMPRESSED_PUBLIC_KEY_SIZE]) }
        )
    }

    /// Creates a new public key from a secret key.
    #[inline]
    pub fn from_secret_key<'a>(sk: &SecretKey<'a>, compressed: bool) -> PublicKey {
        let mut pk = PublicKey::new(compressed);
        let compressed = if compressed {1} else {0};
        let mut len = 0;

        init();
        unsafe {
            // We can assume the return value because it's not possible to construct
            // an invalid `SecretKey` without transmute trickery or something
            assert_eq!(ffi::secp256k1_ecdsa_pubkey_create(
                pk.as_mut_ptr(), &mut len,
                sk.as_ptr(), compressed), 1);
        }
        assert_eq!(len as uint, pk.len()); 
        pk
    }

    /// Creates a public key directly from a slice
    #[inline]
    pub fn from_slice(data: &[u8]) -> Result<PublicKey> {
        match data.len() {
            constants::COMPRESSED_PUBLIC_KEY_SIZE => {
                let mut ret = [0, ..constants::COMPRESSED_PUBLIC_KEY_SIZE];
                unsafe {
                    if ffi::secp256k1_ecdsa_pubkey_verify(data.as_ptr(),
                                                          data.len() as ::libc::c_int) == 0 {
                        return Err(InvalidPublicKey);
                    }
                    copy_nonoverlapping_memory(ret.as_mut_ptr(),
                                               data.as_ptr(),
                                               data.len());
                }
                Ok(PublicKey(Compressed(ret)))
            }
            constants::UNCOMPRESSED_PUBLIC_KEY_SIZE => {
                let mut ret = [0, ..constants::UNCOMPRESSED_PUBLIC_KEY_SIZE];
                unsafe {
                    copy_nonoverlapping_memory(ret.as_mut_ptr(),
                                               data.as_ptr(),
                                               data.len());
                }
                Ok(PublicKey(Uncompressed(ret)))
            }
            _ => Err(InvalidPublicKey)
        }
    }

    /// Returns whether the public key is compressed or uncompressed
    #[inline]
    pub fn is_compressed(&self) -> bool {
        let &PublicKey(ref data) = self;
        match *data {
            Compressed(_) => true,
            Uncompressed(_) => false
        }
    }

    /// Returns the length of the public key
    #[inline]
    pub fn len(&self) -> uint {
        let &PublicKey(ref data) = self;
        match *data {
            Compressed(ref x) => x.len(),
            Uncompressed(ref x) => x.len()
        }
    }

    /// Converts the public key into a byte slice
    #[inline]
    pub fn as_slice<'a>(&'a self) -> &'a [u8] {
        let &PublicKey(ref data) = self;
        data.as_slice()
    }

    /// Converts the public key to a raw pointer suitable for use
    /// with the FFI functions
    #[inline]
    pub fn as_ptr(&self) -> *const u8 {
        let &PublicKey(ref data) = self;
        match *data {
            Compressed(ref x) => x.as_ptr(),
            Uncompressed(ref x) => x.as_ptr()
        }
    }

    /// Converts the public key to a mutable raw pointer suitable for use
    /// with the FFI functions
    #[inline]
    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        let &PublicKey(ref mut data) = self;
        match *data {
            Compressed(ref mut x) => x.as_mut_ptr(),
            Uncompressed(ref mut x) => x.as_mut_ptr()
        }
    }

    #[inline]
    /// Adds the pk corresponding to `other` to the pk `self` in place
    pub fn add_exp_assign<'a>(&mut self, other: &SecretKey<'a>) -> Result<()> {
        init();
        unsafe {
            if ffi::secp256k1_ecdsa_pubkey_tweak_add(self.as_mut_ptr(),
                                                     self.len() as ::libc::c_int,
                                                     other.as_ptr()) != 1 {
                Err(Unknown)
            } else {
                Ok(())
            }
        }
    }
}

impl PublicKeyData {
    #[inline]
    fn as_slice<'a>(&'a self) -> &'a [u8] {
        match *self {
            Compressed(ref x) => x.as_slice(),
            Uncompressed(ref x) => x.as_slice()
        }
    }
}

// We have to do all these impls ourselves as Rust can't derive
// them for arrays
impl fmt::Show for Nonce {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

impl Clone for PublicKeyData {
    fn clone(&self) -> PublicKeyData { *self }
}

impl PartialEq for PublicKeyData {
    fn eq(&self, other: &PublicKeyData) -> bool {
        self.as_slice() == other.as_slice()
    }
}

impl Eq for PublicKeyData {}

impl fmt::Show for PublicKeyData {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

impl<D: Decoder<E>, E> Decodable<D, E> for PublicKey {
    fn decode(d: &mut D) -> ::std::prelude::Result<PublicKey, E> {
        d.read_seq(|d, len| {
            if len == constants::UNCOMPRESSED_PUBLIC_KEY_SIZE {
                unsafe {
                    use std::mem;
                    let mut ret: [u8, ..constants::UNCOMPRESSED_PUBLIC_KEY_SIZE] = mem::uninitialized();
                    for i in range(0, len) {
                        ret[i] = try!(d.read_seq_elt(i, |d| Decodable::decode(d)));
                    }
                    Ok(PublicKey(Uncompressed(ret)))
                }
            } else if len == constants::COMPRESSED_PUBLIC_KEY_SIZE {
                unsafe {
                    use std::mem;
                    let mut ret: [u8, ..constants::COMPRESSED_PUBLIC_KEY_SIZE] = mem::uninitialized();
                    for i in range(0, len) {
                        ret[i] = try!(d.read_seq_elt(i, |d| Decodable::decode(d)));
                    }
                    Ok(PublicKey(Compressed(ret)))
                }
            } else {
                Err(d.error("Invalid length"))
            }
        })
    }
}

impl <E: Encoder<S>, S> Encodable<E, S> for PublicKey {
    fn encode(&self, e: &mut E) -> ::std::prelude::Result<(), S> {
        self.as_slice().encode(e)
    }
}

impl<'a> PartialEq for SecretKey<'a> {
    fn eq(&self, other: &SecretKey<'a>) -> bool {
        self.as_slice() == other.as_slice()
    }
}

impl<'a> Eq for SecretKey<'a> {}

impl<'a> fmt::Show for SecretKey<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "[secret data]")
    }
}

#[cfg(test)]
mod test {
    use serialize::hex::FromHex;
    use std::rand::task_rng;

    use super::super::{InvalidNonce, InvalidPublicKey, InvalidSecretKey};
    use super::{Nonce, PublicKey, SecretKey};

    #[test]
    fn nonce_from_slice() {
        let n = Nonce::from_slice([1, ..31]);
        assert_eq!(n, Err(InvalidNonce));

        let mut n = SecretKey::new();
        assert_eq!(n.init_slice([1, ..32]), Ok(()));
    }

    #[test]
    fn skey_from_slice() {
        let mut sk = SecretKey::new();
        assert_eq!(sk.init_slice([1, ..31]), Err(InvalidSecretKey));
        let mut sk = SecretKey::new();
        assert_eq!(sk.init_slice([1, ..32]), Ok(()));
    }

    #[test]
    fn pubkey_from_slice() {
        assert_eq!(PublicKey::from_slice([]), Err(InvalidPublicKey));
        assert_eq!(PublicKey::from_slice([1, 2, 3]), Err(InvalidPublicKey));

        let uncompressed = PublicKey::from_slice([4, 54, 57, 149, 239, 162, 148, 175, 246, 254, 239, 75, 154, 152, 10, 82, 234, 224, 85, 220, 40, 100, 57, 121, 30, 162, 94, 156, 135, 67, 74, 49, 179, 57, 236, 53, 162, 124, 149, 144, 168, 77, 74, 30, 72, 211, 229, 110, 111, 55, 96, 193, 86, 227, 183, 152, 195, 155, 51, 247, 123, 113, 60, 228, 188]);
        assert!(uncompressed.is_ok());
        assert!(!uncompressed.unwrap().is_compressed());

        let compressed = PublicKey::from_slice([3, 23, 183, 225, 206, 31, 159, 148, 195, 42, 67, 115, 146, 41, 248, 140, 11, 3, 51, 41, 111, 180, 110, 143, 114, 134, 88, 73, 198, 174, 52, 184, 78]);
        assert!(compressed.is_ok());
        assert!(compressed.unwrap().is_compressed());
    }

    #[test]
    fn keypair_slice_round_trip() {
        let mut rng = task_rng();
        let mut sk1 = SecretKey::new();
        sk1.init_rng(&mut rng);
        let mut sk2 = SecretKey::new();
        sk2.clone_from(&sk1);

        assert_eq!(sk1, sk2);

        let pk1 = PublicKey::from_secret_key(&sk1, false);
        assert_eq!(PublicKey::from_slice(pk1.as_slice()), Ok(pk1));
        let pk2 = PublicKey::from_secret_key(&sk1, true);
        assert_eq!(PublicKey::from_slice(pk2.as_slice()), Ok(pk2));
    }

    #[test]
    fn nonce_slice_round_trip() {
        let mut rng = task_rng();
        let nonce = Nonce::new(&mut rng);
        assert_eq!(Nonce::from_slice(nonce.as_slice()), Ok(nonce));
    }

    #[test]
    fn invalid_secret_key() {
        let mut sk = SecretKey::new();
        // Zero
        assert_eq!(sk.init_slice([0, ..32]), Err(InvalidSecretKey));
        // -1
        assert_eq!(sk.init_slice([0xff, ..32]), Err(InvalidSecretKey));
        // Top of range
        assert_eq!(sk.init_slice([0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
                                   0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFE,
                                   0xBA, 0xAE, 0xDC, 0xE6, 0xAF, 0x48, 0xA0, 0x3B,
                                   0xBF, 0xD2, 0x5E, 0x8C, 0xD0, 0x36, 0x41, 0x40]), Ok(()));
        // One past top of range
        assert_eq!(sk.init_slice([0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
                                  0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFE,
                                  0xBA, 0xAE, 0xDC, 0xE6, 0xAF, 0x48, 0xA0, 0x3B,
                                  0xBF, 0xD2, 0x5E, 0x8C, 0xD0, 0x36, 0x41, 0x41]), Err(InvalidSecretKey));
    }

    #[test]
    fn test_addition() {
        let mut rng = task_rng();

        let mut sk1 = SecretKey::new();
        let mut sk2 = SecretKey::new();
        sk1.init_rng(&mut rng);
        sk2.init_rng(&mut rng);
        let mut pk1 = PublicKey::from_secret_key(&sk1, true);
        let mut pk2 = PublicKey::from_secret_key(&sk2, true);

        assert_eq!(PublicKey::from_secret_key(&sk1, true), pk1);
        assert!(sk1.add_assign(&sk2).is_ok());
        assert!(pk1.add_exp_assign(&sk2).is_ok());
        assert_eq!(PublicKey::from_secret_key(&sk1, true), pk1);

        assert_eq!(PublicKey::from_secret_key(&sk2, true), pk2);
        assert!(sk2.add_assign(&sk1).is_ok());
        assert!(pk2.add_exp_assign(&sk1).is_ok());
        assert_eq!(PublicKey::from_secret_key(&sk2, true), pk2);
    }

    #[test]
    fn test_deterministic() {
        // nb code in comments is equivalent python

        // from ecdsa import rfc6979
        // from ecdsa.curves import SECP256k1
        // # This key was generated randomly
        // sk = 0x09e918bbea76205445e9a73eaad2080a135d1e33e9dd1b3ca8a9a1285e7c1f81
        let mut sk = SecretKey::new();
        sk.init_slice(hex_slice_mut!("09e918bbea76205445e9a73eaad2080a135d1e33e9dd1b3ca8a9a1285e7c1f81")).unwrap();
        assert_eq!(sk.as_slice(),
                   hex_slice!("09e918bbea76205445e9a73eaad2080a135d1e33e9dd1b3ca8a9a1285e7c1f81"));

        // "%x" % rfc6979.generate_k(SECP256k1.generator, sk, hashlib.sha512, hashlib.sha512('').digest())
        let nonce = Nonce::deterministic([], &sk);
        assert_eq!(nonce.as_slice(),
                   hex_slice!("d954eddd184cac2b60edcd0e6be9ec54d93f633b28b366420d38ed9c346ffe27"));

        // "%x" % rfc6979.generate_k(SECP256k1.generator, sk, hashlib.sha512, hashlib.sha512('test').digest())
        let nonce = Nonce::deterministic(b"test", &sk);
        assert_eq!(nonce.as_slice(),
                   hex_slice!("609cc24acce2f19e46e38a82afc56c1745dee16e04f2b27e24999e1fefeb08bd"));

        // # Decrease the secret key by one
        // sk = 0x09e918bbea76205445e9a73eaad2080a135d1e33e9dd1b3ca8a9a1285e7c1f80
        sk.init_slice(hex_slice_mut!("09e918bbea76205445e9a73eaad2080a135d1e33e9dd1b3ca8a9a1285e7c1f80")).unwrap();

        // "%x" % rfc6979.generate_k(SECP256k1.generator, sk, hashlib.sha512, hashlib.sha512('').digest())
        let nonce = Nonce::deterministic([], &sk);
        assert_eq!(nonce.as_slice(),
                   hex_slice!("9f45f8d0a28e8956673c8da6db3db86ca4f172f0a2dbd62364fdbf786c7d96df"));

        // "%x" % rfc6979.generate_k(SECP256k1.generator, sk, hashlib.sha512, hashlib.sha512('test').digest())
        let nonce = Nonce::deterministic(b"test", &sk);
        assert_eq!(nonce.as_slice(),
                   hex_slice!("355c589ff662c838aee454d62b12c50a87b7e95ede2431c7cfa40b6ba2fddccd"));
    }
}