/*
 * Copyright 2002-2006 The Apache Software Foundation.
 *
 * 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
 * imitations under the License.
 *
 * @author Werner Dittmann (werner@apache.org)
 */

// #define TESTING

#include <stdio.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/rand.h>
#include <openssl/buffer.h>
#include <jni.h>
#include "org_apache_security_juice_provider_JCEF8CipherOpenSSL.h"

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include "openSSLNative.h"

extern char bufferOpenSSL[1024];

typedef struct _f8_ctx {
    unsigned char *S;
    unsigned char *ivAccent;
    EVP_CIPHER_CTX *cipher;
    int blocksize;
    union {
        unsigned int J;
        unsigned char jb[4];
    } J;
} F8_CIPHER_CTX;

/*
 * The forward declarations
 */
int F8_CipherInit(F8_CIPHER_CTX *f8ctx, unsigned char *key, int keyLen, unsigned char *iv, 
                unsigned char *salt, int saltLent, const EVP_CIPHER *c);

int F8_CipherUpdate(F8_CIPHER_CTX *f8ctx, unsigned char *out, int *outlen, unsigned char *in, int inlen);

int F8_EncryptFinal(F8_CIPHER_CTX *f8ctx);

void F8_CIPHER_CTX_cleanup(F8_CIPHER_CTX *f8ctx);

void F8_CIPHER_CTX_init(F8_CIPHER_CTX *f8ctx);

/**
 * At first we have the JNI / Java wrapper methods.
 */

/*
 * Class:     org_apache_security_juice_provider_JCEF8CipherOpenSSL
 * Method:    initOpenSSLF8
 * Signature: (Ljava/lang/String;[BI[B[BI)J
 */
JNIEXPORT jlong JNICALL Java_org_apache_security_juice_provider_JCEF8CipherOpenSSL_initOpenSSLF8
(JNIEnv *env, jobject jobj, jstring algoName, jbyteArray key, jint keyLen, jbyteArray iv, jbyteArray salt, jint saltLen) {

    jclass clazz;
    jbyte *keyBytes, *ivBytes, *saltBytes;
    F8_CIPHER_CTX *f8ctx;
    const EVP_CIPHER *c;
    int ret;    
    
    const char *algo = (*env)->GetStringUTFChars(env, algoName, 0);

    c = EVP_get_cipherbyname(algo);

    if (!c) {
        clazz = (*env)->FindClass(env, NO_SUCH_ALGO_EX);
        BIO_snprintf(bufferOpenSSL, sizeof(bufferOpenSSL), "JuICE OpenSSL: cipher algorithm not supported: %s", 
                     algo);
        (*env)->ReleaseStringUTFChars(env, algoName, algo);
        (*env)->ThrowNew(env, clazz, bufferOpenSSL);
        return 0L;
    }
    (*env)->ReleaseStringUTFChars(env, algoName, algo);
    
    // Get pointer to key, if NULL we are "OutOfMemory"
    keyBytes = (*env)->GetByteArrayElements(env, key, 0);
    if (keyBytes == NULL) {
        return 0L;
    }

    ivBytes = (*env)->GetByteArrayElements(env, iv, 0);
    if (ivBytes == NULL) {
        return 0L;
    }
    
    saltBytes = (*env)->GetByteArrayElements(env, salt, 0);
    if (saltBytes == NULL) {
        return 0L;
    }
    
    f8ctx = malloc(sizeof(F8_CIPHER_CTX));
    if (f8ctx == NULL) {
	JUICEerr(F_initOpenSSLF8, ERR_R_MALLOC_FAILURE);
	clazz = (*env)->FindClass(env, OUT_OF_MEMORY_EX);
        (*env)->ThrowNew(env, clazz, getErrorStringOpenSSL(OUT_OF_MEMORY_MSG));
	return 0L;
    }
    F8_CIPHER_CTX_init(f8ctx);

    ret = F8_CipherInit(f8ctx, keyBytes, keyLen, ivBytes, saltBytes, saltLen, c);

    (*env)->ReleaseByteArrayElements(env, key, keyBytes, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, iv, ivBytes, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, salt, saltBytes, JNI_ABORT);

    if (ret == 0) {
        clazz = (*env)->FindClass(env, RUNTIME_EX);
        (*env)->ThrowNew(env, clazz, getErrorStringOpenSSL("JuiCE OpenSSL: error during F8 cipher initialize"));
        free(f8ctx);
        return 0L;
    }    
    return (long)f8ctx;
}

/*
 * Class:     org_apache_security_juice_provider_JCEF8CipherOpenSSL
 * Method:    processBytesOpenSSL
 * Signature: (J[BII[BI)I
 */
JNIEXPORT jint JNICALL Java_org_apache_security_juice_provider_JCEF8CipherOpenSSL_processBytesOpenSSL
(JNIEnv *env, jobject jobj, jlong handle, jbyteArray input, jint inOff, jint inLen, jbyteArray output, jint outOff) {

    int ret = 0;
    int outLen = -1;

    jbyte* inBytes, *outBytes;
    jclass clazz;


#if SIZEOF_INT_P == 4
    F8_CIPHER_CTX  *ctx = (F8_CIPHER_CTX  *)((int)(handle & 0xffffffff));
#else
    F8_CIPHER_CTX  *ctx = (F8_CIPHER_CTX  *)handle;
#endif

    if (ctx == NULL) {
        clazz = (*env)->FindClass(env, ILLEGAL_ARGUMENT_EX);
        (*env)->ThrowNew(env, clazz, getErrorStringOpenSSL(CONTEXT_NULL_MSG));
    }
    /*
     * We check if the pointer of the array opeartion is NULL. In that
     * case an "OutOfMemoryError" occured. Return immediatly. The
     * exception will be handled by the java code (according to JNI
     * documentation)
     */

    inBytes = (*env)->GetByteArrayElements(env, input, 0);
    if (inBytes == NULL) {
        return 0;
    }
    // Get pointer to output, if NULL we are "OutOfMemory"
    outBytes = (*env)->GetByteArrayElements(env, output, 0);
    if (outBytes == NULL) {
        return 0;
    }
    ret = F8_CipherUpdate(ctx, outBytes+outOff, &outLen, inBytes+inOff, inLen);

    (*env)->ReleaseByteArrayElements(env, input, inBytes, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, output, outBytes, 0);

    if(ret == 0) {
        free(ctx);
        clazz = (*env)->FindClass(env, RUNTIME_EX);
        (*env)->ThrowNew(env, clazz, getErrorStringOpenSSL("JuiCE OpenSSL: error during F8 cipher update"));
        return 0;
    }
    return outLen;
}

/*
 * Class:     org_apache_security_juice_provider_JCEF8CipherOpenSSL
 * Method:    doFinalOpenSSL
 * Signature: (J[BI)I
 */
JNIEXPORT jint JNICALL Java_org_apache_security_juice_provider_JCEF8CipherOpenSSL_doFinalOpenSSL
(JNIEnv *env, jobject jobj, jlong handle, jbyteArray output, jint outOff) {
    int ret = 0;
    int outLen = 0;
    jclass clazz;

#if SIZEOF_INT_P == 4
    F8_CIPHER_CTX *ctx = (F8_CIPHER_CTX *)((int)(handle & 0xffffffff));
#else
    F8_CIPHER_CTX *ctx = (F8_CIPHER_CTX *)handle;
#endif

    if (ctx == NULL) {
        clazz = (*env)->FindClass(env, ILLEGAL_ARGUMENT_EX);
        (*env)->ThrowNew(env, clazz, getErrorStringOpenSSL(CONTEXT_NULL_MSG));
    }
    ret = F8_EncryptFinal(ctx);

    F8_CIPHER_CTX_cleanup(ctx);
    free(ctx);

    if(ret == 0) {
        clazz = (*env)->FindClass(env, RUNTIME_EX);
        (*env)->ThrowNew(env, clazz, getErrorStringOpenSSL("JuiCE OpenSSL: error during F8 cipher doFinal"));
        return 0;
    }
    return outLen;
}

/**
 * Now the F8 implementation as pure C functions
 */    
    

/**
 * 
 * Initialize the F8 context. 
 * 
 * Refer to RFC3711 (SRTP) how to generate and initialize these data. 
 * We need to store
 * - the computed IV (called ivAccent)
 * - the key stream (called S)
 * - the counter (called J)
 * 
 * The computed IV and the key stream fields must have the same length
 * as the basic cipher's blocksize. The counter is an unsigned int.
 * The value of this counter shall not exceed 2**32 (refer to RFC3711).
 */

int F8_CipherInit(F8_CIPHER_CTX *f8ctx, unsigned char *key, int keyLen, unsigned char *iv, 
                unsigned char *salt, int saltLen, const EVP_CIPHER *c) {
                            
    unsigned char *saltMask;
    unsigned char *maskedKey;
    unsigned char *ivAccent;
    unsigned char *cp_in, *cp_in1, *cp_out;
    int i, ret, blocksize;
    EVP_CIPHER_CTX *ctx;
    
    blocksize = c->block_size;
    
    /*
     * Get memory for the derived IV (IV')
     */
    ivAccent = malloc(blocksize);
    
    /*
     * Get memory for the special key that is input to compute the
     * derived IV (IV').
     */
    saltMask = malloc(keyLen);
    maskedKey = malloc(keyLen);

    if (saltMask == NULL || maskedKey == NULL || ivAccent == NULL) {
        JUICEerr(F_F8_CipherInit, ERR_R_MALLOC_FAILURE);
        return 0;
    }
    /*
     * First copy the salt into the mask field, the fill with 0x55 to
     * get a full key.
     */
    memcpy(saltMask, salt, saltLen);
    memset(saltMask+saltLen, 0x55, keyLen-saltLen);
    
    cp_out = maskedKey;
    cp_in = key;
    cp_in1 = saltMask;
    
    /*
     * XOR the original key with the above created mask to
     * get the special key.
     */
    for (i = 0; i < keyLen; i++) {
        *cp_out++ = *cp_in++ ^ *cp_in1++;
    }    
    /*
     * Prepare the first AES-ECB cipher to compute IV'
     */
    ctx = malloc(sizeof (EVP_CIPHER_CTX));
    if (ctx == NULL) {
        JUICEerr(F_F8_CipherInit, ERR_R_MALLOC_FAILURE);
        return 0;
    }
    /*
     * Even if the basic cipher is in ECB mode we need to give an IV
     * here - the cipher ignores it (if you don't specify this IV we get a
     * NULL pinter exception).
     * 
     * Use the masked key to encrypt the original IV to produce IV'. Because
     * this is ECB mode the following doFinal should not produce any output.
     * However, we need an address there to avoid NPE.
     * 
     * After computing the IV' we don't need this cipher context anymore, free it.
     * 
     */
    EVP_CIPHER_CTX_init(ctx);

    f8ctx->cipher = ctx;

    ret = EVP_CipherInit_ex(ctx, c, NULL, maskedKey, iv, 1);
    if (ret != 0) {
	ret = EVP_CipherUpdate(ctx, ivAccent, &i, iv, blocksize);
    }
    if (ret != 0) {
	ret = EVP_EncryptFinal_ex(ctx, saltMask, &i);
    }

    free(saltMask);
    memset(maskedKey, 0, keyLen);
    free(maskedKey);                    // both values no longer needed
    /*
     * Re-initialize cipher context for use in real processing. This
     * cipher uses the original key to compute the key stream.
     */
    if (ret != 0) {
	ret = EVP_CipherInit_ex(ctx, c, NULL, key, iv, 1);
    }

    if (ret == 0) {
        F8_CIPHER_CTX_cleanup(f8ctx);
	return 0;
    }
    f8ctx->ivAccent = ivAccent;         // save the computed IV'
    f8ctx->J.J = 0;                     // initialize the counter
    f8ctx->blocksize = blocksize;
    f8ctx->S = malloc(blocksize);       // get key stream buffer

    if (f8ctx->S == NULL) {
        JUICEerr(F_F8_CipherInit, ERR_R_MALLOC_FAILURE);
        F8_CIPHER_CTX_cleanup(f8ctx);
        return 0;
    }
    memset(f8ctx->S, 0, blocksize);     // initial value for key stream
    return 1;
}

void F8_CIPHER_CTX_cleanup(F8_CIPHER_CTX *ctx) {
    if (ctx->S != NULL) {
	memset(ctx->S, 0, ctx->blocksize);
        free(ctx->S);
    }
    if (ctx->ivAccent != NULL) {
	memset(ctx->ivAccent, 0, ctx->blocksize);
        free(ctx->ivAccent);
    }
    if (ctx->cipher != NULL) {
        EVP_CIPHER_CTX_cleanup(ctx->cipher);
        free(ctx->cipher);
    }
    ctx->S = NULL;
    ctx->ivAccent = NULL;
    ctx->cipher = NULL;
}

void F8_CIPHER_CTX_init(F8_CIPHER_CTX *f8ctx) {
    f8ctx->S = NULL;
    f8ctx->ivAccent = NULL;
    f8ctx->cipher = NULL;
}

static int processBlock(F8_CIPHER_CTX *f8ctx, unsigned char *in, int length, unsigned char *out) {

    int i, ret, outlen;
    unsigned char *cp_in, *cp_in1, *cp_out;

    /*
     * XOR the previous key stream with IV'
     * ( S(-1) xor IV' )
     */
    cp_in = f8ctx->ivAccent;
    cp_out = f8ctx->S;
    for (i = 0; i < f8ctx->blocksize; i++) {
        *cp_out++ ^= *cp_in++;
    }
    cp_out = f8ctx->S + f8ctx->blocksize;
    cp_in = f8ctx->J.jb;
    /*
     * Now XOR (S(-1) xor IV') with the current counter, then increment the counter
     */
    *--cp_out ^= *cp_in++;
    *--cp_out ^= *cp_in++;
    *--cp_out ^= *cp_in++;                
    *--cp_out ^= *cp_in++;
    f8ctx->J.J++;
    /*
     * Now compute the new key stream using AES-ECB update mode
     */
    ret = EVP_CipherUpdate(f8ctx->cipher, f8ctx->S, &outlen, f8ctx->S, f8ctx->blocksize);
    /*
     * as the last step XOR the plain text with the key stream to produce 
     * the ciphertext.
     */
    cp_out = out;
    cp_in = in;
    cp_in1 = f8ctx->S;
    for (i = 0; i < length; i++) {
        *cp_out++ = *cp_in++ ^ *cp_in1++;
    }
    return length;
}

int F8_EncryptFinal(F8_CIPHER_CTX *f8ctx) {
    // TODO error check
    unsigned char outbuf[64];
    int outlen, ret;
    
    ret = EVP_EncryptFinal_ex(f8ctx->cipher, outbuf, &outlen);

    return 1;
}

int F8_CipherUpdate(F8_CIPHER_CTX *f8ctx, unsigned char *out, int *outlen, unsigned char *in, int inlen) {

    int offset = 0;
  
    *outlen = 0;
    while (inlen >= f8ctx->blocksize) {
        *outlen += processBlock(f8ctx, in+offset, f8ctx->blocksize, out+offset);
        inlen -= f8ctx->blocksize;
        offset += f8ctx->blocksize;
    }
    if (inlen > 0) {
        *outlen += processBlock(f8ctx, in+offset, inlen, out+offset);
    }
    return 1;
}

#ifdef TESTING

void hexdump(FILE *f,const char *title,const unsigned char *s,int l) {
    int n=0;

    fprintf(f,"%s",title);
    for( ; n < l ; ++n)
    {
        if((n%16) == 0)
            fprintf(f,"\n%04x",n);
        fprintf(f," %02x",s[n]);
    }
    fprintf(f,"\n");
}

F8_CIPHER_CTX *AES_F8_CIPHERInit(unsigned char *keyBytes, int keyLen, unsigned char *ivBytes, 
                                 unsigned char *saltBytes, int saltLen) {
    
    char buff[20];
    const EVP_CIPHER *c;
    F8_CIPHER_CTX *f8ctx;
    int ret;
    /*
     * The AES F8 mode uses basic AES ECB encryption mode to produce the key
     * stream
     */
    BIO_snprintf(buff, 20, "AES-%i-ECB", keyLen*8);
    c = EVP_get_cipherbyname(buff);

    if (!c) {
        BIO_snprintf(bufferOpenSSL, sizeof(bufferOpenSSL), 
                     "JuICE OpenSSL: cipher algorithm not supported for AES-f8: %s", 
                     buff);
        return NULL;
    }
    f8ctx = malloc(sizeof(F8_CIPHER_CTX));
    F8_CIPHER_CTX_init(f8ctx);

    ret = F8_CTX_init(f8ctx, keyBytes, keyLen, ivBytes, saltBytes, saltLen, c);
    if (ret == 0) {
        free(f8ctx);
        return NULL;
    }
    return (f8ctx);
}

/*
 * The data of the test vectors according to RFC3711
 */
unsigned char salt[] = {0x32, 0xf2, 0x87, 0x0d};
unsigned char iv[] = {0x00, 0x6e, 0x5c, 0xba, 0x50, 0x68, 0x1d, 0xe5, 0x5c, 0x62, 0x15, 0x99, 0xd4, 0x62, 0x56, 0x4a};
unsigned char key[]= {0x23, 0x48, 0x29, 0x00, 0x84, 0x67, 0xbe, 0x18, 0x6c, 0x3d, 0xe1, 0x4a, 0xae, 0x72, 0xd6, 0x2c};
unsigned char payload[] = {0x70, 0x73, 0x65, 0x75, 0x64, 0x6f, 0x72, 0x61, 0x6e, 0x64, 0x6f, 0x6d, 0x6e, 0x65, 0x73,
                           0x73, 0x20, 0x69, 0x73, 0x20, 0x74, 0x68, 0x65, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x62,
                           0x65, 0x73, 0x74, 0x20, 0x74, 0x68, 0x69, 0x6e, 0x67, 0};  // 39 chars + 1 null

int main() {
    F8_CIPHER_CTX *f8ctx;
    unsigned char outbuf[256];
    int outlen;
    
    printf("start main, clear text: %s\n", payload);
    OpenSSL_add_all_ciphers();
    
    // encrypt
    f8ctx = F8_CipherInit(key, 16, iv, salt, 4);
    F8_CipherUpdate(f8ctx, outbuf, &outlen, payload, 39);
    F8_EncryptFinal(f8ctx);

    // decrypt
    f8ctx = F8_CipherInit(key, 16, iv, salt, 4);
    F8_CipherUpdate(f8ctx, outbuf, &outlen, outbuf, 39);
    F8_EncryptFinal(f8ctx);

    // check
    if (strncmp(payload, outbuf, 39) != 0) {
        printf("failure: decrypted clear text does not match\n");
    }
    return 0;
}
#endif