/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "mpi.h"
#include "mpi-priv.h"

/* #define OLD_WAY 1  */

/* This key is the 1024-bit test key used for speed testing of RSA private
** key ops.
*/

#define CONST const

static CONST unsigned char default_n[128] = {
    0xc2, 0xae, 0x96, 0x89, 0xaf, 0xce, 0xd0, 0x7b, 0x3b, 0x35, 0xfd, 0x0f, 0xb1, 0xf4, 0x7a, 0xd1,
    0x3c, 0x7d, 0xb5, 0x86, 0xf2, 0x68, 0x36, 0xc9, 0x97, 0xe6, 0x82, 0x94, 0x86, 0xaa, 0x05, 0x39,
    0xec, 0x11, 0x51, 0xcc, 0x5c, 0xa1, 0x59, 0xba, 0x29, 0x18, 0xf3, 0x28, 0xf1, 0x9d, 0xe3, 0xae,
    0x96, 0x5d, 0x6d, 0x87, 0x73, 0xf6, 0xf6, 0x1f, 0xd0, 0x2d, 0xfb, 0x2f, 0x7a, 0x13, 0x7f, 0xc8,
    0x0c, 0x7a, 0xe9, 0x85, 0xfb, 0xce, 0x74, 0x86, 0xf8, 0xef, 0x2f, 0x85, 0x37, 0x73, 0x0f, 0x62,
    0x4e, 0x93, 0x17, 0xb7, 0x7e, 0x84, 0x9a, 0x94, 0x11, 0x05, 0xca, 0x0d, 0x31, 0x4b, 0x2a, 0xc8,
    0xdf, 0xfe, 0xe9, 0x0c, 0x13, 0xc7, 0xf2, 0xad, 0x19, 0x64, 0x28, 0x3c, 0xb5, 0x6a, 0xc8, 0x4b,
    0x79, 0xea, 0x7c, 0xce, 0x75, 0x92, 0x45, 0x3e, 0xa3, 0x9d, 0x64, 0x6f, 0x04, 0x69, 0x19, 0x17
};

static CONST unsigned char default_d[128] = {
    0x13, 0xcb, 0xbc, 0xf2, 0xf3, 0x35, 0x8c, 0x6d, 0x7b, 0x6f, 0xd9, 0xf3, 0xa6, 0x9c, 0xbd, 0x80,
    0x59, 0x2e, 0x4f, 0x2f, 0x11, 0xa7, 0x17, 0x2b, 0x18, 0x8f, 0x0f, 0xe8, 0x1a, 0x69, 0x5f, 0x6e,
    0xac, 0x5a, 0x76, 0x7e, 0xd9, 0x4c, 0x6e, 0xdb, 0x47, 0x22, 0x8a, 0x57, 0x37, 0x7a, 0x5e, 0x94,
    0x7a, 0x25, 0xb5, 0xe5, 0x78, 0x1d, 0x3c, 0x99, 0xaf, 0x89, 0x7d, 0x69, 0x2e, 0x78, 0x9d, 0x1d,
    0x84, 0xc8, 0xc1, 0xd7, 0x1a, 0xb2, 0x6d, 0x2d, 0x8a, 0xd9, 0xab, 0x6b, 0xce, 0xae, 0xb0, 0xa0,
    0x58, 0x55, 0xad, 0x5c, 0x40, 0x8a, 0xd6, 0x96, 0x08, 0x8a, 0xe8, 0x63, 0xe6, 0x3d, 0x6c, 0x20,
    0x49, 0xc7, 0xaf, 0x0f, 0x25, 0x73, 0xd3, 0x69, 0x43, 0x3b, 0xf2, 0x32, 0xf8, 0x3d, 0x5e, 0xee,
    0x7a, 0xca, 0xd6, 0x94, 0x55, 0xe5, 0xbd, 0x25, 0x34, 0x8d, 0x63, 0x40, 0xb5, 0x8a, 0xc3, 0x01
};

#define DEFAULT_ITERS 50

typedef clock_t timetype;
#define gettime(x) *(x) = clock()
#define subtime(a, b) a -= b
#define msec(x) ((clock_t)((double)x * 1000.0 / CLOCKS_PER_SEC))
#define sec(x) (x / CLOCKS_PER_SEC)

struct TimingContextStr {
    timetype start;
    timetype end;
    timetype interval;

    int minutes;
    int seconds;
    int millisecs;
};

typedef struct TimingContextStr TimingContext;

TimingContext *
CreateTimingContext(void)
{
    return (TimingContext *)malloc(sizeof(TimingContext));
}

void
DestroyTimingContext(TimingContext *ctx)
{
    free(ctx);
}

void
TimingBegin(TimingContext *ctx)
{
    gettime(&ctx->start);
}

static void
timingUpdate(TimingContext *ctx)
{

    ctx->millisecs = msec(ctx->interval) % 1000;
    ctx->seconds = sec(ctx->interval);
    ctx->minutes = ctx->seconds / 60;
    ctx->seconds %= 60;
}

void
TimingEnd(TimingContext *ctx)
{
    gettime(&ctx->end);
    ctx->interval = ctx->end;
    subtime(ctx->interval, ctx->start);
    timingUpdate(ctx);
}

char *
TimingGenerateString(TimingContext *ctx)
{
    static char sBuf[4096];

    sprintf(sBuf, "%d minutes, %d.%03d seconds", ctx->minutes,
            ctx->seconds, ctx->millisecs);
    return sBuf;
}

static void
dumpBytes(unsigned char *b, int l)
{
    int i;
    if (l <= 0)
        return;
    for (i = 0; i < l; ++i) {
        if (i % 16 == 0)
            printf("\t");
        printf(" %02x", b[i]);
        if (i % 16 == 15)
            printf("\n");
    }
    if ((i % 16) != 0)
        printf("\n");
    printf("\n");
}

static mp_err
testNewFuncs(const unsigned char *modulusBytes, int modulus_len)
{
    mp_err mperr = MP_OKAY;
    mp_int modulus;
    unsigned char buf[512];

    mperr = mp_init(&modulus);
    mperr = mp_read_unsigned_octets(&modulus, modulusBytes, modulus_len);
    mperr = mp_to_fixlen_octets(&modulus, buf, modulus_len);
    mperr = mp_to_fixlen_octets(&modulus, buf, modulus_len + 1);
    mperr = mp_to_fixlen_octets(&modulus, buf, modulus_len + 4);
    mperr = mp_to_unsigned_octets(&modulus, buf, modulus_len);
    mperr = mp_to_signed_octets(&modulus, buf, modulus_len + 1);
    mp_clear(&modulus);
    return mperr;
}

int
testModExp(const unsigned char *modulusBytes,
           const unsigned int expo,
           const unsigned char *input,
           unsigned char *output,
           int modulus_len)
{
    mp_err mperr = MP_OKAY;
    mp_int modulus;
    mp_int base;
    mp_int exponent;
    mp_int result;

    mperr = mp_init(&modulus);
    mperr += mp_init(&base);
    mperr += mp_init(&exponent);
    mperr += mp_init(&result);
    /* we initialize all mp_ints unconditionally, even if some fail.
    ** This guarantees that the DIGITS pointer is valid (even if null).
    ** So, mp_clear will do the right thing below.
    */
    if (mperr == MP_OKAY) {
        mperr = mp_read_unsigned_octets(&modulus,
                                        modulusBytes + (sizeof default_n - modulus_len), modulus_len);
        mperr += mp_read_unsigned_octets(&base, input, modulus_len);
        mp_set(&exponent, expo);
        if (mperr == MP_OKAY) {
#if OLD_WAY
            mperr = s_mp_exptmod(&base, &exponent, &modulus, &result);
#else
            mperr = mp_exptmod(&base, &exponent, &modulus, &result);
#endif
            if (mperr == MP_OKAY) {
                mperr = mp_to_fixlen_octets(&result, output, modulus_len);
            }
        }
    }
    mp_clear(&base);
    mp_clear(&result);

    mp_clear(&modulus);
    mp_clear(&exponent);

    return (int)mperr;
}

int
doModExp(const unsigned char *modulusBytes,
         const unsigned char *exponentBytes,
         const unsigned char *input,
         unsigned char *output,
         int modulus_len)
{
    mp_err mperr = MP_OKAY;
    mp_int modulus;
    mp_int base;
    mp_int exponent;
    mp_int result;

    mperr = mp_init(&modulus);
    mperr += mp_init(&base);
    mperr += mp_init(&exponent);
    mperr += mp_init(&result);
    /* we initialize all mp_ints unconditionally, even if some fail.
    ** This guarantees that the DIGITS pointer is valid (even if null).
    ** So, mp_clear will do the right thing below.
    */
    if (mperr == MP_OKAY) {
        mperr = mp_read_unsigned_octets(&modulus,
                                        modulusBytes + (sizeof default_n - modulus_len), modulus_len);
        mperr += mp_read_unsigned_octets(&exponent, exponentBytes, modulus_len);
        mperr += mp_read_unsigned_octets(&base, input, modulus_len);
        if (mperr == MP_OKAY) {
#if OLD_WAY
            mperr = s_mp_exptmod(&base, &exponent, &modulus, &result);
#else
            mperr = mp_exptmod(&base, &exponent, &modulus, &result);
#endif
            if (mperr == MP_OKAY) {
                mperr = mp_to_fixlen_octets(&result, output, modulus_len);
            }
        }
    }
    mp_clear(&base);
    mp_clear(&result);

    mp_clear(&modulus);
    mp_clear(&exponent);

    return (int)mperr;
}

int
main(int argc, char **argv)
{
    TimingContext *timeCtx;
    char *progName;
    long iters = DEFAULT_ITERS;
    unsigned int modulus_len;
    int i;
    int rv;
    unsigned char buf[1024];
    unsigned char buf2[1024];

    progName = strrchr(argv[0], '/');
    if (!progName)
        progName = strrchr(argv[0], '\\');
    progName = progName ? progName + 1 : argv[0];

    if (argc >= 2) {
        iters = atol(argv[1]);
    }

    if (argc >= 3) {
        modulus_len = atol(argv[2]);
    } else
        modulus_len = sizeof default_n;

    /* no library init function !? */

    memset(buf, 0x41, sizeof buf);

    if (iters < 2) {
        testNewFuncs(default_n, modulus_len);
        testNewFuncs(default_n + 1, modulus_len - 1);
        testNewFuncs(default_n + 2, modulus_len - 2);
        testNewFuncs(default_n + 3, modulus_len - 3);

        rv = testModExp(default_n, 0, buf, buf2, modulus_len);
        dumpBytes((unsigned char *)buf2, modulus_len);

        rv = testModExp(default_n, 1, buf, buf2, modulus_len);
        dumpBytes((unsigned char *)buf2, modulus_len);

        rv = testModExp(default_n, 2, buf, buf2, modulus_len);
        dumpBytes((unsigned char *)buf2, modulus_len);

        rv = testModExp(default_n, 3, buf, buf2, modulus_len);
        dumpBytes((unsigned char *)buf2, modulus_len);
    }
    rv = doModExp(default_n, default_d, buf, buf2, modulus_len);
    if (rv != 0) {
        fprintf(stderr, "Error in modexp operation:\n");
        exit(1);
    }
    dumpBytes((unsigned char *)buf2, modulus_len);

    timeCtx = CreateTimingContext();
    TimingBegin(timeCtx);
    i = iters;
    while (i--) {
        rv = doModExp(default_n, default_d, buf, buf2, modulus_len);
        if (rv != 0) {
            fprintf(stderr, "Error in modexp operation\n");
            exit(1);
        }
    }
    TimingEnd(timeCtx);
    printf("%ld iterations in %s\n", iters, TimingGenerateString(timeCtx));

    return 0;
}