ppvss 0.0.1-alpha

use age::secrecy::ExposeSecret;
use age::{Decryptor, Encryptor, decrypt, encrypt, x25519};
use num_bigint::BigUint;
use secretsharing_shamir::{BitSize, SS, Share};
use serde::{Deserialize, Serialize};
use std::fmt::Error;
use std::fs;
use std::io::{Read, Write};
use std::path::Path;

/// Struct for Pre-constructed Publicly Verifiable Secret Sharing
//#[derive(Debug)]
pub struct PPVSS {
    N: u8,
    T: u8,
    Parties: Vec<Party>, // keeping everythin vector for now, later arrays could be used wherever possible
                         /*H0: BigUint,*/
}
impl PPVSS {
    ///
    pub fn new(
        n: u8,
        t: u8,
        parties: Vec<Party>, /* h0: BigUint*/
    ) -> Result<Self, Box<dyn std::error::Error>> {
        Ok(Self {
            N: n,
            T: t,
            Parties: parties,
            /*H0 : h0,*/
        })
    }

    /// Share geneartion
    /// INPUT: n, t, f_0, {pk_i}_{i=1}^n, h_0 or pk_0
    /// h_0 is used when using commitments
    /// pk_0 is used when using encryption
    ///
    /// OUTPUT: {y_i}_{i=0}^n, Proof_ppvss
    /// Proof_ppvss proves the encryption of f_i to y_i via resp pk_i and f_0 to y_0 via h_0 or pk_0
    /// Step 1: Creating shares of f0: f_1 ... f_n
    /// Step 2: Encrypting each f_i using respective pk_i of parties for all parties
    /// Step 3: Encrypting (pk0,sk0) or committing (h0) f0: y0
    /// Step 4: Generate proof of step 2 and step 3 PV NIZK (Publicly verifiable NIZK) proof
    pub fn share(&self, secret: BigUint) -> (Vec<Share>, BigUint) {
        // secret = f0 in the original paper
        // Initial settings
        let prime_size = BitSize::BN254; // Issue 2: have prime_sizr as parameter for share()
        let mut ss = SS::new(prime_size, true, self.T, &secret).unwrap();
        //let prime = ss.get_prime();
        let mut points: Vec<BigUint> = vec![]; // Issue 1: add parameter to give custom points to evaluate shares on, give a boolean
        for i in 0..self.N as usize {
            points.push(prime_size.n_bit_random()); // Issue 5: can we use public ket as x input to the share generation
            //points.push(self.Parties[i].public_key.)
        }

        // Step 1
        let shares = ss.gen_shares(&points);

        // Step 2
        let mut encrypted_shares: Vec<Share> = vec![];
        for i in 0..self.N as usize {
            let fi = shares[i].y().to_bytes_be();

            let yi_vecu8 = encrypt(&self.Parties[i].public_key, &fi).unwrap();
            let yi = BigUint::from_bytes_be(&yi_vecu8);

            encrypted_shares.push(Share::new(shares[i].x().clone(), yi)); // consume yi here
        }
        // Step 3: Todo. Generate commitment or enc of f0(secret) ---h0 or pk0 --> y0
        //

        // Step 4: Todo. Generate proof for the 2nd nad 3rd step

        // return shares vector and proof string
        (encrypted_shares, BigUint::from(23u32))
    }

    /// Verify proofs
    /// INPUT: n, t, {pk_i}_{i=0}^n, Proof_ppvss, note the pk's are from i=0 to n
    /// OUTPUT: bool (true or false)
    /// Note: Publicly verifiable proof, which means verifier may or may not be a participant in share holder
    pub fn verify(
        n: u8,
        pk: &[BigUint], /*including the dealer's pk */
        t: u8,
        proof: BigUint,
    ) -> u8 {
        11u8
    }

    pub fn reconstruct(&self, shares: Vec<Share>, SK: Vec<KeyPair>) -> Result<BigUint, Error> {
        // Issue 3: Does the parameter for prime bitsize need to be added? Or should it be fixed for the PPVSS
        // Issue 4:
        let mut decrypted_shares: Vec<Share> = vec![];
        for i in 0..SK.len() as usize {
            let fi = shares[i].y().to_bytes_be();

            let fi_vecu8 = decrypt(&SK[i].private_key, &fi).unwrap();
            let fi = BigUint::from_bytes_be(&fi_vecu8);

            decrypted_shares.push(Share::new(shares[i].x().clone(), fi)); // consume yi here
        }
        let rs = SS::reconstruct_secret(&BitSize::BN254.fixed_prime(), &decrypted_shares);

        Ok(rs)
    }
}

/// This struct is to be used for generating keys only,
/// otherwise one only needs public keys from parties to encrypt stuff
/// and private keys are only required during reconstruction
//#[derive(Debug)]
pub struct KeyPair {
    private_key: x25519::Identity,
    public_key: x25519::Recipient,
}

impl KeyPair {
    /// Generate a new keypair
    pub fn generate() -> Self {
        let private_key = x25519::Identity::generate();
        let public_key = private_key.to_public();
        Self {
            private_key,
            public_key,
        }
    }

    /// Save private key to file
    pub fn save_private_key<P: AsRef<Path>>(
        &self,
        path: P,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let sk_str = self.private_key.to_string();
        fs::write(path, sk_str.expose_secret())?;
        Ok(())
    }

    /// Save public key to file
    pub fn save_public_key<P: AsRef<Path>>(
        &self,
        path: P,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let public_key_str = self.public_key.to_string();
        fs::write(path, public_key_str)?;
        Ok(())
    }

    /// Load private key from file and derive public key, returns a new KeyPair
    pub fn load_from_private_key<P: AsRef<Path>>(
        path: P,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let private_key_str = fs::read_to_string(path)?;
        let private_key: x25519::Identity = private_key_str.trim().parse()?;
        let public_key = private_key.to_public();

        Ok(Self {
            private_key,
            public_key,
        })
    }

    /// Get public key as string
    pub fn public_key_string(&self) -> String {
        self.public_key.to_string()
    }
}

/// This struct is used as a field inside the PPVSS struct.
/// It has only the public keys of users and meant for encryption only.
/// The Party can only be instantiated if you give a file as input containing the public key.
#[derive(Debug, Clone)]
pub struct Party {
    public_key: x25519::Recipient,
    share: Share, // share contains an x and encrypt(poly(x))
                  // Issue 5: can we use public ket as x input to the share generation
}
impl Party {
    /// Load public key from file and create new Party (for encryption only)
    pub fn new<P: AsRef<Path>>(path: P) -> Self {
        let public_key_str = fs::read_to_string(path).unwrap();
        let public_key: x25519::Recipient = public_key_str.trim().parse().unwrap();
        //Ok(public_key)
        let empty_share = Share::new(BigUint::from(0u8), BigUint::from(0u8));
        Self {
            public_key: public_key,
            share: empty_share,
        }
    }

    /// Encrypt data using a public key
    pub fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, age::EncryptError> {
        match age::encrypt(&self.public_key, data) {
            Ok(ncryptd) => Ok(ncryptd),
            Err(e) => Err(e), // propagate error manually
        }
    }
    /// Get public key as string
    pub fn public_key_string(&self) -> String {
        self.public_key.to_string()
    }
}

/*
pub struct ReconstructSet{
    Parties: Vec<KeyPair>,
    Shares: Vec<Share>,
}*/