/*
 * bsspeke.c - BS-SPEKE over Curve25519
 *
 * Author: Charles V. Wright <cvwright@futo.org>
 * 
 * Copyright (c) 2022 FUTO Holdings, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

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

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

typedef enum {
    DEBUG,
    INFO,
    WARN,
    ERROR,
    FATAL
} debug_level_t;

debug_level_t curr_level = DEBUG;

void debug(debug_level_t level,
           const char *msg
) {
    //if( level >= curr_level ) {
        puts(msg);
    //}
}

void print_point(const char *label,
                 uint8_t point[32]
) {
    printf("%8s:\t[", label);
    int i = 31;
    for(i=31; i >= 0; i--)
        printf("%02x", point[i]);
    printf("]\n");
}

void
generate_random_bytes(uint8_t *buf, size_t len)
{
#ifdef linux
    getrandom(buf, len, 0);
#else
    arc4random_buf(buf, len);
#endif
}

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_server_init(bsspeke_server_ctx *ctx,
                    const char* server_id, const size_t server_id_len
) {
    ctx->server_id = (uint8_t *)server_id;
    ctx->server_id_len = server_id_len;
}

void
bsspeke_client_generate_message1
(
    bsspeke_login_msg1_t *msg1,
    bsspeke_client_ctx *client
) {
    debug(DEBUG, "Hashing client's password");
    // 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 hash_ctx;
        // Give us a 256 bit (32 byte) hash; Don't use a key
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        // Add the client id, server id, and the password to the hash
        crypto_blake2b_update(&hash_ctx,
                              (const uint8_t *)(client->password),
                              client->password_len);
        crypto_blake2b_update(&hash_ctx,
                              (const uint8_t *)(client->client_id),
                              client->client_id_len);
        crypto_blake2b_update(&hash_ctx,
                              (const uint8_t *)(client->server_id),
                              client->server_id_len);
        // Write the digest value into `scalar_hash`
        crypto_blake2b_final(&hash_ctx, scalar_hash);
    }
    debug(DEBUG, "Mapping password hash onto the curve");
    // Now use Elligator to map our scalar hash to a point on the curve
    crypto_hidden_to_curve(curve_point, scalar_hash);

    print_point("H(pass)", curve_point);

    // 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...
    debug(DEBUG, "Generating random blind `r`");
    //arc4random_buf(client->r, 32);
    generate_random_bytes(client->r, 32);
    print_point("r", client->r);
    debug(DEBUG, "Clamping r");
    crypto_x25519_clamp(client->r);
    print_point("r", client->r);

    debug(DEBUG, "Multiplying curve point by r");
    // 3. Multiply our curve point by r
    crypto_x25519_scalarmult(msg1->blind, client->r, curve_point, 256);
    print_point("blind", msg1->blind);
    debug(DEBUG, "Done");

    msg1->client_id = client->client_id;
    return;
}

void
bsspeke_client_setup_generate_message1
(
    bsspeke_setup_msg1_t *msg1,
    bsspeke_client_ctx *client
) {
    bsspeke_client_generate_message1(msg1, client);
}

void
bsspeke_server_setup_generate_message2
    (
        bsspeke_setup_msg2_t *msg2,
        const bsspeke_setup_msg1_t *msg1,
        // uint8_t *salt, const size_t salt_len,
        bsspeke_user_info_t *user_info,
        bsspeke_server_ctx *server_ctx
    )
{
    // We're setting up a new account
    // So we have to create a new random salt for the user
    debug(DEBUG, "Generating new salt");
    user_info->salt_len = 32;
    //arc4random_buf(user_info->salt, user_info->salt_len);
    generate_random_bytes(user_info->salt, user_info->salt_len);
    print_point("salt", user_info->salt);

    // Hash the salt
    debug(DEBUG, "Hashing the salt");
    uint8_t H_salt[32];
    crypto_blake2b_general(H_salt, 32,
                           NULL, 0,
                           user_info->salt,
                           user_info->salt_len);
    // Use clamp() to ensure we stay on the curve in the multiply below
    crypto_x25519_clamp(H_salt);
    print_point("H_salt", H_salt);
    // Multiply H(salt) by msg1->blind, save into msg2->blind_salt
    debug(DEBUG, "Multiplying H_salt by the user's blind");
    crypto_x25519_scalarmult(msg2->blind_salt, H_salt, msg1->blind, 256);
    print_point("blndsalt", msg2->blind_salt);
}

int
bsspeke_client_setup_generate_message3
    (
        bsspeke_setup_msg3_t *msg3,
        const bsspeke_setup_msg2_t *msg2,
        uint32_t phf_blocks, uint32_t phf_iterations,
        bsspeke_client_ctx *client
    )
{
    // Sanity checks first, before we do any work
    if( phf_blocks < BSSPEKE_MIN_PHF_BLOCKS ) {
        return -1;
    }
    if( 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
    debug(DEBUG, "Removing the blind from the oblivious salt");
    crypto_x25519_inverse(oblivious_salt, client->r, msg2->blind_salt);
    print_point("obv_salt", oblivious_salt);

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

    debug(DEBUG, "Creating the salt for the PHF");
    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->client_id,
                              client->client_id_len);
        crypto_blake2b_update(&hash_ctx,
                              client->server_id,
                              client->server_id_len);
        crypto_blake2b_final(&hash_ctx, phf_salt);
    }
    print_point("phf_salt", phf_salt);

    debug(DEBUG, "Running the PHF to generate p and v");
    void *work_area;
    if ((work_area = malloc(phf_blocks * 1024)) == NULL) {
        return -1;
    }
    crypto_argon2i(password_hash, 64, work_area,
                   phf_blocks, phf_iterations,
                   client->password, client->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]);
    // FIXME Should we clamp() v before we multiply???
    crypto_x25519_clamp(v);
    print_point("p", p);
    print_point("v", v);

    // Hash p onto the curve to get this user's base point P
    //uint8_t P[32];
    debug(DEBUG, "Hashing p onto the curve to get P");
    crypto_hidden_to_curve(msg3->P, p);
    print_point("P", msg3->P);

    // Generate our long-term public key V = v * P
    debug(DEBUG, "V = v * P");
    crypto_x25519_scalarmult(msg3->V, v, msg3->P, 256);
    print_point("V", msg3->V);

    // Send our PHF params so the server can store them for us
    msg3->phf_blocks = phf_blocks;
    msg3->phf_iterations = phf_iterations;

    return 0;
}

void
bsspeke_server_setup_process_message3
    (
        bsspeke_setup_msg3_t *msg3,
        bsspeke_user_info_t *user_info,
        bsspeke_server_ctx *server
    )
{
    // Setup Message 3 contains stuff that we need to remember.
    // The user info struct is where we store stuff that
    // needs to persist across sessions.
    memcpy(user_info->P, msg3->P, 32);
    memcpy(user_info->V, msg3->V, 32);
    user_info->phf_blocks = msg3->phf_blocks;
    user_info->phf_iterations = msg3->phf_iterations;
}

void
bsspeke_client_login_generate_message1
(
    bsspeke_login_msg1_t *msg1,
    bsspeke_client_ctx *client
) {
    bsspeke_client_generate_message1(msg1, client);
}

void
bsspeke_server_login_generate_message2(bsspeke_login_msg2_t *msg2,
                                 const bsspeke_login_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,
                                 const bsspeke_user_info_t *user_info,
                                 bsspeke_server_ctx *server
) {
    // Record the client's id
    server->client_id_len = strlen(msg1->client_id);
    server->client_id = (uint8_t *)malloc(server->client_id_len + 1);
    strncpy(server->client_id, msg1->client_id, server->client_id_len);
    printf("Server got msg1 from [%s]\n", (char *)(server->client_id));

    // Hash the salt
    debug(DEBUG, "Hashing the salt");
    uint8_t H_salt[32];
    crypto_blake2b_general(H_salt, 32,
                           NULL, 0,
                           user_info->salt,
                           user_info->salt_len);
    // Use clamp() to ensure we stay on the curve in the multiply below
    crypto_x25519_clamp(H_salt);
    print_point("H_salt", H_salt);
    // Multiply H(salt) by msg1->blind, save into msg2->blind_salt
    debug(DEBUG, "Multiplying H_salt by the user's blind");
    crypto_x25519_scalarmult(msg2->blind_salt, H_salt, msg1->blind, 256);
    print_point("blndsalt", msg2->blind_salt);

    // Generate random ephemeral private key b, save it in ctx->b
    debug(DEBUG, "Generating ephemeral private key b");
    //arc4random_buf(server->b, 32);
    generate_random_bytes(server->b, 32);
    crypto_x25519_clamp(server->b);
    print_point("b", server->b);

    debug(DEBUG, "Using user's base point P");
    print_point("P", user_info->P);

    // Compute public key B = b * P, save it in msg2->B
    debug(DEBUG, "Computing ephemeral public key B = b * P");
    crypto_x25519_scalarmult(server->B, server->b, user_info->P, 256);
    print_point("B", server->B);

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

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

    return;
}


int
bsspeke_client_login_generate_message3(bsspeke_login_msg3_t *msg3,
                                 const bsspeke_login_msg2_t *msg2,
                                 bsspeke_client_ctx *client
) {
    // Sanity checks first, before we do any work
    if( msg2->phf_blocks < BSSPEKE_MIN_PHF_BLOCKS ) {
        debug(ERROR, "phf_blocks is too low.  Aborting.");
        return -1;
    }
    if( msg2->phf_iterations < BSSPEKE_MIN_PHF_ITERATIONS ) {
        debug(ERROR, "phf_iterations is too low.  Aborting.");
        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
    debug(DEBUG, "Removing the blind from the oblivious salt");
    crypto_x25519_inverse(oblivious_salt, client->r, msg2->blind_salt);
    print_point("obv_salt", oblivious_salt);

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

    debug(DEBUG, "Creating the salt for the PHF");
    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->client_id,
                              client->client_id_len);
        crypto_blake2b_update(&hash_ctx,
                              client->server_id,
                              client->server_id_len);
        crypto_blake2b_final(&hash_ctx, phf_salt);
    }
    print_point("phf_salt", phf_salt);

    debug(DEBUG, "Running the PHF to generate p and v");
    void *work_area;
    if ((work_area = malloc(msg2->phf_blocks * 1024)) == NULL) {
        debug(ERROR, "Failed to malloc() memory for PHF work area");
        return -1;
    }
    crypto_argon2i(password_hash, 64, work_area,
                   msg2->phf_blocks, msg2->phf_iterations,
                   client->password, client->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]);
    // In the setup protocol, we clamped v before using it
    // So we should do the same here
    crypto_x25519_clamp(v);
    print_point("p", p);
    print_point("v", v);

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

    // Generate a random ephemeral private key a, store it in ctx->a
    debug(DEBUG, "Generating ephemeral private key a");
    //arc4random_buf(client->a, 32);
    generate_random_bytes(client->a, 32);
    crypto_x25519_clamp(client->a);
    print_point("a", client->a);
    // Generate the ephemeral public key A = a * P, store it in msg3->A
    debug(DEBUG, "Generating ephemeral public key A = a * P");
    crypto_x25519_scalarmult(msg3->A, client->a, P, 256);
    print_point("A", msg3->A);

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

    // Hash everything we know so far to generate our key, save it in ctx->K_c
    debug(DEBUG, "Hashing current state to get key K_c");
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx,
                              client->client_id,
                              client->client_id_len);
        crypto_blake2b_update(&hash_ctx,
                              client->server_id,
                              client->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->K_c);
    }
    print_point("K_c", client->K_c);

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

    return 0;
}

int
bsspeke_server_login_generate_message4(bsspeke_login_msg4_t *msg4,
                                 const bsspeke_login_msg3_t *msg3,
                                 const bsspeke_user_info_t *user_info,
                                 bsspeke_server_ctx *server
) {
    // Compute the two Diffie-Hellman shared secrets
    debug(DEBUG, "Computing Diffie-Hellman shared secrets");
    // DH shared secret from b * A
    uint8_t b_A[32];
    crypto_x25519(b_A, server->b, msg3->A);
    print_point("b * A", b_A);
    // DH shared secret from b * V
    //print_point("V", user_info->V);
    uint8_t b_V[32];
    crypto_x25519(b_V, server->b, user_info->V);
    print_point("b * V", b_V);

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

    // Check that the client's hash is correct
    // Compute H( k || VERIFY_CLIENT_MODIFIER )
    debug(DEBUG, "Checking client's verifier hash");
    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->K_s, 32);
        crypto_blake2b_update(&hash_ctx,
                              (uint8_t *)BSSPEKE_VERIFY_CLIENT_MODIFIER,
                              BSSPEKE_VERIFY_CLIENT_MODIFIER_LEN);
        crypto_blake2b_final(&hash_ctx, my_client_verifier);
    }
    print_point("client's", msg3->client_verifier);
    print_point("mine", my_client_verifier);

    // Compare vs msg3->client_verifier
    if( crypto_verify32(msg3->client_verifier, my_client_verifier) != 0 ) {
        debug(ERROR, "Client's verifier doesn't match!");
        return -1;
    }
    debug(DEBUG, "Client's verifier checks out");

    // Compute our own verifier H( k || VERIFY_SERVER_MODIFIER ), save it in msg4->server_verifier
    debug(DEBUG, "Computing server verifier hash");
    {
        crypto_blake2b_ctx hash_ctx;
        crypto_blake2b_general_init(&hash_ctx, 32, NULL, 0);
        crypto_blake2b_update(&hash_ctx, server->K_s, 32);
        crypto_blake2b_update(&hash_ctx,
                              (uint8_t *)BSSPEKE_VERIFY_SERVER_MODIFIER,
                              BSSPEKE_VERIFY_SERVER_MODIFIER_LEN);
        crypto_blake2b_final(&hash_ctx, msg4->server_verifier);
    }
    print_point("server_v", msg4->server_verifier);

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

int
bsspeke_client_login_verify_message4(const bsspeke_login_msg4_t *msg4,
                               const bsspeke_client_ctx *client
) {
    // Compute our own version of the server's verifier hash
    debug(DEBUG, "Verifying hash from the server");
    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->K_c, 32);
        crypto_blake2b_update(&hash_ctx,
                              (uint8_t *)BSSPEKE_VERIFY_SERVER_MODIFIER,
                              BSSPEKE_VERIFY_SERVER_MODIFIER_LEN);
        crypto_blake2b_final(&hash_ctx, my_server_verifier);
    }
    print_point("mine", my_server_verifier);
    print_point("server's", msg4->server_verifier);

    // If the hashes don't match, return failure
    if( crypto_verify32(msg4->server_verifier, my_server_verifier) != 0 ) {
        debug(WARN, "Server's hash doesn't match.  Aborting.");
        return -1;
    }
    debug(DEBUG, "Server's hash checks out.  SUCCESS!");

    // Otherwise, return success
    return 0;
}