ns.cc

#include <ns.h>

static int generate_key(int key_size, const std::string pkey, const std::string ppkey)
{
    RSA *rsa = nullptr;
    BIGNUM *b;
    FILE *fp;
    int ret;

    rsa = RSA_new();

    b = BN_new();
    BN_set_word(b, RSA_F4);

    ret = RSA_generate_key_ex(rsa, key_size, b, nullptr);
    if (ret != 1) {
        printf("RSA: failed to generate key @ %s %u\n",
                            __func__, __LINE__);
        return -1;
    }

    fp = fopen(ppkey.c_str(), "w");
    if (!fp) {
        printf("RSA: failed to open private key file to write @ %s %u\n",
                            __func__, __LINE__);
        return -1;
    }

    ret = PEM_write_RSAPrivateKey(fp, rsa, nullptr, nullptr, 0, nullptr, nullptr);
    if (ret != 1) {
        printf("RSA: failed to write private key @ %s %u\n",
                            __func__, __LINE__);
        fclose(fp);
        return -1;
    }

    fflush(fp);
    fclose(fp);

    fp = fopen(pkey.c_str(), "w");
    if (!fp) {
        printf("RSA: failed to open public key file to write @ %s %u\n",
                            __func__, __LINE__);
        return -1;
    }

    ret = PEM_write_RSA_PUBKEY(fp, rsa);
    if (ret != 1) {
        printf("RSA: failed to write RSA public key @ %s %u\n",
                            __func__, __LINE__);
        fclose(fp);
        return -1;
    }

    BN_free(b);
    RSA_free(rsa);
    fflush(fp);
    fclose(fp);

    return 0;
}

static int sign_msg(const std::string priv_key,
                    const uint8_t *msg,
                    size_t msg_len,
                    uint8_t *signature,
                    size_t &signature_len)
{
    const EVP_MD *md;
    EVP_MD_CTX *ctx;
    EVP_PKEY *key;
    FILE *fp;
    int ret;

    md = EVP_sha256();

    fp = fopen(priv_key.c_str(), "r");
    if (!fp) {
        printf("RSA: failed to open rsa private key\n");
        return -1;
    }

    key = PEM_read_PrivateKey(fp, nullptr, nullptr, nullptr);
    if (!key) {
        ERR_print_errors_fp(stderr);
        printf("RSA: failed to load RSA private key\n");
        goto err;
    }

    ctx = EVP_MD_CTX_create();
    if (!ctx) {
        printf("RSA: failed to create md context\n");
        goto err;
    }

    ret = EVP_DigestInit_ex(ctx, md, nullptr);
    if (ret != 1) {
        printf("RSA: failed to initialize digest context\n");
        goto err;
    }

    ret = EVP_DigestSignInit(ctx, nullptr, md, nullptr, key);
    if (ret != 1) {
        printf("RSA: failed to initialize digest for signing\n");
        goto err;
    }

    ret = EVP_SignUpdate(ctx, msg, msg_len);
    if (ret != 1) {
        printf("RSA: failed to update digest signature\n");
        goto err;
    }

    ret = EVP_DigestSignFinal(ctx, NULL, &signature_len);
    if (ret != 1) {
        goto err;
    }

    ret = EVP_DigestSignFinal(ctx, signature, &signature_len);
    if (ret != 1) {
        printf("RSA: failed to finalize the digest\n");
        goto err;
    }

    EVP_PKEY_free(key);
    EVP_MD_CTX_free(ctx);
    fclose(fp);

    return 0;

err:
    if (fp) {
        fclose(fp);
    }
    EVP_MD_CTX_free(ctx);
    EVP_PKEY_free(key);

    return -1;
}

static int verify_msg(std::string pub_key,
                      const uint8_t *msg,
                      size_t msg_len,
                      uint8_t *signature,
                      size_t &signature_len)
{
    const EVP_MD *md;
    EVP_MD_CTX *ctx = nullptr;
    EVP_PKEY *key;
    FILE *fp;
    int ret;

    md = EVP_sha256();

    fp = fopen(pub_key.c_str(), "r");
    if (!fp) {
        printf("RSA: failed to open publik key %s\n", pub_key.c_str());
        return -1;
    }

    key = PEM_read_PUBKEY(fp, nullptr, nullptr, nullptr);
    if (!key) {
        printf("RSA: failed to load public key to memory\n");
        goto err;
    }

    ctx = EVP_MD_CTX_create();
    if (!ctx) {
        printf("RSA: failed to create digest context\n");
        goto err;
    }

    ret = EVP_DigestInit_ex(ctx, md, nullptr);
    if (ret != 1) {
        printf("RSA: failed to initialize digest\n");
        goto err;
    }

    ret = EVP_DigestVerifyInit(ctx, nullptr, md, nullptr, key);
    if (ret != 1) {
        goto err;
    }

    ret = EVP_VerifyUpdate(ctx, msg, msg_len);
    if (ret != 1) {
        goto err;
    }

    ret = EVP_DigestVerifyFinal(ctx, signature, signature_len);
    if (ret != 1) {
        goto err;
    }

    if (fp) {
        fclose(fp);
    }
    EVP_PKEY_free(key);
    EVP_MD_CTX_free(ctx);

    return 0;

err:
    if (fp) {
        fclose(fp);
    }
    EVP_MD_CTX_destroy(ctx);
    EVP_PKEY_free(key);

    return -1;
}

static int public_encrypt(const std::string pub_key, const uint8_t *msg, size_t msg_len, uint8_t *encbuf, size_t &encbuf_size)
{
    RSA *rsa = nullptr;
    FILE *fp;
    int ret;

    fp = fopen(pub_key.c_str(), "r");
    if (!fp) {
        return -1;
    }

    PEM_read_RSA_PUBKEY(fp, &rsa, nullptr, nullptr);

    ret = RSA_public_encrypt(msg_len, msg, encbuf, rsa, RSA_PKCS1_PADDING);
    if (ret < 0) {
        return -1;
    }

    encbuf_size = ret;

    RSA_free(rsa);
    fclose(fp);
    return 0;
}

static int private_decrypt(const std::string priv_key, const uint8_t *msg, size_t msg_len, uint8_t *decbuf, size_t &decbuf_size)
{
    RSA *rsa = nullptr;
    FILE *fp;
    int ret;

    fp = fopen(priv_key.c_str(), "r");
    if (!fp) {
        return -1;
    }

    PEM_read_RSAPrivateKey(fp, &rsa, nullptr, nullptr);

    ret = RSA_private_decrypt(msg_len, msg, decbuf, rsa, RSA_PKCS1_PADDING);
    if (ret < 0) {
        return -1;
    }

    decbuf_size = ret;

    RSA_free(rsa);
    fclose(fp);
    return 0;
}


static std::string alice_privkey = "./alice_priv_key.pem";
static std::string alice_pubkey = "./alice_pub_key.pem";

static std::string bob_privkey = "./bob_priv_key.pem";
static std::string bob_pubkey = "./bob_pub_key.pem";

static std::string server_privkey = "./server_priv_key.pem";
static std::string server_pubkey = "./server_pub_key.pem";

static uint32_t alice_uid = 0x12345678;
static uint32_t bob_uid = 0x12345688;


static uint8_t nonce_a[16];
static uint8_t nonce_b[16];

static uint8_t msg[4096];
static uint32_t msg_len = 0;

static int alice_rsa_keypair_create()
{
    return generate_key(2048, alice_pubkey, alice_privkey);
}

static int server_rsa_keypair_create()
{
    return generate_key(2048, server_pubkey, server_privkey);
}

static int bob_rsa_keypair_create()
{
    return generate_key(2048, bob_pubkey, bob_privkey);
}

static int get_random(uint8_t *data, size_t data_len)
{
    int fd;
    int ret;

    fd = open("/dev/urandom", O_RDONLY);
    if (fd < 0) {
        return -1;
    }

    ret = read(fd, data, data_len);
    if (ret < 0) {
        return -1;
    }

    close(fd);
    return 0;
}

static int ns_setup()
{
    int ret;

    ret = alice_rsa_keypair_create();
    if (ret < 0) {
        return -1;
    }

    ret = get_random(nonce_a, sizeof(nonce_a));
    if (ret < 0) {
        return -1;
    }

    ret = bob_rsa_keypair_create();
    if (ret < 0) {
        return -1;
    }

    ret = get_random(nonce_b, sizeof(nonce_b));
    if (ret < 0) {
        return -1;
    }

    ret = server_rsa_keypair_create();
    if (ret < 0) {
        return -1;
    }

    return 0;
}

static int step1_alice_initiate_communication()
{
    int off = 0;

    memcpy(msg, &alice_uid, sizeof(alice_uid));
    off += sizeof(alice_uid);

    memcpy(msg + off, &bob_uid, sizeof(bob_uid));
    off += sizeof(bob_uid);

    printf("alice:\t send uid 0x%x bob uid 0x%x to server A->S: A, B\n", alice_uid, bob_uid);
    return off;
}

static int step2_server_response_to_alice(uint8_t *signature, size_t &signature_len)
{
    int off = 0;
    int fd;
    int ret;

    printf("server:\t read bob's public key\n");
    fd = open(bob_pubkey.c_str(), O_RDONLY);
    if (fd < 0) {
        return -1;
    }

    ret = read(fd, msg, sizeof(msg));
    if (ret < 0) {
        return -1;
    }

    off = ret;

    printf("server:\t copy bob's key\n");

    memcpy(msg + off, &bob_uid, sizeof(bob_uid));
    off += sizeof(bob_uid);

    ret = sign_msg(server_privkey, msg, off, signature, signature_len);
    if (ret < 0) {
        printf("server: failed to sign message with server private %s\n", server_privkey.c_str());
        return -1;
    }

    printf("server:\t send response to A, S -> A {KPB, B}KSS\n");

    msg_len = off;

    return 0;
}

static int step3_alice_send_msg_to_bob(uint8_t *server_signature, size_t ss_len,
                                       uint8_t *enc_msg, size_t &enc_msg_len)
{
    int off = 0;
    int ret;


    ret = verify_msg(server_pubkey, msg, msg_len, server_signature, ss_len);
    if (ret < 0) {
        printf("alice:\t failed to verify message\n");
        return -1;
    }

    printf("alice:\t server signature verified.. extract bob's key\n");

    memset(msg, 0, sizeof(msg));

    memcpy(msg, nonce_a, sizeof(nonce_a));
    off += sizeof(nonce_a);

    memcpy(msg + off, &alice_uid, sizeof(alice_uid));
    off += sizeof(alice_uid);

    msg_len = off;

    printf("alice:\t set nonce and alice uid\n");

    ret = public_encrypt(bob_pubkey, msg, msg_len, enc_msg, enc_msg_len);
    if (ret < 0) {
        printf("alice:\t failed to encrypt with %s\n", bob_pubkey.c_str());
        return -1;
    }

    printf("alice:\t encrypted ok %lu A -> B: {NA, A}KPB\n", enc_msg_len);

    return 0;
}

static int step4_bob_request_server_about_alice(uint8_t *bob_msg, size_t bob_msg_len)
{
    int off = 0;

    memcpy(bob_msg, &bob_uid, sizeof(bob_uid));
    off += 4;

    memcpy(bob_msg, &alice_uid, sizeof(alice_uid));
    off += 4;

    printf("bob:\t set bob_uid 0x%x alice_uid 0x%x B->S: {B, A}\n", bob_uid, alice_uid);

    return 0;
}

static int step5_server_responds_to_bob_about_alice(uint8_t *bob_msg, size_t &bob_msg_len,
                                                    uint8_t *signature, size_t &signature_len)
{
    int off = 0;
    int fd;
    int ret;

    printf("server:\t load alice public key\n");
    fd = open(alice_pubkey.c_str(), O_RDONLY);
    if (fd < 0) {
        return -1;
    }

    printf("server:\t copy alice public key\n");
    ret = read(fd, bob_msg, 1024);
    if (ret < 0) {
        return -1;
    }

    close(fd);

    off = ret;

    printf("server:\t copy alice uid\n");
    memcpy(bob_msg + off, &alice_uid, sizeof(alice_uid));
    off += sizeof(alice_uid);

    bob_msg_len = off;

    ret = sign_msg(server_privkey, bob_msg, bob_msg_len,
                   signature, signature_len);
    if (ret < 0) {
        printf("server:\t failed to sign msg\n");
        return -1;
    }

    printf("server:\t signed ok S -> B: {KPA, A}KSS\n");
    return 0;
}

int step6_bob_responds_to_alice(uint8_t *server_msg, size_t server_msg_len,
                                uint8_t *signature, size_t signature_len,
                                uint8_t *encrypted_msg, size_t &enc_msg_len)
{
    int off = 0;
    uint8_t dec_msg[1024];
    size_t dec_msg_len = 0;
    uint8_t a_nonce[16];
    int ret;

    ret = verify_msg(server_pubkey, server_msg, server_msg_len, signature, signature_len);
    if (ret < 0) {
        printf("bob:\t failed to verify message\n");
        return -1;
    }

    printf("bob:\t verified server message ok\n");

    printf("bob:\t try to verify alice len %d\n", msg_len);

    ret = private_decrypt(bob_privkey, encrypted_msg, enc_msg_len, dec_msg, dec_msg_len);
    if (ret < 0) {
        printf("bob:\t failed to decrypt msg\n");
        return -1;
    }

    memcpy(a_nonce, msg, sizeof(a_nonce));

    memset(msg, 0, sizeof(msg));
    memcpy(msg, nonce_a, sizeof(nonce_a));
    off += sizeof(nonce_a);

    memcpy(msg + off, nonce_b, sizeof(nonce_b));
    off += sizeof(nonce_b);

    msg_len = off;
    ret = public_encrypt(alice_pubkey, msg, msg_len, encrypted_msg, enc_msg_len);
    if (ret < 0) {
        printf("bob:\t failed to encrypt\n");
        return -1;
    }

    printf("bob:\t encrypt {NA,NB}KPA\n");
    return 0;
}

static int step7_alice_responds_to_bob(uint8_t *bob_msg, size_t bob_msg_len,
                                       uint8_t *enc_msg, size_t &enc_msg_len)
{
    uint8_t dec_msg[1024];
    size_t dec_msg_len = 0;
    uint8_t a_nonce[16];
    size_t a_nonce_len = 0;
    uint8_t b_nonce[16];
    size_t b_nonce_len = 0;
    int off = 0;
    int ret;

    ret = private_decrypt(alice_privkey, bob_msg, bob_msg_len, dec_msg, dec_msg_len);
    if (ret < 0) {
        printf("alice:\t failed to decrypt bob's response\n");
        return -1;
    }

    printf("alice:\t decrypt bob's msg ok\n");

    memcpy(a_nonce, dec_msg, sizeof(a_nonce));
    a_nonce_len = sizeof(a_nonce);
    off += a_nonce_len;

    memcpy(b_nonce, dec_msg + off, sizeof(b_nonce));
    b_nonce_len = sizeof(b_nonce);
    off += b_nonce_len;

    if (memcmp(a_nonce, nonce_a, sizeof(nonce_a)) != 0) {
        printf("alice:\t failed to verify a_nonce\n");
        return -1;
    }

    printf("alice:\t **** bob is verified ok **** \n");
    if (memcmp(b_nonce, nonce_b, sizeof(nonce_b)) != 0) {
        printf("alice:\t failed to verify b_nonce\n");
        return -1;
    }

    printf("alice:\t b_nonce verified ok\n");

    memset(msg, 0, sizeof(msg));

    memcpy(msg, b_nonce, b_nonce_len);
    off = sizeof(b_nonce);

    msg_len = off;

    ret = public_encrypt(bob_pubkey, msg, msg_len, enc_msg, enc_msg_len);
    if (ret < 0) {
        printf("alice:\t failed to encrypt with bob's public key\n");
        return -1;
    }

    printf("alice:\t encrypt ok A -> B {NB}KPB\n");
    return 0;
}

static int step8_bob_verify_alice_msg(uint8_t *enc_msg, size_t enc_msg_len)
{
    uint8_t dec_msg[1024];
    size_t dec_msg_len = 0;
    uint8_t b_nonce[16];
    int ret;

    ret = private_decrypt(bob_privkey, enc_msg, enc_msg_len, dec_msg, dec_msg_len);
    if (ret < 0) {
        printf("bob:\t failed to decrypt alice message\n");
        return -1;
    }

    printf("bob:\t decrypt alice msg ok\n");
    memcpy(b_nonce, dec_msg, sizeof(b_nonce));

    if (memcmp(b_nonce, nonce_b, sizeof(b_nonce)) != 0) {
        printf("bob:\t failed to verify b_nonce from received alice message\n");
        return -1;
    }

    printf("bob:\t **** alice is verified ok ****\n");
    return 0;
}

int main()
{
    uint8_t signature[4096];
    size_t signature_len = 0;
    uint8_t encrypted_msg[4096];
    size_t enc_msg_len = 0;
    uint8_t tmp[1024];
    size_t tmp_len = 0;
    int ret;

    printf("main:\t setup communication\n\n");
    ret = ns_setup();
    if (ret < 0) {
        return -1;
    }

    printf("step1:\n");
    printf("======\n");
    printf("alice:\t initiate communication\n");
    ret = step1_alice_initiate_communication();
    if (ret < 0) {
        return -1;
    }

    printf("\n\n");
    printf("step2:\n");
    printf("======\n");
    printf("server:\t prepare response\n");
    ret = step2_server_response_to_alice(signature, signature_len);
    if (ret < 0) {
        return -1;
    }

    printf("\n\n");
    printf("step3:\n");
    printf("=====\n");
    printf("alice:\t send first message to bob\n");
    ret = step3_alice_send_msg_to_bob(signature, signature_len,
                                      encrypted_msg, enc_msg_len);
    if (ret < 0) {
        return -1;
    }

    printf("\n\n");
    printf("step4:\n");
    printf("=====\n");
    printf("bob:\t received alice request\n");
    ret = step4_bob_request_server_about_alice(tmp, tmp_len);
    if (ret < 0) {
        return -1;
    }

    printf("bob:\t request server about alice identity\n");

    printf("\n\n");
    printf("step5:\n");
    printf("=====\n");

    ret = step5_server_responds_to_bob_about_alice(tmp, tmp_len,
                                                   signature, signature_len);
    if (ret < 0) {
        return -1;
    }

    printf("server:\t respond to bob with KPA, A signed with server priv key\n");

    printf("\n\n");
    printf("step6:\n");
    printf("=====\n");

    ret = step6_bob_responds_to_alice(tmp, tmp_len,
                                      signature, signature_len,
                                      encrypted_msg, enc_msg_len);
    if (ret < 0) {
        return -1;
    }

    printf("bob:\t sends {NA, NB}KPA to alice\n");

    printf("\n\n");
    printf("step7:\n");
    printf("=====\n");

    ret = step7_alice_responds_to_bob(encrypted_msg, enc_msg_len,
                                      tmp, tmp_len);
    if (ret < 0) {
        return -1;
    }

    printf("alice:\t sends {NB}KPB to bob\n");

    printf("\n\n");
    printf("step8:\n");
    printf("=====\n");

    ret = step8_bob_verify_alice_msg(tmp, tmp_len);
    if (ret < 0) {
        return -1;
    }

    printf("main:\t mutual verification complete\n");
    return 0;
}