#include "kex.h"
#include "fips202.h"
#include "verify.h"

void kyber_uake_initA(u8 *send, u8 *tk, u8 *sk, const u8 *pkb) {
    crypto_kem_keypair(send, sk);
    crypto_kem_enc(send + KYBER_PUBLICKEYBYTES, tk, pkb);
}

void kyber_uake_sharedB(u8 *send, u8 *k, const u8 *recv, const u8 *skb) {
    unsigned char buf[2 * KYBER_SYMBYTES];
    crypto_kem_enc(send, buf, recv);
    crypto_kem_dec(buf + KYBER_SYMBYTES, recv + KYBER_PUBLICKEYBYTES, skb);
    shake256(k, KYBER_SYMBYTES, buf, 2 * KYBER_SYMBYTES);
}

void kyber_uake_sharedA(u8 *k, const u8 *recv, const u8 *tk, const u8 *sk) {
    unsigned char buf[2 * KYBER_SYMBYTES];
    int i;
    crypto_kem_dec(buf, recv, sk);
    for (i = 0; i < KYBER_SYMBYTES; i++) {
        buf[i + KYBER_SYMBYTES] = tk[i];
    }
    shake256(k, KYBER_SYMBYTES, buf, 2 * KYBER_SYMBYTES);
}

void kyber_ake_initA(u8 *send, u8 *tk, u8 *sk, const u8 *pkb) {
    crypto_kem_keypair(send, sk);
    crypto_kem_enc(send + KYBER_PUBLICKEYBYTES, tk, pkb);
}

void kyber_ake_sharedB(u8 *send, u8 *k, const u8 *recv, const u8 *skb,
                       const u8 *pka) {
    unsigned char buf[3 * KYBER_SYMBYTES];
    crypto_kem_enc(send, buf, recv);
    crypto_kem_enc(send + KYBER_CIPHERTEXTBYTES, buf + KYBER_SYMBYTES, pka);
    crypto_kem_dec(buf + 2 * KYBER_SYMBYTES, recv + KYBER_PUBLICKEYBYTES, skb);
    shake256(k, KYBER_SYMBYTES, buf, 3 * KYBER_SYMBYTES);
}

void kyber_ake_sharedA(u8 *k, const u8 *recv, const u8 *tk, const u8 *sk,
                       const u8 *ska) {
    unsigned char buf[3 * KYBER_SYMBYTES];
    int i;
    crypto_kem_dec(buf, recv, sk);
    crypto_kem_dec(buf + KYBER_SYMBYTES, recv + KYBER_CIPHERTEXTBYTES, ska);
    for (i = 0; i < KYBER_SYMBYTES; i++) {
        buf[i + 2 * KYBER_SYMBYTES] = tk[i];
    }
    shake256(k, KYBER_SYMBYTES, buf, 3 * KYBER_SYMBYTES);
}