/*
 * 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 "minimonocypher.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,
                    const char* client_id, const size_t client_id_len
) {
    ctx->server_id = (uint8_t *)server_id;
    ctx->server_id_len = server_id_len;

    ctx->client_id = (uint8_t *)client_id;
    ctx->client_id_len = client_id_len;
}

void
bsspeke_client_generate_blind
(
    uint8_t blind[32],
    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(blind, client->r, curve_point, 256);
    print_point("blind", blind);
    debug(DEBUG, "Done");

    return;
}

void
bsspeke_server_blind_salt
    (
        uint8_t blind_salt[32],
        const uint8_t blind[32],
        const uint8_t *salt, const size_t salt_len
    )
{
    /*
    // FIXME This function no longer generates its own salt!
    //       We MUST do this before calling the function!!!
    // We're setting up a new account
    // So we have to create a new random salt for the user
    debug(DEBUG, "Generating new salt");
    salt_len = 32;
    generate_random_bytes(user_info->salt, user_info->salt_len);
    */
    print_point("salt", salt);

    // Hash the salt
    debug(DEBUG, "Hashing 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);
    print_point("H_salt", H_salt);
    // Multiply H(salt) by blind, save into blind_salt
    debug(DEBUG, "Multiplying H_salt by the user's blind");
    crypto_x25519_scalarmult(blind_salt, H_salt, blind, 256);
    print_point("blndsalt", blind_salt);
}

void
bsspeke_server_generate_B(const uint8_t P[32],
                          bsspeke_server_ctx *server)
{
    // Generate random ephemeral private key b, save it in server->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", 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, P, 256);
    print_point("B", server->B);

}

int
bsspeke_client_generate_keys_from_password(const uint8_t blind_salt[32],
                                           uint32_t phf_blocks, uint32_t phf_iterations,
                                           bsspeke_client_ctx *client)
{
    // Sanity checks first, before we do any work
    if( phf_blocks < BSSPEKE_ARGON2_MIN_PHF_BLOCKS ) {
        debug(ERROR, "Requested PHF blocks below the minimum");
        return -1;
    }
    if( phf_iterations < BSSPEKE_ARGON2_MIN_PHF_ITERATIONS ) {
        debug(ERROR, "Requested PHF iterations below the minimum");
        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, blind_salt);
    print_point("obv_salt", oblivious_salt);

    // Hash the oblivious salt together with the id's to create the salt for the PHF
    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");
    uint8_t password_hash[64];
    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 *tmp_p = &(password_hash[0]);
    uint8_t *tmp_v = &(password_hash[32]);
    // clamp() v before we do anything else with it
    crypto_x25519_clamp(tmp_v);

    memcpy(client->p, tmp_p, 32);
    memcpy(client->v, tmp_v, 32);
    crypto_wipe(password_hash, 64);

    print_point("p", client->p);
    print_point("v", client->v);

    return 0;
}
        

int
bsspeke_client_generate_P_and_V
    (
        uint8_t P[32], uint8_t V[32],
        const uint8_t blind_salt[32],
        uint32_t phf_blocks, uint32_t phf_iterations,
        bsspeke_client_ctx *client
    )
{
    int rc = bsspeke_client_generate_keys_from_password(blind_salt,
                                                        phf_blocks, phf_iterations,
                                                        client);
    if( rc != 0 ) {
        debug(ERROR, "Password hashing function failed");
        return -1;
    }

    // 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, client->p);
    print_point("P", P);

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

    return 0;
}


int
bsspeke_client_generate_A(const uint8_t blind_salt[32],
                          uint32_t phf_blocks, uint32_t phf_iterations,
                          bsspeke_client_ctx *client
) {
    int rc = bsspeke_client_generate_keys_from_password(blind_salt,
                                                        phf_blocks, phf_iterations,
                                                        client);
    if( rc != 0 ) {
        debug(ERROR, "Password hashing failed");
        return -1;
    }

    // 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, client->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(client->A, client->a, P, 256);
    print_point("A", client->A);

    return 0;
}

void
bsspeke_client_derive_shared_key(const uint8_t B[32],
                                 bsspeke_client_ctx *client)
{
    // 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, B);
    print_point("a * B", a_B);
    // DH shared secret from v * B
    uint8_t v_B[32];
    crypto_x25519(v_B, client->v, 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, client->A, 32);
        crypto_blake2b_update(&hash_ctx, 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);
}

void
bsspeke_client_generate_verifier(uint8_t client_verifier[32],
                                 bsspeke_client_ctx *client)
{
    // 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, client_verifier);
    }
    print_point("client_v", client_verifier);
}

void
bsspeke_server_derive_shared_key(const uint8_t A[32],
                                 const uint8_t V[32],
                                 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, A);
    print_point("b * A", b_A);
    // DH shared secret from b * V
    uint8_t b_V[32];
    crypto_x25519(b_V, server->b, 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, 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);
}

int
bsspeke_server_verify_client(uint8_t client_verifier[32],
                             bsspeke_server_ctx *server
) {

    // 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", client_verifier);
    print_point("mine", my_client_verifier);

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

    return 0;
}

void
bsspeke_server_generate_verifier(uint8_t server_verifier[32],
                                 bsspeke_server_ctx *server
) {
    // Compute our own verifier H( k || VERIFY_SERVER_MODIFIER ), save it in 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, server_verifier);
    }
    print_point("server_v", server_verifier);
}

int
bsspeke_client_verify_server(const uint8_t server_verifier[32],
                             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", server_verifier);

    // If the hashes don't match, return failure
    if( crypto_verify32(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;
}