//! `AW11` scheme by Lewko, Allison, and Brent Waters.
//!
//! * Developped by Lewko, Allison, and Brent Waters, "Decentralizing Attribute-Based Encryption.", see Appendix D
//! * Published in Eurocrypt 2011
//! * Available from <http://eprint.iacr.org/2010/351.pdf>
//! * Type:			encryption (identity-based)
//! * Setting:		bilinear groups (asymmetric)
//! * Authors:		Georg Bramm
//! * Date:			04/2018
//!
//! # Examples
//!
//! ```
//!use rabe::schemes::aw11::*;
//! use rabe::utils::policy::pest::PolicyLanguage;
//!let gk = setup();
//!let (pk, msk) = authgen(&gk, &vec!["A".to_string(), "B".to_string()]).unwrap();
//!let plaintext = String::from("our plaintext!").into_bytes();
//!let policy = String::from(r#""A" or "B""#);
//!let bob = keygen(&gk, &msk, &String::from("bob"), &vec!["A".to_string()]).unwrap();
//!let ct: Aw11Ciphertext = encrypt(&gk, &vec![pk], &policy, PolicyLanguage::HumanPolicy, &plaintext).unwrap();
//!let matching = decrypt(&gk, &bob, &ct).unwrap();
//!assert_eq!(matching, plaintext);
//! ```
use std::string::String;
use rand::Rng;
use rabe_bn::{Fr, G1, G2, Gt, pairing};
use utils::{
    secretsharing::{
        calc_coefficients,
        calc_pruned
    },
    policy::msp::AbePolicy,
    tools::*,
    aes::*,
    hash::sha3_hash
};
use utils::policy::pest::{PolicyLanguage, parse, PolicyType};
use utils::secretsharing::gen_shares_policy;
use crate::error::RabeError;
#[cfg(not(feature = "borsh"))]
use serde::{Serialize, Deserialize};
#[cfg(feature = "borsh")]
use borsh::{BorshSerialize, BorshDeserialize};

/// An AW11 Global Parameters Key (GK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "borsh", derive(BorshSerialize, BorshDeserialize))]
#[cfg_attr(not(feature = "borsh"), derive(Serialize, Deserialize))]
pub struct Aw11GlobalKey {
    pub _g1: G1,
    pub _g2: G2,
}

/// An AW11 Public Key (PK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "borsh", derive(BorshSerialize, BorshDeserialize))]
#[cfg_attr(not(feature = "borsh"), derive(Serialize, Deserialize))]
pub struct Aw11PublicKey {
    pub _attr: Vec<(String, Gt, G2)>,
}

/// An AW11 Master Key (MK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "borsh", derive(BorshSerialize, BorshDeserialize))]
#[cfg_attr(not(feature = "borsh"), derive(Serialize, Deserialize))]
pub struct Aw11MasterKey {
    pub _attr: Vec<(String, Fr, Fr)>,
}

/// An AW11 Ciphertext (CT)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "borsh", derive(BorshSerialize, BorshDeserialize))]
#[cfg_attr(not(feature = "borsh"), derive(Serialize, Deserialize))]
pub struct Aw11Ciphertext {
    pub _policy: (String, PolicyLanguage),
    pub _c_0: Gt,
    pub _c: Vec<(String, Gt, G2, G2)>,
    pub _ct: Vec<u8>,
}

/// An AW11 Secret Key (SK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "borsh", derive(BorshSerialize, BorshDeserialize))]
#[cfg_attr(not(feature = "borsh"), derive(Serialize, Deserialize))]
pub struct Aw11SecretKey {
    pub _gid: String,
    pub _attr: Vec<(String, G1)>,
}

/// A global Context for an AW11 Global Parameters Key (GP)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "borsh", derive(BorshSerialize, BorshDeserialize))]
#[cfg_attr(not(feature = "borsh"), derive(Serialize, Deserialize))]
pub struct Aw11GlobalContext {
    pub _gk: Aw11GlobalKey,
}

/// Sets up a new AW11 Scheme by creating a Global Parameters Key (GK)
pub fn setup() -> Aw11GlobalKey {
    // random number generator
    let mut _rng = rand::thread_rng();
    // generator of group G1: g1 and generator of group G2: g2
    let _gk = Aw11GlobalKey {
        _g1: _rng.gen(),
        _g2: _rng.gen(),
    };
    // return PK and MSK
    return _gk;
}

/// Sets up a new AW11 Authority by creating a key pair: Public Parameters Key (PK) and Master Key (MK)
///
/// # Arguments
///
///	* `_gk` - A Global Parameters Key (GK), generated by the function setup()
///	* `_attributes` - A Vector of String attributes assigned to this Authority
///
/// # Remarks
///
/// In this scheme, all attributes are converted to upper case bevor calculation, i.e. they are case insensitive
/// This means that all attributes "tEsT", "TEST" and "test" are the same in this scheme.
pub fn authgen(
    _gk: &Aw11GlobalKey,
    _attributes: &Vec<String>,
) -> Option<(Aw11PublicKey, Aw11MasterKey)> {
    // if no attibutes or an empty policy
    // maybe add empty msk also here
    if _attributes.is_empty() {
        return None;
    }
    // random number generator
    let mut _rng = rand::thread_rng();
    // generator of group G1: g and generator of group G2: h
    let mut _sk: Vec<(String, Fr, Fr)> = Vec::new(); //dictionary of {s: {alpha_i, y_i}}
    let mut _pk: Vec<(String, Gt, G2)> = Vec::new(); // dictionary of {s: {e(g,g)^alpha_i, g1^y_i}}
    // now calculate attribute values
    for _attr in _attributes {
        // calculate randomness
        let _alpha_i:Fr = _rng.gen();
        let _y_i:Fr = _rng.gen();
        _sk.push((_attr.clone().to_uppercase(), _alpha_i, _y_i));
        _pk.push((
            _attr.clone().to_uppercase(),
            pairing(_gk._g1, _gk._g2).pow(_alpha_i),
            _gk._g2 * _y_i,
        ));
    }
    // return PK and MSK
    return Some((Aw11PublicKey { _attr: _pk }, Aw11MasterKey { _attr: _sk }));
}

/// Sets up and generates a new User by creating a secret user key (SK). The key is created for a user with a given "name" on the given set of attributes.
///
/// # Arguments
///
///	* `_gk` - A Global Parameters Key (GK), generated by setup()
///	* `_msk` - A Master Key (MK), associated with an authority and generated by authgen()
///	* `_name` - The name of the user the key is associated with. Must be unique.
///	* `_attributes` - A Vector of String attributes assigned to this User
///
/// # Remarks
///
/// In this scheme, all attributes are converted to upper case bevor calculation, i.e. they are case insensitive
/// This means that all attributes "tEsT", "TEST" and "test" are treated the same in this scheme.
pub fn keygen(
    _gk: &Aw11GlobalKey,
    _msk: &Aw11MasterKey,
    _name: &String,
    _attributes: &Vec<String>,
) -> Result<Aw11SecretKey, RabeError> {
    // if no attibutes or no gid
    if _attributes.is_empty() {
        Err(RabeError::new("empty _attributes"))
    }
    else if _name.is_empty() {
        Err(RabeError::new("empty _name"))
    }
    else {
        let mut _sk: Aw11SecretKey = Aw11SecretKey {
            _gid: _name.clone(),
            _attr: Vec::new(),
        };
        for _attribute in _attributes {
            if let Err(e) =  add_to_attribute(_gk, _msk, _attribute, &mut _sk) {
                return Err(e);
            }
        }
        Ok(_sk)
    }
}

/// This function does not create a new User key, but adds a new attribute to an already generated key (SK).
///
/// # Arguments
///
///	* `_gk` - A Global Parameters Key (GK), generated by setup()
///	* `_msk` - A Master Key (MK), associated with an authority and generated by authgen()
///	* `_attribute` - A String attribute that should be added to the already existing key (SK)
///	* `_sk` - The secret user key (SK)
pub fn add_to_attribute(
    _gk: &Aw11GlobalKey,
    _msk: &Aw11MasterKey,
    _attribute: &String,
    _sk: &mut Aw11SecretKey,
) -> Result<(), RabeError> {
    // if no attibutes or no gid
    if _attribute.is_empty() {
        Err(RabeError::new("empty _attributes"))
    }
    else if _sk._gid.is_empty() {
        Err(RabeError::new("empty _gid"))
    }
    else {
        match sha3_hash(_gk._g1, &_sk._gid) {
            Ok(hash) => {
                let _auth_attribute = _msk
                    ._attr
                    .iter()
                    .filter(|_attr| _attr.0 == _attribute.to_string())
                    .nth(0)
                    .unwrap();
                _sk._attr.push((
                    _auth_attribute.0.clone().to_uppercase(),
                    (_gk._g1 * _auth_attribute.1) + (hash * _auth_attribute.2),
                ));
                Ok(())
            },
            Err(e) => Err(e)
        }
    }
}

/// This function encrypts plaintext data using a given JSON String policy and produces a 'Aw11Ciphertext' if successfull.
///
/// # Arguments
///
///	* `_gk` - A Global Parameters Key (GK), generated by setup()
///	* `_pk` - A Public Parameters Key (MK), associated with an authority and generated by authgen()
///	* `_policy` - A JSON String policy describing the access rights
///	* `_plaintext` - The plaintext data given as a Vector of u8.
pub fn encrypt(
    _gk: &Aw11GlobalKey,
    _pks: &Vec<Aw11PublicKey>,
    _policy: &String,
    _language: PolicyLanguage,
    _plaintext: &[u8],
) -> Result<Aw11Ciphertext, RabeError> {
    // random number generator
    let mut _rng = rand::thread_rng();
    match parse(_policy, _language) {
        Ok(pol) => {
            // an msp policy from the given String
            let msp: AbePolicy = AbePolicy::from_policy(&pol).unwrap();
            let _num_cols = msp._m[0].len();
            let _num_rows = msp._m.len();
            // pick randomness
            let _s:Fr = _rng.gen();
            // and calculate shares "s" and "zero"
            let _s_shares = gen_shares_policy(_s, &pol, None).unwrap();
            let _w_shares = gen_shares_policy(Fr::zero(), &pol, None).unwrap();
            // calculate c0 with a randomly selected "msg"
            let _msg: Gt = _rng.gen();
            let _c_0 = _msg * pairing(_gk._g1, _gk._g2).pow(_s);
            // now calculate the C1,x C2,x and C3,x parts
            let mut _c: Vec<(String, Gt, G2, G2)> = Vec::new();
            for (_i, (_attr_name, _attr_share)) in _s_shares.into_iter().enumerate() {
                let _r_x:Fr = _rng.gen();
                let _pk_attr = find_pk_attr(_pks, &_attr_name.to_uppercase());
                match _pk_attr {
                    None => {},
                    Some(_attr) => {
                        _c.push((
                            _attr_name.clone().to_uppercase(),
                            pairing(_gk._g1, _gk._g2).pow(_attr_share) * _attr.1.pow(_r_x),
                            _gk._g2 * _r_x,
                            (_attr.2 * _r_x) + (_gk._g2 * _w_shares[_i].1),
                        ));
                    }
                }
            }
            //Encrypt plaintext using derived key from secret
            let _policy = _policy.to_string();
            let _ct = encrypt_symmetric(_msg, &_plaintext.to_vec()).unwrap();
            Ok(Aw11Ciphertext { _policy: (_policy, _language), _c_0, _c, _ct })
        },
        Err(e) => Err(e)
    }
}

/// This function decrypts a 'Aw11Ciphertext' if the attributes in SK match the policy of CT. If successfull, returns the plaintext data as a Vetor of u8's.
///
/// # Arguments
///
///	* `_gk` - A Global Parameters Key (GK), generated by setup()
///	* `_sk` - A secret user key (SK), associated with a set of attributes.
///	* `_ct` - A Aw11Ciphertext
pub fn decrypt(
    gk: &Aw11GlobalKey,
    sk: &Aw11SecretKey,
    ct: &Aw11Ciphertext) -> Result<Vec<u8>, RabeError> {
    let _str_attr = sk
        ._attr
        .iter()
        .map(|_values| {
            let (_str, _g2) = _values.clone();
            _str
        })
        .collect::<Vec<_>>();
    return match parse(ct._policy.0.as_ref(), ct._policy.1) {
        Ok(pol) => {
            return if traverse_policy(&_str_attr, &pol, PolicyType::Leaf) == false {
                Err(RabeError::new("Error: attributes in sk do not match policy in ct."))
            } else {
                let _pruned = calc_pruned(&_str_attr, &pol, None);
                match _pruned {
                    Err(e) => Err(e),
                    Ok(_p) => {
                        let (_match, _list) = _p;
                        let _coeffs = calc_coefficients(&pol, Some(Fr::one()), None).unwrap();
                        if _match {
                            match sha3_hash(gk._g1, &sk._gid) {
                                Ok(hash) => {
                                    let mut _egg_s = Gt::one();
                                    for _current in _list.iter() {
                                        let _sk_attr = sk
                                            ._attr
                                            .iter()
                                            .filter(|_attr| _attr.0 == _current.to_string())
                                            .nth(0)
                                            .unwrap();
                                        let _ct_attr = ct
                                            ._c
                                            .iter()
                                            .filter(|_attr| _attr.0 == _current.to_string())
                                            .nth(0)
                                            .unwrap();
                                        let num = _ct_attr.1 * pairing(hash, _ct_attr.3);
                                        let dem = pairing(_sk_attr.1, _ct_attr.2);
                                        let _coeff = _coeffs
                                            .iter()
                                            .filter(|_c| _c.0 == _current.to_string())
                                            .map(|_c| _c.1)
                                            .nth(0)
                                            .unwrap();
                                        _egg_s = _egg_s * ((num * dem.inverse()).pow(_coeff));
                                    }
                                    let _msg = ct._c_0 * _egg_s.inverse();
                                    //println!("dec: {:?}", serde_json::to_string(&_msg).unwrap());
                                    // Decrypt plaintext using derived secret from cp-abe scheme
                                    decrypt_symmetric(_msg, &ct._ct)
                                },
                                Err(e) => Err(e)
                            }
                        } else {
                            Err(RabeError::new("Error in aw11/decrypt: attributes in sk do not match policy in ct."))
                        }
                    }
                }
            }
        },
        Err(e) => Err(e)
    }
}
/// private function. finds the value vector of a specific attribute in a vector of various public keys
///
/// # Arguments
///
///	* `_pks` - A vector of Aw11PublicKeys
///	* `_attr` - An attribute
///
fn find_pk_attr(_pks: &Vec<Aw11PublicKey>, _attr: &String) -> Option<(String, Gt, G2)> {
    for _pk in _pks.into_iter() {
        let _pk_attr = _pk
            ._attr
            .clone()
            .into_iter()
            .filter(|_tuple| _tuple.0 == _attr.to_string())
            .nth(0);
        if _pk_attr.is_some() {
            return _pk_attr;
        }
    }
    return None;
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn and() {
        // global setup
        let _gp = setup();
        // setup attribute authority 1 with
        // a set of two attributes "A" "B" "C"
        let mut att_authority1: Vec<String> = Vec::new();
        att_authority1.push(String::from("A"));
        att_authority1.push(String::from("B"));
        att_authority1.push(String::from("C"));
        let (_auth1_pk, _auth1_msk) = authgen(&_gp, &att_authority1).unwrap();
        // setup attribute authority 1 with
        // a set of two attributes "D" "E" "F"
        let mut att_authority2: Vec<String> = Vec::new();
        att_authority2.push(String::from("D"));
        att_authority2.push(String::from("E"));
        att_authority2.push(String::from("F"));
        let (_auth2_pk, _auth2_msk) = authgen(&_gp, &att_authority2).unwrap();
        // setup attribute authority 1 with
        // a set of two attributes "G" "H" "I"
        let mut att_authority3: Vec<String> = Vec::new();
        att_authority3.push(String::from("G"));
        att_authority3.push(String::from("H"));
        att_authority3.push(String::from("I"));
        let (_auth3_pk, _auth3_msk) = authgen(&_gp, &att_authority3).unwrap();

        // setup a user "bob" and give him some attribute-keys
        let mut att_bob: Vec<String> = Vec::new();
        att_bob.push(String::from("H"));
        att_bob.push(String::from("I"));
        let mut _bob = keygen(&_gp, &_auth3_msk, &String::from("bob"), &att_bob).unwrap();
        // our plaintext
        let _plaintext =
            String::from("dance like no one's watching, encrypt like everyone is!").into_bytes();
        // our policy
        let _policy = String::from(r#"{"name": "and", "children": [{"name": "H"}, {"name": "B"}]}"#);

        // add attribute "B" of auth1 to bobs key
        if let Err(e) = add_to_attribute(&_gp, &_auth1_msk, &String::from("B"), &mut _bob) {
            panic!("Error: {}", e.to_string())
        }

        // a vector of public attribute keys
        let mut _pks: Vec<Aw11PublicKey> = Vec::new();
        _pks.push(_auth3_pk);
        _pks.push(_auth1_pk);

        // cp-abe ciphertext
        let ct_cp: Aw11Ciphertext = encrypt(&_gp, &_pks, &_policy, PolicyLanguage::JsonPolicy, &_plaintext).unwrap();
        // and now decrypt again with mathcing sk
        let _matching = decrypt(&_gp, &_bob, &ct_cp).unwrap();
        assert_eq!(_matching, _plaintext);
    }

    #[test]
    fn or() {
        // global setup
        let _gp = setup();
        // setup attribute authority 1 with
        // a set of two attributes "A" "B"
        let mut att_authority1: Vec<String> = Vec::new();
        att_authority1.push(String::from("A"));
        att_authority1.push(String::from("B"));
        let (_auth1_pk, _auth1_msk) = authgen(&_gp, &att_authority1).unwrap();
        // setup attribute authority 1 with
        // a set of two attributes "C" "D"
        let mut att_authority2: Vec<String> = Vec::new();
        att_authority2.push(String::from("C"));
        att_authority2.push(String::from("D"));
        let (_auth2_pk, _auth2_msk) = authgen(&_gp, &att_authority2).unwrap();

        // setup a user "bob" and give him some attribute-keys
        let mut _bob = keygen(
            &_gp,
            &_auth1_msk,
            &String::from("bob"),
            &vec![String::from("A")],
        )
        .unwrap();

        // our plaintext
        let _plaintext =
            String::from("dance like no one's watching, encrypt like everyone is!").into_bytes();
        // our policy
        let _policy = String::from(r#"{"name": "or", "children": [{"name": "B"}, {"name": "C"}]}"#);

        // a vector of public attribute keys
        let mut _pks: Vec<Aw11PublicKey> = Vec::new();
        _pks.push(_auth2_pk);
        _pks.push(_auth1_pk);

        // add attribute "C" of auth2 to bobs key
        if let Err(e) = add_to_attribute(&_gp, &_auth2_msk, &String::from("C"), &mut _bob) {
            panic!("Error: {}", e.to_string())
        }
        // cp-abe ciphertext
        let ct_cp: Aw11Ciphertext = encrypt(&_gp, &_pks, &_policy,PolicyLanguage::JsonPolicy,  &_plaintext).unwrap();
        // and now decrypt again with mathcing sk
        let _matching = decrypt(&_gp, &_bob, &ct_cp).unwrap();
        assert_eq!(_matching, _plaintext);
    }

    #[test]
    fn or_and() {
        // global setup
        let _gp = setup();
        // setup attribute authority 1 with
        // a set of two attributes "A" "B"
        let mut att_authority1: Vec<String> = Vec::new();
        att_authority1.push(String::from("A"));
        att_authority1.push(String::from("B"));
        let (_auth1_pk, _auth1_msk) = authgen(&_gp, &att_authority1).unwrap();
        // setup attribute authority 1 with
        // a set of two attributes "C" "D"
        let mut att_authority2: Vec<String> = Vec::new();
        att_authority2.push(String::from("C"));
        att_authority2.push(String::from("D"));
        let (_auth2_pk, _auth2_msk) = authgen(&_gp, &att_authority2).unwrap();

        // setup a user "bob" and give him some attribute-keys
        let mut _bob = keygen(
            &_gp,
            &_auth1_msk,
            &String::from("bob"),
            &vec![String::from("A")],
        )
        .unwrap();

        // our plaintext
        let _plaintext =
            String::from("dance like no one's watching, encrypt like everyone is!").into_bytes();
        // our policy
        let _policy =
            String::from(r#"{"name": "or", "children": [{"name": "B"}, {"name": "and", "children": [{"name": "C"}, {"name": "D"}]}]}"#);

        // a vector of public attribute keys
        let mut _pks: Vec<Aw11PublicKey> = Vec::new();
        _pks.push(_auth2_pk);
        _pks.push(_auth1_pk);

        // add attribute "C" and "D" of auth2 to bobs key
        if let Err(e) = add_to_attribute(&_gp, &_auth2_msk, &String::from("C"), &mut _bob) {
            panic!("Error: {}", e.to_string())
        }
        if let Err(e) = add_to_attribute(&_gp, &_auth2_msk, &String::from("D"), &mut _bob) {
            panic!("Error: {}", e.to_string())
        }

        // cp-abe ciphertext
        let ct_cp: Aw11Ciphertext = encrypt(&_gp, &_pks, &_policy, PolicyLanguage::JsonPolicy, &_plaintext).unwrap();
        // and now decrypt again with mathcing sk
        let _matching = decrypt(&_gp, &_bob, &ct_cp).unwrap();
        assert_eq!(_matching, _plaintext);
    }

    #[test]
    fn not() {
        // global setup
        let _gp = setup();
        // setup attribute authority 1 with
        // a set of two attributes "A" "B"
        let mut att_authority1: Vec<String> = Vec::new();
        att_authority1.push(String::from("A"));
        att_authority1.push(String::from("B"));
        let (_auth1_pk, _auth1_msk) = authgen(&_gp, &att_authority1).unwrap();
        // setup attribute authority 1 with
        // a set of two attributes "C" "D"
        let mut att_authority2: Vec<String> = Vec::new();
        att_authority2.push(String::from("C"));
        att_authority2.push(String::from("D"));
        let (_auth2_pk, _auth2_msk) = authgen(&_gp, &att_authority2).unwrap();
        // setup a user "bob" and give him some attribute-keys
        let mut _bob = keygen(
            &_gp,
            &_auth1_msk,
            &String::from("bob"),
            &vec![String::from("A")],
        )
        .unwrap();
        // our plaintext
        let _plaintext =
            String::from("dance like no one's watching, encrypt like everyone is!").into_bytes();
        // our policy
        let _policy =
            String::from(r#"{"name": "or", "children": [{"name": "B"}, {"name": "and", "children": [{"name": "C"}, {"name": "A"}]}]}"#);
        // a vector of public attribute keys
        let mut _pks: Vec<Aw11PublicKey> = Vec::new();
        _pks.push(_auth2_pk);
        _pks.push(_auth1_pk);
        // add attribute "C" and "D" of auth2 to bobs key

        if let Err(e) = add_to_attribute(&_gp, &_auth2_msk, &String::from("D"), &mut _bob) {
            panic!("Error: {}", e.to_string())
        }
        // cp-abe ciphertext
        let ct_cp: Aw11Ciphertext = encrypt(&_gp, &_pks, &_policy, PolicyLanguage::JsonPolicy, &_plaintext).unwrap();
        // and now decrypt again
        let pt = decrypt(&_gp, &_bob, &ct_cp);
        assert_eq!(pt.is_ok(), false);
    }
}
