/*
 * bsspeke.c - BS-SPEKE over Curve25519
 *
 * Author: Charles V. Wright <cvwright@futo.org>
 * 
 * Copyright (c) 2022 FUTO Holdings, Inc.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "monocypher.h"
#include "bsspeke.h"

void
bsspeke_client_init(bsspeke_client_ctx *ctx,
                    const char* client_id, const size_t client_id_len,
                    const char* server_id, const size_t server_id_len,
                    const char* password, const size_t password_len
) {
    ctx->client_id = (uint8_t *)client_id;
    ctx->client_id_len = client_id_len;

    ctx->server_id = (uint8_t *)server_id;
    ctx->server_id_len = server_id_len;

    ctx->password = (uint8_t *)password;
    ctx->password_len = password_len;
}

void
bsspeke_client_generate_message1(bsspeke_msg1_t *msg1,
                                 bsspeke_client_ctx *client
) {
    // 1. Hash the client's password, client_id, server_id to a point on the curve
    uint8_t scalar_hash[32];
    uint8_t curve_point[32];
    {
        crypto_blake2b_ctx bctx;
        // Give us a 256 bit (32 byte) hash; Don't use a key
        crypto_blake2b_general_init(&bctx, 32, NULL, 0);
        // Add the client id, server id, and the password to the hash
        crypto_blake2b_update(&bctx,
                              (const uint8_t *)(client->password),
                              client->password_len);
        crypto_blake2b_update(&bctx,
                              (const uint8_t *)(client->client_id),
                              client->client_id_len);
        crypto_blake2b_update(&bctx,
                              (const uint8_t *)(client->server_id),
                              client->server_id_len);
        // Write the digest value into `scalar_hash`
        crypto_blake2b_final(&bctx, scalar_hash);
    }
    // Now use Elligator to map our scalar hash to a point on the curve
    crypto_hidden_to_curve(curve_point, scalar_hash);

    // 2. Generate random r
    //    * Actually generate 1/r first, and clamp() it
    //      That way we know it will always lead us back to a point on the curve
    //    * Then use the inverse of 1/r as `r`
    //  FIXME: On second thought, monocypher seems to handle all of this complexity for us.  Let's see what happens if we just do things the straightforward way for now...
    arc4random_buf(client->r, 32);
    crypto_x25519_clamp(client->r);
    // 3. Multiply our curve point by r
    crypto_x25519_scalarmult(msg1->blind, client->r, curve_point, 256);
    return;
}

void
bsspeke_server_init(bsspeke_server_ctx *ctx,
                    const char* server_id
) {
    ctx->server_id = server_id;
}

void
bsspeke_server_generate_message2(bsspeke_msg2_t *msg2,
                                 const bsspeke_msg1_t *msg1,
                                 const uint8_t *salt, const size_t salt_len,
                                 uint8_t P[32], uint8_t V[32],
                                 uint32_t phf_blocks,
                                 uint32_t phf_iterations,
                                 bsspeke_server_ctx *server_ctx
) {
    // Hash the salt
    uint8_t H_salt[32];
    crypto_blake2b_general(H_salt, 32, NULL, 0, salt, salt_len);
    // Use clamp() to ensure we stay on the curve in the multiply below
    crypto_x25519_clamp(H_salt);
    // Multiply H(salt) by msg1->blind, save into msg2->blind_salt
    crypto_x25519_scalarmult(msg2->blind_salt, H_salt, msg1->blind, 256);

    // Generate random ephemeral private key b, save it in ctx->b
    arc4random_buf(server_ctx->b, 32);
    crypto_x25519_clamp(server_ctx->b);

    // Compute public key B = b * P, save it in msg2->B
    crypto_x25519_scalarmult(server_ctx->B, server_ctx->b, P, 256);

    // Copy the public key into the outgoing message as well
    memcpy(msg2->B, server_ctx->B, 32);

    // Copy the PHF params too
    msg2->phf_blocks = phf_blocks;
    msg2->phf_iterations = phf_iterations;

    return;
}

int
bsspeke_client_generate_message3(bsspeke_msg3_t *msg3,
                                 const bsspeke_msg2_t *msg2,
                                 bsspeke_client_ctx *client_ctx
) {
    // Sanity checks first, before we do any work
    if( msg2->phf_blocks < BSSPEKE_MIN_PHF_BLOCKS ) {
        return -1;
    }
    if( msg2->phf_iterations < BSSPEKE_MIN_PHF_ITERATIONS ) {
        return -1;
    }

    uint8_t oblivious_salt[32];
    // Multiply the blinded salt by 1/r to get the oblivious salt
    // Here we rely on Monocypher to do the heavy lifting for us
    crypto_x25519_inverse(oblivious_salt, client_ctx->r, msg2->blind_salt);

    // Hash the oblivious salt together with the id's to create the salt for the PHF
    uint8_t password_hash[64];

    uint8_t phf_salt[32];
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx, oblivious_salt, 32);
        crypto_blake2b_update(&hash_ctx,
                              client_ctx->client_id,
                              client_ctx->client_id_len);
        crypto_blake2b_update(&hash_ctx,
                              client_ctx->server_id,
                              client_ctx->server_id_len);
        crypto_blake2b_final(&hash_ctx, phf_salt);
    }

    void *work_area;
    if ((work_area = malloc(msg2->phf_blocks * 1024)) == NULL) {
        return -1;
    }
    crypto_argon2i(password_hash, 64, work_area,
                   msg2->phf_blocks, msg2->phf_iterations,
                   client_ctx->password, client_ctx->password_len,
                   phf_salt, 32);
    free(work_area);

    // p || v = pwKdf(password, BlindSalt, idC, idS, settings)
    uint8_t *p = &(password_hash[0]);
    uint8_t *v = &(password_hash[32]);

    // Hash p onto the curve to get this user's base point P
    uint8_t P[32];
    crypto_hidden_to_curve(P, p);

    // Generate a random ephemeral private key a, store it in ctx->a
    arc4random_buf(client_ctx->a, 32);
    crypto_x25519_clamp(client_ctx->a);
    // Generate the ephemeral public key A = a * P, store it in msg3->A
    crypto_x25519_scalarmult(msg3->A, client_ctx->a, P, 256);

    // Compute the two Diffie-Hellman shared secrets 
    // DH shared secret from a * B
    uint8_t a_B[32];
    crypto_x25519(a_B, client_ctx->a, msg2->B);
    // DH shared secret from v * B
    uint8_t v_B[32];
    crypto_x25519(v_B, v, msg2->B);

    // Hash everything we know so far to generate our key, save it in ctx->K_c
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx,
                              client_ctx->client_id,
                              client_ctx->client_id_len);
        crypto_blake2b_update(&hash_ctx,
                              client_ctx->server_id,
                              client_ctx->server_id_len);
        crypto_blake2b_update(&hash_ctx, msg3->A, 32);
        crypto_blake2b_update(&hash_ctx, msg2->B, 32);
        crypto_blake2b_update(&hash_ctx, a_B, 32);
        crypto_blake2b_update(&hash_ctx, v_B, 32);
        crypto_blake2b_final(&hash_ctx, client_ctx->K_c);
    }

    // Hash k and the client modifier to get our verifier, save it in msg3->client_verifier
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx, client_ctx->K_c, 32);
        crypto_blake2b_update(&hash_ctx,
                              BSSPEKE_VERIFY_CLIENT_MODIFIER,
                              BSSPEKE_VERIFY_CLIENT_MODIFIER_LEN);
        crypto_blake2b_final(&hash_ctx, msg3->client_verifier);
    }

    return 0;
}

int
bsspeke_server_generate_message4(bsspeke_msg4_t *msg4,
                                 const bsspeke_msg3_t *msg3,
                                 bsspeke_server_ctx *server_ctx
) {
    // Compute the two Diffie-Hellman shared secrets
    // DH shared secret from b * A
    uint8_t b_A[32];
    crypto_x25519(b_A, server_ctx->b, msg3->A);
    // DH shared secret from b * V
    uint8_t b_V[32];
    crypto_x25519(b_V, server_ctx->b, server_ctx->V);

    // Hash everything we've learned so far to generate k, save it in ctx->k
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx,
                              server_ctx->client_id,
                              server_ctx->client_id_len);
        crypto_blake2b_update(&hash_ctx,
                              server_ctx->server_id,
                              server_ctx->server_id_len);
        crypto_blake2b_update(&hash_ctx, msg3->A, 32);
        crypto_blake2b_update(&hash_ctx, server_ctx->B, 32);
        crypto_blake2b_update(&hash_ctx, b_A, 32);
        crypto_blake2b_update(&hash_ctx, b_V, 32);
        crypto_blake2b_final(&hash_ctx, server_ctx->K_s);
    }

    // Check that the client's hash is correct
    // Compute H( k || VERIFY_CLIENT_MODIFIER )
    uint8_t my_client_verifier[32];
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx, server_ctx->K_s, 32);
        crypto_blake2b_update(&hash_ctx,
                              BSSPEKE_VERIFY_CLIENT_MODIFIER,
                              BSSPEKE_VERIFY_CLIENT_MODIFIER_LEN);
        crypto_blake2b_final(&hash_ctx, my_client_verifier);
    }

    // Compare vs msg3->client_verifier
    if( crypto_verify32(msg3->client_verifier, my_client_verifier) != 0 ) {
        return -1;
    }

    // Compute our own verifier H( k || VERIFY_SERVER_MODIFIER ), save it in msg4->server_verifier
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx, server_ctx->K_s, 32);
        crypto_blake2b_update(&hash_ctx,
                              BSSPEKE_VERIFY_SERVER_MODIFIER,
                              BSSPEKE_VERIFY_SERVER_MODIFIER_LEN);
        crypto_blake2b_final(&hash_ctx, msg4->server_verifier);
    }

    // If we made it this far, return success
    return 0;
}

int
bsspeke_client_verify_message4(const bsspeke_msg4_t *msg4,
                               const bsspeke_client_ctx *client_ctx
) {
    // Compute our own version of the server's verifier hash
    uint8_t my_server_verifier[32];
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx, client_ctx->K_c, 32);
        crypto_blake2b_update(&hash_ctx,
                              BSSPEKE_VERIFY_SERVER_MODIFIER,
                              BSSPEKE_VERIFY_SERVER_MODIFIER_LEN);
        crypto_blake2b_final(&hash_ctx, my_server_verifier);
    }

    // If the hashes don't match, return failure
    if( crypto_verify32(msg4->server_verifier, my_server_verifier) != 0 ) {
        return -1;
    }

    // Otherwise, return success
    return 0;
}

int main(int argc, char *argv[])
{
    // Before execution of the protocol
    // Both have:
    //       idS = server identity
    char *server_id;
    // Client has:
    //       idC = client identity
    char *client_id;
    char *password;
    // Server has these for "idC":
    //       salt
    //       settings
    //       P = hashToPoint(p)
    //       V = v * P
    uint8_t salt[32];
    uint8_t P[32];   // curve point
    uint8_t V[32];   // curve point


    // Step 1: Client hashes password, maps to a point on the curve, blinds with a random value
    // C:    r = random()
    // C:    R = r * hashToPoint(H(password, idC, idS))
    // C->S: idC, R
    uint8_t r[32];   // scalar
    uint8_t R[32];   // curve point


    // Step 2: Server generates response with blind salt
    // S: b = random()
    // S: B = b * P
    // S: R' = H(salt) * R
    // C<-S: B, R', settings
    uint8_t b[32];   // scalar
    uint8_t B[32];   // curve point
    uint8_t R_prime[32];  // curve point

    // Step 3: 
    // C:    BlindSalt = (1/r) * R'
    // C:    p || v = pwKdf(password, BlindSalt, idC, idS, settings)
    // C:    P = hashToPoint(p)
    // C:    a = random()
    // C:    A = a * P
    // C:    K_c = H(idC, idS, A, B, a * B, v * B)
    // C:    verifierC = H(K_c, verifyCModifier)
    // C->S: A, verifierC[, encryptedDataC]
    uint8_t blind_salt[32];  // curve point
    uint8_t password_hash[64]; // hash
    uint8_t client_P[32];    // curve point
    uint8_t a[32];           // scalar
    uint8_t A[32];           // curve point
    uint8_t K_c[32];         // hash
    uint8_t verifierC[32];   // hash


    // Step 4:
    //    S: K_s = H(idC, idS, A, B, b * A, b * V)
    //    S: Checks verifierC == H(K_s, verifyCModifier)
    //    S: verifierS = H(K_s, verifySModifier)
    // C<-S: verifierS[, encryptedDataS]
    uint8_t K_s[32];         // hash
    uint8_t server_hash[32]; // hash -- should match verifierC

    
    // Step 5:
    // C:    Checks verifierS == H(K_c, verifySModifier)
    uint8_t client_hash[32]; // hash -- should match verifierS


    return 0;
}