#define JM_UTIL_LIB

#include <lua.h>
#include <lauxlib.h>

#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include "sw_aes.h"

#define FREE_PTR(ptr) \
    if (NULL != (ptr)) {\
        free (ptr);\
        (ptr) = NULL;\
    }

static const unsigned int kAesKeySize = 16;
static const unsigned int kAesIVSize = 16;

static const char *
buffer2lstring(lua_State *L, size_t *sz, int index) {
    const char * ptr;
    if (lua_isuserdata(L,index)) {
        ptr = (const char *)lua_touserdata(L,index);
        *sz = (size_t)luaL_checkinteger(L, index+1);
    } else {
        ptr = luaL_checklstring(L, index, sz);
    }
    return ptr;
}

char *AES_CBCDecrypt( const char * sSource, const uint32_t iSize,
        const char * sKey, uint32_t iKeySize, 
        const char * sIv, uint32_t iIvSize)
{
// printf("\r\n ====================%d, %d, %d, %d, %d\r\n",!sSource,!sKey,iSize < kAesKeySize,iSize % kAesKeySize != 0,!poResult);
    if ( !sSource || !sKey || iSize < kAesKeySize || iSize % kAesKeySize != 0) return NULL;

    uint8_t * out = (uint8_t*)malloc(iSize);
    uint8_t *key = (uint8_t*)malloc(kAesKeySize);
    uint8_t *iv = (uint8_t*)malloc(kAesIVSize);
    memcpy( key, sKey, iKeySize > kAesKeySize ? kAesKeySize : iKeySize );
    memcpy( iv, sIv, iIvSize > kAesIVSize ? kAesIVSize : iIvSize );

    AES_CTX context;
    AES_set_key(&context, (const uint8_t *)key, (const uint8_t *)iv, AES_MODE_128);
    AES_convert_key(&context);
    AES_cbc_decrypt(&context, (const uint8_t *)sSource, out, iSize);
// printf("\r\n -----------out[%s]\r\n", out);
    char * poResult = NULL;
    if( out[iSize-1] > 0 && out[iSize-1] <= kAesKeySize && (iSize - out[iSize-1]) > 0 )
    {
        int len = iSize - out[iSize-1];
        poResult = (char*)malloc(len);
        memcpy( poResult, out, len);
    }

    FREE_PTR(out);
    FREE_PTR(key);
    FREE_PTR(iv);
    return poResult;
}

static int
lstr2lightuserdata(lua_State *L) {
    size_t len;
    const char * ptr = buffer2lstring(L, &len, 1);
    if (len >= 0x10000) {
        return luaL_error(L, "Invalid size (too long) of data : %d", (int)len);
    }

    uint8_t * buffer = malloc(len);

    buffer[0] = (len >> 8) & 0xff;
    buffer[1] = len & 0xff;

    memcpy(buffer, ptr, len);

    lua_pushlightuserdata(L, buffer);
    lua_pushinteger(L, len);
    
    return 2;
}

static int
lwxdecryptdata(lua_State *L) {
    size_t sEncryptedDataSize, sIvSize, sSessionKeySize;

    const char * sEncryptedData = buffer2lstring(L, &sEncryptedDataSize, 1);
// printf("\r\n sEncryptedData[%zu]\r\n",sEncryptedDataSize);
    const char * sIv = buffer2lstring(L, &sIvSize, 2);
// printf("\r\n sIv[%zu]\r\n",sIvSize);
    const char * sSessionkey = buffer2lstring(L, &sSessionKeySize, 3);
// printf("\r\n sSessionkey[%zu]\r\n",sSessionKeySize);

    char *res = AES_CBCDecrypt(sEncryptedData, sEncryptedDataSize, sSessionkey, sSessionKeySize, sIv, sIvSize);
    // printf("\r\n out[%s]\r\n", res);
    lua_pushstring(L, res);
    return 1;
}

LUAMOD_API int
luaopen_jmutil(lua_State *L) {
    luaL_checkversion(L);
    luaL_Reg l[] = {
        { "str2lightuserdata", lstr2lightuserdata},
        { "wxdecryptdata", lwxdecryptdata},
        { NULL, NULL },
    };
    luaL_newlib(L,l);

    return 1;
}