rabe 0.4.2

use std::{
    string::String,
    ops::Neg
};
use rabe_bn::{Group, Gt, G1, G2, Fr, pairing};

use rand::Rng;
use utils::{
    tools::*,
    secretsharing::*,
    aes::*,
    hash::sha3_hash
};
use utils::policy::pest::{PolicyLanguage, parse, PolicyType};
use crate::error::RabeError;
#[cfg(feature = "serde")]
use serde::{Serialize, Deserialize};

/// An Ghw11 Public Key (PK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Ghw11PublicKey {
    pub g1: G1,
    pub g2: G2,
    pub g1_a: G1,
    pub g2_a: G2,
    pub e_gg_alpha: Gt,
}

/// An Ghw11 Master Key (MSK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Ghw11MasterKey {
    pub g2_alpha: G2,
    pub pk: Ghw11PublicKey,
}

/// An Ghw11 Secret Key (SK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Ghw11SecretKey {
    //g2^alpha * g2^ar
    pub k: G2,
    //g2^r
    pub l: G2,
    //F(x)^r
    pub attr_key: Vec<Ghw11Attribute>,
}

/// A Ghw11 Attribute
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Ghw11Attribute {
    pub string: String,
    pub k_x: G2,
}

/// An Ghw11 Transform Key (TK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Ghw11TransformKey {
    pub k_z: G2,
    pub l_z: G2,
    pub attr_key_z: Vec<Ghw11Attribute>,
}

/// An Ghw11 Retrieve Key (RK)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Ghw11RetrieveKey {
    pub z: Fr,
}

/// A Ghw11 Ciphertext (CT)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Ghw11Ciphertext {
    pub policy: (String, PolicyLanguage),
    pub c : Gt,
    pub c1: G1,
    pub ci_di: Vec<(String, G1, G1)>,
    pub data: Vec<u8>,
}

/// A Ghw11 Transform Ciphertext (TCT)
#[derive(Clone, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Ghw11TransformCiphertext {
    pub c : Gt,
    pub t : Gt,
}

/// The setup algorithm of Ghw11. Generates a Ghw11PublicKey and a Ghw11MasterKey.
pub fn setup() -> (Ghw11PublicKey, Ghw11MasterKey) {
    // random number generator
    let mut rng = rand::thread_rng();
    let g1:G1 = rng.gen();
    let g2:G2 = rng.gen();

    let a:Fr = rng.gen();
    let g1_a = g1 * a;
    let g2_a = g2 * a;

    let alpha:Fr = rng.gen();
    let e_gg_alpha = pairing(g1, g2).pow(alpha);
    let g2_alpha = g2 * alpha;
    let pk = Ghw11PublicKey { g1, g1_a, g2, g2_a, e_gg_alpha};
    // return pk and mk
    (
        pk.clone(),
        Ghw11MasterKey { g2_alpha, pk },
    )
}

/// The key generation algorithm of Ghw11 CP-ABE. Generates a Ghw11SecrectKey using a Ghw11PublicKey, a Ghw11MasterKey and a set of attributes given as Vec<String>.
///
/// # Arguments
///
///	* `pk` - A Public Key (PK), generated by the function setup()
///	* `msk` - A Master Key (MSK), generated by the function setup()
///	* `attributes` - A Vector of String attributes assigned to this user key
///
pub fn keygen(
    pk: &Ghw11PublicKey,
    msk: &Ghw11MasterKey,
    attributes: &[String],
) -> Option<Ghw11SecretKey> {
    // if no attibutes or an empty policy
    // maybe add empty msk also here
    if attributes.is_empty() || attributes.len() == 0 {
        return None;
    }
    // random number generator
    let mut rng = rand::thread_rng();
    // generate random r
    let r:Fr = rng.gen();

    //L
    let g2_r = pk.g2 * r;

    //K
    let k = msk.g2_alpha + pk.g2_a * r;

    let mut k_x: Vec<Ghw11Attribute> = Vec::new();
    for j in attributes {
        k_x.push(Ghw11Attribute {
            string: j.to_string(), // attribute name
            k_x:  (sha3_hash(pk.g2, j.as_str()).expect("could not hash _j") * r), // K_x
        });
    }
    return Some(Ghw11SecretKey { k, l: g2_r, attr_key: k_x });
}

/// The tansform key generation algorithm of Ghw11 CP-ABE. 
/// Tansfrom Secretkey with a random RetrieveKey z to TransformKey, return (Ghw11TransformKey, Ghw11RetrieveKey).
pub fn tkgen(
    sk: Ghw11SecretKey,
) -> Option<(Ghw11TransformKey, Ghw11RetrieveKey)> {
        // random number generator
        let mut rng = rand::thread_rng();

        // generate random z
        let z:Fr = rng.gen();
        let z_inverse = z.inverse().unwrap();

        let k_z = sk.k * z_inverse;

        let l_z = sk.l * z_inverse;

        let mut attr_key_z: Vec<Ghw11Attribute> = Vec::new();
        for k_i in sk.attr_key.iter() {
            attr_key_z.push(Ghw11Attribute {
                string: k_i.string.clone(), // attribute name
                k_x:  k_i.k_x * z_inverse, 
            });
        }
        return Some((Ghw11TransformKey { k_z, l_z, attr_key_z }, Ghw11RetrieveKey{z}));
}

/// The encrypt algorithm of Ghw11 CP-ABE. Generates a new Ghw11Ciphertext using an Ghw11PublicKey, an access policy given as String and some plaintext data given as [u8].
///
/// # Arguments
///
///	* `pk` - A Public Key (PK), generated by the function setup()
///	* `policy` - An access policy given as JSON String
///	* `language` - The policy language
///	* `plaintext` - plaintext data given as a Vector of u8
///
pub fn encrypt(
    pk: &Ghw11PublicKey,
    policy: &str,
    language: PolicyLanguage,
    plaintext: &[u8],
) -> Result<Ghw11Ciphertext, RabeError> {
    if plaintext.is_empty() || policy.is_empty() {
        RabeError::new("Error in bsw/encrypt: data or policy is empty.");
    }
    let mut rng = rand::thread_rng();
    // the shared root secret
    let secret:Fr = rng.gen();

    let msg: Gt = rng.gen();

    match parse(policy, language) {
        Ok(policy_value) => {
            let shares: Vec<(String, Fr)> = gen_shares_policy(secret, &policy_value, None).unwrap();

            let c = pk.e_gg_alpha.pow(secret) * msg;
            let c1 = pk.g1 * secret;

            let mut ci_di: Vec<(String, G1, G1)> = Vec::new();
            for (node, i_val) in shares.clone() {
                let t_i:Fr = rng.gen();
                let j = remove_index(&node);
                ci_di.push((node.clone(), pk.g1_a * i_val + sha3_hash(pk.g1, &j).unwrap() * (t_i.neg()), pk.g1 * t_i));
            }
            match encrypt_symmetric(msg, &plaintext.to_vec()) {
                Ok(data) => Ok(Ghw11Ciphertext { policy: (policy.to_string(), language), c, c1, ci_di, data }),
                Err(e) => Err(e)
            }
        }
        Err(e) => Err(e)
    }
}

/// The transform algorithm of Ghw11 CP-ABE. Generates a new Ghw11TransformCiphertext using an Ghw11TransformKey and Ghw11Ciphertext.
pub fn transform(
    ct: Ghw11Ciphertext,
    tk: Ghw11TransformKey,
) -> Result<Ghw11TransformCiphertext, RabeError> {
    let str_attr: Vec<String> = tk
        .attr_key_z
        .iter()
        .map(|_values| {
            _values.clone().string
            
        })
        .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 tk 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 mut coeff_list: Vec<(String, Fr)> = Vec::new();
                        coeff_list = calc_coefficients(&pol, Some(Fr::one()), coeff_list,None).unwrap();
                        if _match {
                            let mut t = Gt::one();
                            let mut ci_wi = G1::zero();
                            for _current in _list.iter() {
                                //w_i
                                let _coeff = coeff_list
                                    .iter()
                                    .filter(|_c| _c.0 == _current.1.to_string())
                                    .map(|_c| _c.1)
                                    .nth(0)
                                    .unwrap();
                                //k_pi
                                let _tk_attr = tk
                                    .attr_key_z
                                    .iter()
                                    .filter(|_attr| _attr.string == _current.0.to_string())
                                    .nth(0)
                                    .unwrap();
                                //ci di
                                let _ct_attr = ct
                                    .ci_di
                                    .iter()
                                    .filter(|_attr| _attr.0 == _current.1.to_string())
                                    .nth(0)
                                    .unwrap();
                                //add ci^wi
                                ci_wi = ci_wi + _ct_attr.1 * _coeff;
                                //mul 
                                t = t * pairing(_ct_attr.2 * _coeff, _tk_attr.k_x);
                            }
                            t = t * pairing(ci_wi, tk.l_z);
                            t = pairing(ct.c1, tk.k_z) * t.inverse();

                            Ok(Ghw11TransformCiphertext{c: ct.c, t})
                        } else {
                            Err(RabeError::new("Error in Ghw11/decrypt: attributes in sk do not match policy in ct."))
                        }
                    }
                }
            }
        },
        Err(e) => Err(e)
    }    
}

/// The decrypt_out algorithm of GHW11 CP-ABE. Reconstructs the original plaintext data as Vec<u8>, given a Ghw11TransformCiphertext with a matching Ghw11RetrieveKey.
pub fn decrypt_out(
    pct: Ghw11TransformCiphertext,
    rk: Ghw11RetrieveKey,
    data: Vec<u8>,
) -> Result<Vec<u8>, RabeError> {
    let msg = pct.c * (pct.t.pow(rk.z)).inverse();
    decrypt_symmetric(msg, &data)
}

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

    #[test]
    fn or() {
        // setup scheme
        let (pk, msk) = setup();
        // a set of two attributes matching the policy
        let mut att_matching = Vec::new();
        att_matching.push(String::from("D"));
        att_matching.push(String::from("B"));

        // a set of two attributes NOT matching the policy
        let mut att_not_matching = Vec::new();
        att_not_matching.push(String::from("C"));
        att_not_matching.push(String::from("D"));

        //match key gen
        let match_sk = keygen(&pk, &msk, &att_matching).unwrap();

        //not match key gen
        let not_match_sk = keygen(&pk, &msk, &att_not_matching).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": "A"}, {"name": "B"}]}"#);

        // cp-abe ciphertext
        let ct_cp = encrypt(&pk, &policy, PolicyLanguage::JsonPolicy, &plaintext).unwrap();

        //tk gen
        let (match_tk, match_rk) = tkgen(match_sk).unwrap();

        let (not_match_tk, _not_match_rk) = tkgen(not_match_sk).unwrap();

        //transform
        let transform_ct = transform(ct_cp.clone(), match_tk).unwrap();

        let not_match_transform_ct = transform(ct_cp.clone(), not_match_tk);
        assert_eq!(not_match_transform_ct.is_ok(), false);

        let _match = decrypt_out(transform_ct, match_rk, ct_cp.data.clone());
        assert_eq!(_match.is_ok(), true);
        assert_eq!(_match.unwrap(), plaintext);
    }

    #[test]
    fn and2() -> (){
        // global setup
        let (pk, msk) = setup();

        // setup a user "bob" and give him some attribute-keys

        let att_bob = vec![String::from("attr0"),String::from("attr1")];

        let bob_sk = keygen(&pk, &msk, &att_bob).unwrap();

        // our plaintext
        let plaintext =
            String::from("ghw11 OABE").into_bytes();

        // our policy
        let policy = String::from(r#"{"name": "and", "children": [{"name": "attr0"}, {"name": "attr1"}]}"#);

        // cp-abe ciphertext
        let ct_cp = encrypt(&pk, &policy, PolicyLanguage::JsonPolicy, &plaintext).unwrap();

        //tk gen
        let (tk, rk) = tkgen(bob_sk).unwrap();

        //transform
        let transform_ct = transform(ct_cp.clone(), tk).unwrap();

        // and now decrypt again with mathcing tk
        let _matching = decrypt_out(transform_ct, rk, ct_cp.data).unwrap();

        assert_eq!(_matching, plaintext);

        let decrypt_plaintext = String::from_utf8(_matching).unwrap();

        println!("decrypt plaintext:{:?}", decrypt_plaintext);
    }


    #[test]
    fn and10() {
        // setup scheme
        let (pk, msk) = setup();
        // a set of two attributes matching the policy
        let mut att_matching: Vec<String> = Vec::new();
        for n in 1..11 {
            att_matching.push(["attr".to_string(), n.to_string()].concat());
        }
        // a set of two attributes NOT matching the policy
        let mut att_not_matching: Vec<String> = Vec::new();
        att_not_matching.push(String::from("attr201"));
        att_not_matching.push(String::from("attr200"));

        //match key gen
        let match_sk = keygen(&pk, &msk, &att_matching).unwrap();

        //not match key gen
        let not_match_sk = keygen(&pk, &msk, &att_not_matching).unwrap();

        // our plaintext
        let plaintext = String::from("dance like no one's watching, encrypt like everyone is!")
            .into_bytes();

        let mut _policy = String::from("{\"name\": \"and\", \"children\": [");
        for n in 1..11 {
            let mut _current = String::from("{\"name\": \"attr");
            if n < 10 {
                _current.push_str(&n.to_string());
                _current.push_str(&String::from("\"}, "));
            } else {
                _current.push_str(&n.to_string());
                _current.push_str(&String::from("\"}]"));
            }
            _policy.push_str(&_current);
        }
        _policy.push_str(&String::from("}"));
        // cp-abe ciphertext
        let ct_cp = encrypt(&pk, &_policy, PolicyLanguage::JsonPolicy, &plaintext).unwrap();

        //tk gen
        let (match_tk, match_rk) = tkgen(match_sk).unwrap();

        let (not_match_tk, _not_match_rk) = tkgen(not_match_sk).unwrap();

        //transform
        let transform_ct = transform(ct_cp.clone(), match_tk).unwrap();

        let not_match_transform_ct = transform(ct_cp.clone(), not_match_tk);
        assert_eq!(not_match_transform_ct.is_ok(), false);

        let _match = decrypt_out(transform_ct, match_rk, ct_cp.data.clone());
        assert_eq!(_match.is_ok(), true);
        assert_eq!(_match.unwrap(), plaintext);
    }

}