#include <arm9_utils.h>
#include <misc.h>
#include <sha1.h>
#include <tea.h>
#include <tokendefs.h>

#define	NOR			((unsigned char *) 0xA0000000)
#define	NORWORD			((unsigned short int *) 0xA0000000)

#define CHIPID_LOC		((unsigned char *) 0xF440006C)
#define NORID_LOC		((unsigned char *) 0xFFFF3BE2)

#define SECZONE_OFFSET		0x3FA000
#define SECZONE_SIZE		0x1000

#define LOCK_TABLE_OFFSET	0xC10
#define LOCK_TABLE_SIZE		0x210

void eraseblock(int addr);
void writeblock(int block, int length);

unsigned char virgin[] = {
    0x00, 0x01, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
    0x01, 0x01, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00,
    0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x05, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x05, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x05, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x05, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
};

unsigned char NOR_ID[] = {
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
};

unsigned char CHIP_ID[] = {
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00
};

unsigned char payload[SECZONE_SIZE];

void fredtool_main() __attribute__ ((section (".fredtool_entry")));

void
fredtool_entry()
{
    unsigned long sha1[5], key[4], *ltp;
    sha1_context ctx;

    /* Copy the original seczone */
    memcpy(payload, (unsigned char *)SECZONE_OFFSET, SECZONE_SIZE);

    /* Set a pointer to the lock table */
    ltp = (unsigned long *)(payload + LOCK_TABLE_OFFSET);

    /* Zero out lock table */
    memset((unsigned char *)ltp, 0, LOCK_TABLE_SIZE);

    /* Copy the NOR and HWID */
    memcpy(NOR_ID, NORID_LOC, sizeof(NOR_ID));
    memcpy(CHIP_ID, CHIPID_LOC, sizeof(CHIP_ID));
    
    /* Create the hash */
    sha1_starts(&ctx); 
    sha1_update(&ctx, NOR_ID, sizeof(NOR_ID));
    sha1_update(&ctx, CHIP_ID, sizeof(CHIP_ID));
    sha1_finish(&ctx, (void *)sha1);

    /* Endian swap the first 16 bits of the SHA1 output */
    key[0] = htonl(sha1[0]);
    key[1] = htonl(sha1[1]);
    key[2] = htonl(sha1[2]);
    key[3] = htonl(sha1[3]);    
    
    /* Re-encrypt the lock table */
    encrypt_cbc(ltp, LOCK_TABLE_SIZE, key, ltp);

    eraseblock(SECZONE_OFFSET);
    writeblock(SECZONE_OFFSET, SECZONE_SIZE);

    /* Hang until reset */
    arm9_stop_mode();

    /* This should never be reached */
    return;
}

void
eraseblock(int addr)
{
    NOR[addr] = 0xFF;
    NOR[addr] = 0x60;
    NOR[addr] = 0xD0; /* Unlock */
    NOR[addr] = 0x20;
    NOR[addr] = 0xD0; /* Erase */
    NOR[addr] = 0x50;
    while(NOR[addr] != 0x80); /* Wait for it to be ready */
}

void
writeblock(int addr, int length)
{
    extern unsigned char *payload;
    
    unsigned int i;
    unsigned short int j;
    
    /* First unlock the blocks */
    for(i = 0; i < length; i+= 0x4000) {
	NOR[addr + i] = 0xFF;
	NOR[addr + i] = 0x60;

	/* Unlock */
	NOR[addr + i] = 0xD0;
	NOR[addr + i] = 0xFF;
    }

    NOR[0] = 0xFF;
    
    for(i = 0; i < length; i += 4) {
	j  = *(payload++);
	j |= (*(payload++) << 8);
	
	NOR[addr + i] = 0x40;
	NORWORD[(addr + i) / 2] = j;     //0x1100
	NOR[addr] = 0x70;

	/* Wait for it to be ready */
	while(NOR[addr] != 0x80);
	NOR[addr] = 0xFF;
	
	j  = *(payload++);		//2
	j |= (*(payload++) << 8);	//3

	NOR[addr + i] = 0x40;
	NORWORD[((addr + i) / 2) + 1] = j;	   //0x3322
	NOR[addr] = 0x70;
	
	/* Wait for it to be ready */
	while(NOR[addr] != 0x80);
	NOR[addr] = 0xFF;
    }
}