//! This module implements "The Simplest Protocol for Oblivious Transfer" due to
//! Orlandi and Chou.
//! (cf. https://eprint.iacr.org/archive/2015/267/1527602042.pdf)
//!
//! The protocol works as follows in an elliptic curve group G with base point `B` and scalars `Scalars`
//!
//! ```text
//! Sender(l, r)            Receiver(c)
//! y <-$ Scalars
//! S := yB
//! T := yS    -- S -->     x <-$ Scalars
//!                         R := cS + xB
//!            <-- R --
//! k_l                     k = H(S, R, xS)
//!  = H(S, R, yR)       
//! k_r
//!  = H(S, R, yR - T)
//!
//! c_l = E(k_l, l)
//! c_r = E(k_r, r)
//!
//!          -- c_l -->
//!          -- c_r -->     output = D(k, c_l) if decryption successful
//!                         otherwise output = D(k, c_r)
//! ```
//! We instantiate the primitives as follows:
//!     - H: HKDF(SHA-256)
//!     - group G: P256
//!     - Encryption scheme: Chacha20Poly1305

use hacspec_chacha20poly1305::{ChaChaPolyIV, ChaChaPolyKey};
use hacspec_lib::Randomness;
use p256::{p256_point_mul, p256_point_mul_base, P256Point};

use crate::Error;

/// The state of the sender
pub struct OTSender {
    y: p256::P256Scalar,
    s: p256::P256Point,
    t: p256::P256Point,
    dst: Vec<u8>,
}

/// The state of the receiver
pub struct OTReceiver {
    x: p256::P256Scalar,
    r: p256::P256Point,
    s: P256Point,
    dst: Vec<u8>,
}

/// The OT sender's first message.
#[derive(Debug)]
pub struct OTSenderInit(p256::P256Point);

/// The OT receiver's first message.
#[derive(Debug)]
pub struct OTReceiverSelect(p256::P256Point);

/// The encryption of an OT input.
#[derive(Debug)]
pub struct OTCiphertext {
    iv: ChaChaPolyIV,
    ciphertext: Vec<u8>,
    tag: [u8; 16],
}
/// The OT sender's second message.
#[derive(Debug)]
pub struct OTSenderSend {
    left: OTCiphertext,
    right: OTCiphertext,
}

impl OTSender {
    /// Generate the first sender message.
    ///
    /// Initiates an OT sender by picking a random P256 scalar `y` and deriving
    /// `S = yB` and `T = yS`, where `B` is the P256 base point. These values
    /// will later be used to derive encryption keys in the send stage of the
    /// protocol. In addition, the domain separation tag `dst`, which will be
    /// used in key generation is stored in the receiver and `S` is prepared for
    /// sending to the receiver by wrapping it in an `OTSenderInit` message.
    pub fn init(entropy: &mut Randomness, dst: &[u8]) -> Result<(Self, OTSenderInit), Error> {
        let y = p256::random_scalar(entropy, dst)?;

        let s = p256::p256_point_mul_base(y)?;
        let t = p256_point_mul(y, s)?;

        Ok((
            OTSender {
                y,
                s,
                t,
                dst: dst.to_vec(),
            },
            OTSenderInit(s),
        ))
    }

    /// Generate the second sender message based on the receiver's selection.
    ///
    /// Given the `OTReceiverSelect` message and the two sender inputs, the
    /// sender can generate the transfer messages. It does so by deriving two
    /// domain separated encryption keys, based on the values `S` and `T`
    /// generated during initiation and on the masked choice bit sent by the
    /// receiver. It then encrypts the left and right inputs under their
    /// respective keys and prepares an `OTSenderSend` message with both
    /// ciphertexts. This finishes the OT session for the sender. By the
    /// security of the protocol, the receiver will only be able to generate one
    /// of the decryption keys, namely that one corresponding to its choice bit.
    pub fn send(
        &self,
        left_input: &[u8],
        right_input: &[u8],
        selection: &OTReceiverSelect,
        entropy: &mut Randomness,
    ) -> Result<OTSenderSend, Error> {
        debug_assert_eq!(
            left_input.len(),
            right_input.len(),
            "Left and right inputs to the OT must be of the same length."
        );
        let OTReceiverSelect(r) = selection;

        let (left_key, right_key) = self.derive_keys(r)?;

        let (left, right) = encrypt_inputs(entropy, left_key, left_input, right_key, right_input);

        Ok(OTSenderSend { left, right })
    }

    /// Derive a pair of encryption keys for creating the `OTSenderSend`
    /// message.
    ///
    /// Given the points `S`,`T` from the sender initialization, as well as the
    /// masked receiver choice `R`, a pre-key is generated based on `S` and `R`.
    /// Then the keys are generated by first extracting
    /// `HKDF-SHA256-extract(pre-key||yR)` for the right input and
    /// `HKDF-SHA256-extract(pre-key||yR - T)` for the left input where the
    /// elliptic curve points are serialized and concatenated to the pre-key and
    /// a fixed salt value is used to extract. From this, two 32-byte keys are
    /// expanded for use in ChaCha20Poly1305.
    fn derive_keys(
        &self,
        receiver_selection: &p256::P256Point,
    ) -> Result<(ChaChaPolyKey, ChaChaPolyKey), Error> {
        let (salt, ikm) = derive_prk(&self.s, receiver_selection);

        let input_right = p256_point_mul(self.y, *receiver_selection)?;

        let input_left = p256::point_add(input_right, -self.t)?;

        let input_left_serialized = p256::serialize_point(&input_left);
        let input_right_serialized = p256::serialize_point(&input_right);

        let mut ikm_left = ikm.clone();
        let mut ikm_right = ikm;

        ikm_left.extend_from_slice(&input_left_serialized);
        ikm_right.extend_from_slice(&input_right_serialized);

        let prk_left = hmac::hkdf_extract(&salt, &ikm_left);
        let prk_right = hmac::hkdf_extract(&salt, &ikm_right);
        Ok((
            hmac::hkdf_expand(&prk_left, &self.dst, 32)
                .try_into()
                .expect(
                    "should have received the right number of bytes, because we requested them",
                ),
            hmac::hkdf_expand(&prk_right, &self.dst, 32)
                .try_into()
                .expect(
                    "should have received the right number of bytes, because we requested them",
                ),
        ))
    }
}

/// Create the pre-key for encrypting and decrypting `OTSenderSend` ciphertexts.
///
/// Based on points `S` and `R`, use `HKDF-SHA256-extract(S||R)` to generate a
/// pre-key with a fixed salt value by serializing and concatenating the points.
fn derive_prk(
    sender_commitment: &p256::P256Point,
    receiver_selection: &p256::P256Point,
) -> (Vec<u8>, Vec<u8>) {
    let serialized_s = p256::serialize_point(sender_commitment);
    let serialized_r = p256::serialize_point(receiver_selection);
    let salt = b"no-salt";
    let mut ikm = Vec::from(serialized_s);
    ikm.extend_from_slice(&serialized_r);
    (salt.to_vec(), ikm)
}

/// Encrypt the OT sender inputs.
///
/// Using the keys generated by `derive_keys()` encrypt the sender's inputs with
/// Chacha20Poly1305, without any additional authenticated data.
fn encrypt_inputs(
    entropy: &mut Randomness,
    left_key: [u8; 32],
    left_input: &[u8],
    right_key: [u8; 32],
    right_input: &[u8],
) -> (OTCiphertext, OTCiphertext) {
    let left_iv = entropy
        .bytes(12)
        .expect("sufficient randomness should have been provided externally")
        .try_into()
        .expect("should have received the right number of bytes, because we requested them");

    let right_iv = entropy
        .bytes(12)
        .expect("sufficient randomness should have been provided externally")
        .try_into()
        .expect("should have received the right number of bytes, because we requested them");

    let (left_enc, left_tag) =
        hacspec_chacha20poly1305::chacha20_poly1305_encrypt(left_key, left_iv, &[], left_input);
    let (right_enc, right_tag) =
        hacspec_chacha20poly1305::chacha20_poly1305_encrypt(right_key, right_iv, &[], right_input);
    (
        OTCiphertext {
            iv: left_iv,
            ciphertext: left_enc,
            tag: left_tag,
        },
        OTCiphertext {
            iv: right_iv,
            ciphertext: right_enc,
            tag: right_tag,
        },
    )
}

impl OTReceiver {
    /// Generate the first receiver message.
    ///
    /// Initiates the OT receiver by generating a random P256 Scalar `x` which
    /// is used to mask the receivers choice bit `c` of output in the
    /// computation of the masked receiver selection `R = cS + xB`, where `S` is
    /// obtained from the initial sender message. The values required in the
    /// decryption of the sender messages are stored in the `OTReceiver` struct
    /// and the masked choice is wrapped in an `OTReceiverSelect` message for
    /// sending to the sender.
    pub fn select(
        entropy: &mut Randomness,
        dst: &[u8],
        sender_message: OTSenderInit,
        choose_left: bool,
    ) -> Result<(Self, OTReceiverSelect), Error> {
        let x = p256::random_scalar(entropy, dst)?;
        let OTSenderInit(s) = sender_message;

        let mut res = p256_point_mul_base(x)?;
        let r = if choose_left {
            res = p256::point_add(res, s)?;
            res
        } else {
            res
        };

        Ok((
            OTReceiver {
                x,
                r,
                s,
                dst: dst.to_vec(),
            },
            OTReceiverSelect(r),
        ))
    }

    /// Receive the selected input from the sender.
    ///
    /// A decryption key is generated based on the initialization message
    /// received from the sender and the masked choice generated during receiver
    /// initialization. Then trial-decryption of the sender messages are
    /// attempted. Correctness of the protocol guarantees that one decryption
    /// will be successful (exactly one, by security of the protocol) and will
    /// yield the receiver's chosen output ending the OT session for the
    /// receiver.
    pub fn receive(&self, sender_message: OTSenderSend) -> Result<Vec<u8>, Error> {
        let key = self.derive_key()?;

        let dec = hacspec_chacha20poly1305::chacha20_poly1305_decrypt(
            key,
            sender_message.left.iv,
            &[],
            &sender_message.left.ciphertext,
            sender_message.left.tag,
        )
        .or_else(|_| {
            hacspec_chacha20poly1305::chacha20_poly1305_decrypt(
                key,
                sender_message.right.iv,
                &[],
                &sender_message.right.ciphertext,
                sender_message.right.tag,
            )
        });

        dec.map_err(|e| e.into())
    }

    /// Derive a decryption key for trial decrypting the `OTSenderSend`
    /// messages.
    ///
    /// Given the point `S` from the initial sender message, as well as the
    /// scalar `x` generated during receiver initialization and the masked
    /// receiver choice `R`, a pre-key is generated based on `S` and `R`. Then
    /// the key is generated by first extracting
    ///  `HKDF-SHA256-extract(pre-key||xS)` where the elliptic curve points are
    /// serialized and concatenated to the pre-key and a fixed salt value is
    /// used to extract. From this, a 32-byte key is expanded for use in
    /// ChaCha20Poly1305.
    fn derive_key(&self) -> Result<ChaChaPolyKey, Error> {
        let (salt, mut ikm) = derive_prk(&self.s, &self.r);

        let input = p256_point_mul(self.x, self.s)?;
        let input_serialized = p256::serialize_point(&input);

        ikm.extend_from_slice(&input_serialized);

        let prk = hmac::hkdf_extract(&salt, &ikm);

        Ok(hmac::hkdf_expand(&prk, &self.dst, 32)
            .try_into()
            .expect("should have received the required number of bytes because we requested them"))
    }
}

#[test]
fn simple() {
    use rand::RngCore;
    let mut rng = rand::thread_rng();
    let mut entropy = [0u8; 88];
    rng.fill_bytes(&mut entropy);
    let mut entropy = Randomness::new(entropy.to_vec());

    let dst = b"test-context";
    let left_input = b"lefto";
    let right_input = b"right";
    let (sender, commitment) = OTSender::init(&mut entropy, dst).unwrap();
    let (receiver, selection) = OTReceiver::select(&mut entropy, dst, commitment, false).unwrap();

    let send_message = sender
        .send(left_input, right_input, &selection, &mut entropy)
        .unwrap();

    let receiver_output = receiver.receive(send_message).unwrap();
    debug_assert_eq!(right_input.to_vec(), receiver_output);
}