#include <stdlib.h>
#include <stdio.h>

#include "pi_util.h"

#include "atb.h"
#include "ecc/pi_ecc.h"


/*
 * Generate 64MB file to initialize flash with randoms
 *
 *  Format has 512B 16B 512B 16B ...
 *
 */

/* 64MB data file */
#define FLASH_SIZE 0x4000000
//#define FLASH_SIZE 0x800

void flashInit(u8 *key, u8 *iv, u8 *data, int num_kbytes, FILE *fdata);

int main(int argc,char *argv[])
{
    int i,j;
    FILE *fdata;
    u8 expkey[44*4],*indata=NULL;
    u8 iv[32];
    
    
    unsigned char key[] = {
    0x10,0x11,0x12,0x13,
    0x14,0x15,0x16,0x17,
    0x18,0x19,0x1a,0x1b,
    0x1c,0x1d,0x1e,0x1f
};

    unsigned char first_iv[] = {
    0x10,0x11,0x12,0x13,
    0x14,0x15,0x16,0x17,
    0x18,0x19,0x1a,0x1b,
    0x1c,0x1d,0x1e,0x1f
};
    
    if(argc!=2){
        printf("Need one argument for output filename\n");
        exit(1);
    }

    if ((fdata=fopen(argv[1],"w"))==0) {
        printf("cannot open %s\n",argv[1]);
        exit(1);
    }



    aesKeyExpand(key,expkey);
  
    /* 
     * setup PI_FLASH_CTRL for ATB DMA
     *  (assumes PI_FLASH_CONF already initialized) 
     */

    /* 
     * setup AES for ATB DMA
     */

    
    /* write encrypted data to buffer */
    indata = (u8 *)malloc(sizeof(u8) *64*1024*1024);

    
    for(i=0;i<64*1024*1024;i+=4){
      //    *((u32 *)(indata+i))= (u32) i;
      *((u32 *)(indata+i))= (u32) rand();
    }
    
    /*initialise flash according to atb, 16K repeating pattern, 64MB flash
      (input is number of kilo bytes) */
    for(i=0;i<16;i++){
      iv[i] = first_iv[i];
    }

    flashInit(key, iv, indata, 64*1024, fdata);
    /*
    for(i=0;i<FLASH_SIZE/512;i++){
        fprintf(fdata,"// page %d\n",i);
        for(j=0;j<528;j++){
            fprintf(fdata,"%02x ",(unsigned char)rand());
            if((j&0xf)==0xf)
                fprintf(fdata,"\n");
        }
        fprintf(fdata,"\n");
    }
    */
    fclose(fdata);
    free(indata);
    return 0;
}



/* initialisation of flash contents with encrypted known content.
key is allocated and expanded and passed to the function below to create
the encrypted data . Last is size of flash to fill, data_size is in bytes*/

void flashInit(u8 *key, u8 *iv, u8 *data, int num_kbytes, FILE *fdata){
  u8 encdata[512];
  u8 writedata[512];
  u8 indata[512];
  u8 ecc[16];
  u32 num_pages_in_data_block;
  u32 num_AES_blocks_in_page;
  int i,j,k;
  
  num_pages_in_data_block = num_kbytes*1024/512;
  num_AES_blocks_in_page = 512/16;
  
  for(i=0;i<num_pages_in_data_block; i++){
    /*load page data */
    for(j=0; j<512;j++){
      indata[j] = data[512*i + j];
    }
    
    /*encrypt a page*/
    aesEncrypt(key, iv, indata, 512, encdata);
    
    /*last 128 bits of data is IV */
    for(k=0;k<16; k++){
      iv[k] = encdata[512-16+k];
    }
        
    for(k=0; k< 512; k++){
      writedata[k] = encdata[k];
    }
    
    
    /* this is to make it work with later tests: write page number */
    fprintf(fdata,"// page %d\n",i);
    /* save writedata in file */
    for(j=0;j<512;j++){
      fprintf(fdata,"%02x ",(unsigned char)writedata[j]);
      if((j&0xf)==0xf)
	fprintf(fdata,"\n");
    }
    
    /* calculate and save ecc */
    
    calculate_ecc(encdata,ecc+13);
    calculate_ecc(encdata+256,ecc+8);
    
    for(j=0;j<16;j++){
      fprintf(fdata, "%02x ",(unsigned char)ecc[j]);
    }
    

    fprintf(fdata, "\n");
    fflush(fdata);
  }
  
  
}