import re


primitives = ['mustfail']
for size in '6960119','6688128','8192128','460896','348864':
    for primitive in [size, size+'f']:
        primitives.append(primitive)

cfile = ''' /* autogenerated ./testvalgrind.py > testvalgrind.c */

#include <unistd.h>
#include <string.h>
#include <sys/random.h>
#include <valgrind/memcheck.h>
#include <mceliece.h>
#include <stdlib.h>
#include <assert.h>
#include <randombytes.h>

void secret(void *xv, long long xlen) {
    VALGRIND_MAKE_MEM_UNDEFINED(xv, xlen);
}

void public(void *xv, long long xlen) {
    VALGRIND_MAKE_MEM_DEFINED(xv, xlen);
}

void getrandom_wrapper(void *xv, long long xlen) {

    unsigned char *x = (unsigned char *) xv;

    while (xlen > 0) {
        long long todo = 1048576;
        if (xlen < 1048576) todo = xlen;
        todo = getrandom(x, todo, 0);
        if (todo < 1) {
            sleep(1);
            continue;
        }
        x += todo;
        xlen -= todo;
    }
}

void randombytes(void *xv, long long xlen) {
    getrandom_wrapper(xv, xlen);
    secret(xv, xlen);
}

static void *alloc(long long len) {
    unsigned char *x = malloc(len);
    assert(x);
    return x;
}


#define mceliece_kem_mustfail_PUBLICKEYBYTES 256
#define mceliece_kem_mustfail_SECRETKEYBYTES 256
#define mceliece_kem_mustfail_CIPHERTEXTBYTES 256
#define mceliece_kem_mustfail_BYTES 256

void mceliece_kem_mustfail_keypair(unsigned char *pk, unsigned char *sk) {

    long long i, j;

    randombytes(sk, 256);
    for (i = 0; i < 256; ++i) {
        for (j = 0; j < 256; ++j) {
            if (sk[i] == j) pk[i] = sk[i];
        }
    }
}

int mceliece_kem_mustfail_enc(unsigned char *c, unsigned char *k, const unsigned char *pk) {

    long long i;

    randombytes(c, 256);
    for (i = 0; i < 256; ++i) {
        k[i] = pk[c[i]];
    }
    return 0;
}

int mceliece_kem_mustfail_dec(unsigned char *k, const unsigned char *c, const unsigned char *sk) {

    long long i;

    for (i = 0; i < 256; ++i) {
        k[i] = sk[c[i]];
    }
    return 0;
}

long long mceliece_numimpl_kem_mustfail(void) {
    return 1;
}
void (*mceliece_dispatch_kem_mustfail_keypair(long long i))(unsigned char *, unsigned char *) {
    (void) i;
    return  mceliece_kem_mustfail_keypair;
}
int (*mceliece_dispatch_kem_mustfail_enc(long long i))(unsigned char *, unsigned char *, const unsigned char *) {
    (void) i;
    return  mceliece_kem_mustfail_enc;
}
int (*mceliece_dispatch_kem_mustfail_dec(long long i))(unsigned char *, const unsigned char *, const unsigned char *) {
    (void) i;
    return  mceliece_kem_mustfail_dec;
}
const char *mceliece_dispatch_kem_mustfail_implementation(long long i) {
    (void) i;
    return "ref";
}
const char *mceliece_kem_mustfail_implementation(void) {
    return "ref";
}
'''


testf = '''
void test_PRIMITIVE(const char *name, const char *implname) {
    void (*crypto_kem_keypair)(unsigned char *,unsigned char *) = 0;
    int (*crypto_kem_enc)(unsigned char *,unsigned char *,const unsigned char *) = 0;
    int (*crypto_kem_dec)(unsigned char *,const unsigned char *,const unsigned char *) = 0;
    long long pklen = mceliece_kem_PRIMITIVE_PUBLICKEYBYTES;
    long long sklen = mceliece_kem_PRIMITIVE_SECRETKEYBYTES;
    long long clen =  mceliece_kem_PRIMITIVE_CIPHERTEXTBYTES;
    long long klen = mceliece_kem_PRIMITIVE_BYTES;
    unsigned char *pk, *sk, *c, *k1, *k2, *space;
    int result;

    if (strcmp(name, "PRIMITIVE")) return;

    space = alloc(pklen + sklen + clen + klen + klen);
    pk = space; space += pklen;
    sk = space; space += sklen;
    c = space; space += clen;
    k1 = space; space += klen;
    k2 = space; space += klen;

    for (long long impl = 0; impl < mceliece_numimpl_kem_PRIMITIVE(); ++impl) {
        if (!strcmp(implname, mceliece_dispatch_kem_PRIMITIVE_implementation(impl))) {
            crypto_kem_keypair = mceliece_dispatch_kem_PRIMITIVE_keypair(impl);
            crypto_kem_enc = mceliece_dispatch_kem_PRIMITIVE_enc(impl);
            crypto_kem_dec = mceliece_dispatch_kem_PRIMITIVE_dec(impl);
            break;
        }
    }
    if (!crypto_kem_keypair) exit(10);

    /* keypair */
    crypto_kem_keypair(pk, sk);
    public(pk, pklen);
    public(sk, sklen);

    /* enc */
    secret(pk, pklen);
    result = crypto_kem_enc(c, k1, pk);
    public(&result, sizeof result);
    assert(result == 0);
    public(c, clen);
    public(k1, klen);
    public(pk, pklen);

    /* dec */
    secret(sk, sklen);
    secret(c, clen);
    result = crypto_kem_dec(k2, c, sk);
    public(&result, sizeof result);
    assert(result == 0);
    public(k2, klen);
    public(sk, sklen);
    public(c, clen);

    assert(!memcmp(k1, k2, klen));
    exit(0);
}
'''

for primitive in primitives:
    cfile += re.sub('PRIMITIVE', primitive, testf)


cfile += '''
int main(int argc, char **argv) {

    assert(argc == 3);
    assert(argv[0]);
    assert(argv[1]);
    assert(argv[2]);

'''

for primitive in primitives:
        cfile += '''    test_%s(argv[1], argv[2]);
''' % (primitive)

cfile += '''
    return 100;
}
'''

print(cfile)
