#include <PR/bcp.h>
#include <PR/bbvirage.h>
#include <PR/bbnand.h>
#include <PR/bbskapi.h>
#include "skreencrypt.h"
#include <bb_nn.h>
#include <aes.h>
#include <algorithms.h>
#include "util.h"

typedef struct {
    BbEccSig signature; /* over all data including numFields and below */
    u32 numFields; 
} BbReEncryptionHead;

typedef struct {
    BbContentId cid;
    BbAesKey key;
    u32 status;
    u32 padding[2]; /* Each field must be a multiple of 16 bytes, since 
                       they are encrypted with AES-128 */
} BbReEncryptionField;

#define GET_LIST_SIZE(__num_fields)  \
    (sizeof(BbReEncryptionHead) + \
     (__num_fields) * sizeof(BbReEncryptionField))

#define GET_LIST_SIG_SIZE(__num_fields)  \
    (sizeof(BbReEncryptionHead) - sizeof(BbEccSig) + \
     (__num_fields) * sizeof(BbReEncryptionField))

static int 
verifyList(BbReEncryptionHead *pList)
{
    int res;
    BbEccPublicKey pubkey;

    int size = GET_LIST_SIG_SIZE(pList->numFields);

    getBbPublicKey(pubkey);
    res = verifyEccSig(((u8 *)(pList))+sizeof(BbEccSig), 
		 size, pubkey, pList->signature, SK_IDENTITY);
    if(res != SK_API_SUCCESS){
        message("Recrypt: List verify failed\n");
        return SK_API_FAIL;
    }
    return SK_API_SUCCESS;
}

static void
signList(BbReEncryptionHead *pList)
{
    int size;
    size = GET_LIST_SIG_SIZE(pList->numFields);
    computeEccSig(((u8 *)(pList)) + sizeof(BbEccSig), 
		  size, v2->privateKey, pList->signature, SK_IDENTITY);
}

static int 
newList(BbReEncryptionHead *pList)
{
    pList->numFields = 0;
    signList(pList);

    return SK_API_SUCCESS;
}

static void
getFieldFromList(BbReEncryptionField *field, BbReEncryptionHead *pKeyList, int num)
{
  BbAesIv iv;
  int i;

  /* hard code iv to be bbid +i */
  for(i =0 ; i < sizeof(BbAesIv)/sizeof(u32); i++){
    iv[i] = v2->bbId + i;
  }

  aes_SwDecrypt((u8 *) v2->recryptListKey, (u8 *)&iv, 
                ((u8 *) pKeyList) + sizeof(BbReEncryptionHead) + num *sizeof(BbReEncryptionField),
                sizeof(BbReEncryptionField), 
                (u8 *) field);
}

static void
putFieldToList(BbReEncryptionField *field, BbReEncryptionHead *pKeyList, int num)
{
  BbReEncryptionField fieldcopy;
  int i;
  BbAesIv iv;

  /* hard code iv to be bbid +i */
  for(i =0 ; i < sizeof(BbAesIv)/sizeof(u32); i++){
    iv[i] = v2->bbId + i;
  }
  aes_SwEncrypt((u8 *)v2->recryptListKey, (u8 *) &iv, 
                (u8 *) field, sizeof(BbReEncryptionField), 
                (u8 *) &fieldcopy);
  memcpy(((u8 *) pKeyList) + sizeof(BbReEncryptionHead) + 
	 num *sizeof(BbReEncryptionField), (u8 *)&fieldcopy, 
	 sizeof(BbReEncryptionField));
  
}

/*
 * Search list for field with certain cid
 * This function assumes the list exists and has been verified.
 */ 
static int
searchList(BbReEncryptionHead *pList, int cid, BbReEncryptionField *pField)
{
    int i;
    
    for(i = 0; i < pList->numFields; i++){
	getFieldFromList(pField, pList, i);
        if (pField->cid == cid) {
            return i;
        }
    }
    return SK_API_FAIL;

}

int 
isValidList(void *pKeyList)
{
    BbReEncryptionHead *pList = (BbReEncryptionHead *) pKeyList;
    int size = GET_LIST_SIZE(pList->numFields);

    if (size < SK_RECRYPT_LIST_MAX_SIZE) {
        return verifyList(pList);
    } else {
        return SK_API_FAIL;
    }

}

int 
updateRecryptEntry(void *pKeyList, BbContentId cid, int status)
{
    BbReEncryptionHead *headptr;
    BbReEncryptionField field;
    int ret;

    headptr = (BbReEncryptionHead *) pKeyList;

    /* verify first */
    if (isValidList(pKeyList) != SK_API_SUCCESS) {
        return SK_API_FAIL;
    }

    if ((ret = searchList(pKeyList, cid, &field)) < 0) {
        return ret;
    } else {
#ifdef DEBUG
        message("Recrypt update: key match\n");
#endif
        field.status = status;
        putFieldToList(&field, pKeyList, ret);
    }

    signList(pKeyList);
    return SK_API_SUCCESS;
}

int 
getRecryptKey(void *pKeyList, BbAesKey key, BbContentId cid)
{
    int ret;
    BbReEncryptionHead *headptr = (BbReEncryptionHead *) pKeyList;
    BbReEncryptionField field;

    if (isValidList(pKeyList) != SK_API_SUCCESS) {
#ifdef DEBUG
        message("Recrypt: New list\n");
#endif
        newList(pKeyList);
    }

    if ((ret = searchList(pKeyList, cid, &field)) < 0) {
        /*
         * Make a new entry
         */
#ifdef DEBUG
        message("New entry\n");
#endif
        field.cid = cid;
        getRandoms(field.key, sizeof(BbAesKey)/sizeof(u32));
        field.status = SK_API_RECRYPT_INCOMPLETE;
        putFieldToList(&field, pKeyList, headptr->numFields);
        headptr->numFields++;
        signList(headptr);
        ret = SK_API_RECRYPT_NEW;
    } else {
        ret = field.status;
    }

    memcpy(key, field.key, sizeof(BbAesKey));
    return ret;

}