/*
	Draan proudly presents:
	
	With huge help from community:
	coyotebean, Davee, hitchhikr, kgsws, liquidzigong, Mathieulh, Proxima, SilverSpring
	
	******************** KIRK-ENGINE ********************
	An Open-Source implementation of KIRK (PSP crypto engine) algorithms and keys.
	Includes also additional routines for hash forging.
	
	********************
	
	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "crypto.h"

/* ------------------------- KEY VAULT ------------------------- */

u8 kirk1_key[] =   {0x98, 0xC9, 0x40, 0x97, 0x5C, 0x1D, 0x10, 0xE8, 0x7F, 0xE6, 0x0E, 0xA3, 0xFD, 0x03, 0xA8, 0xBA};

u8 kirk7_key03[] = {0x98, 0x02, 0xC4, 0xE6, 0xEC, 0x9E, 0x9E, 0x2F, 0xFC, 0x63, 0x4C, 0xE4, 0x2F, 0xBB, 0x46, 0x68};
u8 kirk7_key04[] = {0x99, 0x24, 0x4C, 0xD2, 0x58, 0xF5, 0x1B, 0xCB, 0xB0, 0x61, 0x9C, 0xA7, 0x38, 0x30, 0x07, 0x5F};
u8 kirk7_key05[] = {0x02, 0x25, 0xD7, 0xBA, 0x63, 0xEC, 0xB9, 0x4A, 0x9D, 0x23, 0x76, 0x01, 0xB3, 0xF6, 0xAC, 0x17};
u8 kirk7_key0C[] = {0x84, 0x85, 0xC8, 0x48, 0x75, 0x08, 0x43, 0xBC, 0x9B, 0x9A, 0xEC, 0xA7, 0x9C, 0x7F, 0x60, 0x18};
u8 kirk7_key0D[] = {0xB5, 0xB1, 0x6E, 0xDE, 0x23, 0xA9, 0x7B, 0x0E, 0xA1, 0x7C, 0xDB, 0xA2, 0xDC, 0xDE, 0xC4, 0x6E};
u8 kirk7_key0E[] = {0xC8, 0x71, 0xFD, 0xB3, 0xBC, 0xC5, 0xD2, 0xF2, 0xE2, 0xD7, 0x72, 0x9D, 0xDF, 0x82, 0x68, 0x82};
u8 kirk7_key0F[] = {0x0A, 0xBB, 0x33, 0x6C, 0x96, 0xD4, 0xCD, 0xD8, 0xCB, 0x5F, 0x4B, 0xE0, 0xBA, 0xDB, 0x9E, 0x03};
u8 kirk7_key10[] = {0x32, 0x29, 0x5B, 0xD5, 0xEA, 0xF7, 0xA3, 0x42, 0x16, 0xC8, 0x8E, 0x48, 0xFF, 0x50, 0xD3, 0x71};
u8 kirk7_key11[] = {0x46, 0xF2, 0x5E, 0x8E, 0x4D, 0x2A, 0xA5, 0x40, 0x73, 0x0B, 0xC4, 0x6E, 0x47, 0xEE, 0x6F, 0x0A};
u8 kirk7_key12[] = {0x5D, 0xC7, 0x11, 0x39, 0xD0, 0x19, 0x38, 0xBC, 0x02, 0x7F, 0xDD, 0xDC, 0xB0, 0x83, 0x7D, 0x9D};
u8 kirk7_key38[] = {0x12, 0x46, 0x8D, 0x7E, 0x1C, 0x42, 0x20, 0x9B, 0xBA, 0x54, 0x26, 0x83, 0x5E, 0xB0, 0x33, 0x03};
u8 kirk7_key39[] = {0xC4, 0x3B, 0xB6, 0xD6, 0x53, 0xEE, 0x67, 0x49, 0x3E, 0xA9, 0x5F, 0xBC, 0x0C, 0xED, 0x6F, 0x8A};
u8 kirk7_key3A[] = {0x2C, 0xC3, 0xCF, 0x8C, 0x28, 0x78, 0xA5, 0xA6, 0x63, 0xE2, 0xAF, 0x2D, 0x71, 0x5E, 0x86, 0xBA};
u8 kirk7_key4B[] = {0x0C, 0xFD, 0x67, 0x9A, 0xF9, 0xB4, 0x72, 0x4F, 0xD7, 0x8D, 0xD6, 0xE9, 0x96, 0x42, 0x28, 0x8B}; //1.xx game eboot.bin
u8 kirk7_key53[] = {0xAF, 0xFE, 0x8E, 0xB1, 0x3D, 0xD1, 0x7E, 0xD8, 0x0A, 0x61, 0x24, 0x1C, 0x95, 0x92, 0x56, 0xB6};
u8 kirk7_key57[] = {0x1C, 0x9B, 0xC4, 0x90, 0xE3, 0x06, 0x64, 0x81, 0xFA, 0x59, 0xFD, 0xB6, 0x00, 0xBB, 0x28, 0x70};
u8 kirk7_key5D[] = {0x11, 0x5A, 0x5D, 0x20, 0xD5, 0x3A, 0x8D, 0xD3, 0x9C, 0xC5, 0xAF, 0x41, 0x0F, 0x0F, 0x18, 0x6F};
u8 kirk7_key63[] = {0x9C, 0x9B, 0x13, 0x72, 0xF8, 0xC6, 0x40, 0xCF, 0x1C, 0x62, 0xF5, 0xD5, 0x92, 0xDD, 0xB5, 0x82};
u8 kirk7_key64[] = {0x03, 0xB3, 0x02, 0xE8, 0x5F, 0xF3, 0x81, 0xB1, 0x3B, 0x8D, 0xAA, 0x2A, 0x90, 0xFF, 0x5E, 0x61};

/* ------------------------- KEY VAULT END ------------------------- */

/* ------------------------- INTERNAL STUFF ------------------------- */

typedef struct header_keys
{
    u8 AES[16];
    u8 CMAC[16];
}header_keys;  //small struct for temporary keeping AES & CMAC key from CMD1 header

u8 fuseID[16]; //Emulate FUSEID	

AES_ctx aes_kirk1; //global

char is_kirk_initialized; //"init" emulation

/* ------------------------- INTERNAL STUFF END ------------------------- */


/* ------------------------- IMPLEMENTATION ------------------------- */

int kirk_CMD0(u8* outbuff, u8* inbuff, int size, int generate_trash)
{
	if(is_kirk_initialized == 0) return KIRK_NOT_INITIALIZED;
	
    KIRK_CMD1_HEADER* header = (KIRK_CMD1_HEADER*)outbuff;
    
    memcpy(outbuff, inbuff, size);
    
	if(header->mode != KIRK_MODE_CMD1) return KIRK_INVALID_MODE;
	
	header_keys *keys = (header_keys *)outbuff; //0-15 AES key, 16-31 CMAC key
	
	//FILL PREDATA WITH RANDOM DATA
	if(generate_trash) kirk_CMD14(outbuff+sizeof(KIRK_CMD1_HEADER), header->data_offset);
	
	//Make sure data is 16 aligned
	int chk_size = header->data_size;
	if(chk_size % 16) chk_size += 16 - (chk_size % 16);
	
	//ENCRYPT DATA
	AES_ctx k1;
	libkirk_AES_set_key(&k1, keys->AES, 128);
	
	libkirk_AES_cbc_encrypt(&k1, inbuff+sizeof(KIRK_CMD1_HEADER)+header->data_offset, outbuff+sizeof(KIRK_CMD1_HEADER)+header->data_offset, chk_size);
	
	//CMAC HASHES
	AES_ctx cmac_key;
	libkirk_AES_set_key(&cmac_key, keys->CMAC, 128);
	    
	u8 cmac_header_hash[16];
	u8 cmac_data_hash[16];
		
	libkirk_AES_CMAC(&cmac_key, outbuff+0x60, 0x30, cmac_header_hash);
	
	libkirk_AES_CMAC(&cmac_key, outbuff+0x60, 0x30 + chk_size + header->data_offset, cmac_data_hash);
	
	memcpy(header->CMAC_header_hash, cmac_header_hash, 16);
	memcpy(header->CMAC_data_hash, cmac_data_hash, 16);
	
	//ENCRYPT KEYS
	
	libkirk_AES_cbc_encrypt(&aes_kirk1, inbuff, outbuff, 16*2);
	return KIRK_OPERATION_SUCCESS;
}

int kirk_CMD1(u8* outbuff, u8* inbuff, int size, int do_check)
{
	if(is_kirk_initialized == 0) return KIRK_NOT_INITIALIZED;
	
    KIRK_CMD1_HEADER* header = (KIRK_CMD1_HEADER*)inbuff;
	if(header->mode != KIRK_MODE_CMD1) return KIRK_INVALID_MODE;
	
	header_keys keys; //0-15 AES key, 16-31 CMAC key
	
	libkirk_AES_cbc_decrypt(&aes_kirk1, inbuff, (u8*)&keys, 16*2); //decrypt AES & CMAC key to temp buffer
	
	// HOAX WARRING! I have no idea why the hash check on last IPL block fails, so there is an option to disable checking
	if(do_check)
	{
       int ret = kirk_CMD10(inbuff, size);
       if(ret != KIRK_OPERATION_SUCCESS) return ret;
    }
	
	AES_ctx k1;
	libkirk_AES_set_key(&k1, keys.AES, 128);
	
	libkirk_AES_cbc_decrypt(&k1, inbuff+sizeof(KIRK_CMD1_HEADER)+header->data_offset, outbuff, header->data_size);	
	
	return KIRK_OPERATION_SUCCESS;
}

int kirk_CMD4(u8* outbuff, u8* inbuff, int size)
{
	if(is_kirk_initialized == 0) return KIRK_NOT_INITIALIZED;
	
	KIRK_AES128CBC_HEADER *header = (KIRK_AES128CBC_HEADER*)inbuff;
	if(header->mode != KIRK_MODE_ENCRYPT_CBC) return KIRK_INVALID_MODE;
	if(header->data_size == 0) return KIRK_DATA_SIZE_ZERO;
	
	u8* key = kirk_4_7_get_key(header->keyseed);
	if(key == (u8*)KIRK_INVALID_SIZE) return KIRK_INVALID_SIZE;
	
	//Set the key
	AES_ctx aesKey;
	libkirk_AES_set_key(&aesKey, key, 128);
	libkirk_AES_cbc_encrypt(&aesKey, inbuff+sizeof(KIRK_AES128CBC_HEADER), outbuff+sizeof(KIRK_AES128CBC_HEADER), size);
	
	return KIRK_OPERATION_SUCCESS;
}

int kirk_CMD7(u8* outbuff, u8* inbuff, int size)
{
	if(is_kirk_initialized == 0) return KIRK_NOT_INITIALIZED;
	
	KIRK_AES128CBC_HEADER *header = (KIRK_AES128CBC_HEADER*)inbuff;
	if(header->mode != KIRK_MODE_DECRYPT_CBC) return KIRK_INVALID_MODE;
	if(header->data_size == 0) return KIRK_DATA_SIZE_ZERO;
	
	u8* key = kirk_4_7_get_key(header->keyseed);
	if(key == (u8*)KIRK_INVALID_SIZE) return KIRK_INVALID_SIZE;
	
	//Set the key
	AES_ctx aesKey;
	libkirk_AES_set_key(&aesKey, key, 128);
	
 	libkirk_AES_cbc_decrypt(&aesKey, inbuff+sizeof(KIRK_AES128CBC_HEADER), outbuff, size);
	
	return KIRK_OPERATION_SUCCESS;
}

int kirk_CMD10(u8* inbuff, int insize)
{
	if(is_kirk_initialized == 0) return KIRK_NOT_INITIALIZED;
	
    KIRK_CMD1_HEADER* header = (KIRK_CMD1_HEADER*)inbuff;
    
	if(!(header->mode == KIRK_MODE_CMD1 || header->mode == KIRK_MODE_CMD2 || header->mode == KIRK_MODE_CMD3)) return KIRK_INVALID_MODE;
	if(header->data_size == 0) return KIRK_DATA_SIZE_ZERO;
	
	if(header->mode == KIRK_MODE_CMD1)
	{
        header_keys keys; //0-15 AES key, 16-31 CMAC key
        
        libkirk_AES_cbc_decrypt(&aes_kirk1, inbuff, (u8*)&keys, 32); //decrypt AES & CMAC key to temp buffer
	    
	    AES_ctx cmac_key;
	    libkirk_AES_set_key(&cmac_key, keys.CMAC, 128);
	    
		u8 cmac_header_hash[16];
		u8 cmac_data_hash[16];
		
		libkirk_AES_CMAC(&cmac_key, inbuff+0x60, 0x30, cmac_header_hash);
	
		//Make sure data is 16 aligned
		int chk_size = header->data_size;
		if(chk_size % 16) chk_size += 16 - (chk_size % 16);
		libkirk_AES_CMAC(&cmac_key, inbuff+0x60, 0x30 + chk_size + header->data_offset, cmac_data_hash);
	
		if(memcmp(cmac_header_hash, header->CMAC_header_hash, 16) != 0) return KIRK_HEADER_HASH_INVALID;
		if(memcmp(cmac_data_hash, header->CMAC_data_hash, 16) != 0) return KIRK_DATA_HASH_INVALID;
	
		return KIRK_OPERATION_SUCCESS;
	}
	return KIRK_SIG_CHECK_INVALID; //Checks for cmd 2 & 3 not included right now
}

int kirk_CMD11(u8* outbuff, u8* inbuff, int size)
{
    if(is_kirk_initialized == 0) return KIRK_NOT_INITIALIZED;
	KIRK_SHA1_HEADER *header = (KIRK_SHA1_HEADER *)inbuff;
	if(header->data_size == 0 || size == 0) return KIRK_DATA_SIZE_ZERO;
	
    SHA1Context sha;
    libkirk_SHA1Reset(&sha);
    size <<= 4;
    size >>= 4;
	size = size < header->data_size ? size : header->data_size;
    libkirk_SHA1Input(&sha, inbuff+sizeof(KIRK_SHA1_HEADER), size);
    memcpy(outbuff, (u8*)sha.Message_Digest, 16);
    return KIRK_OPERATION_SUCCESS;
}
int kirk_CMD14(u8* outbuff, int size)
{
    if(is_kirk_initialized == 0) return KIRK_NOT_INITIALIZED;
    int i;
    u8* buf = (u8*)outbuff;
    for(i = 0; i < size; i++)
    {
          buf[i] = rand()%255;
    }
    return KIRK_OPERATION_SUCCESS;
}

int kirk_init()
{
    libkirk_AES_set_key(&aes_kirk1, kirk1_key, 128);
	is_kirk_initialized = 1;
	srand(time(0));
    return KIRK_OPERATION_SUCCESS;
}

u8* kirk_4_7_get_key(int key_type)
{
    switch(key_type)
	{
		case(0x03): return kirk7_key03; break;
		case(0x04): return kirk7_key04; break;
		case(0x05): return kirk7_key05; break;
		case(0x0C): return kirk7_key0C; break;
		case(0x0D): return kirk7_key0D; break;
		case(0x0E): return kirk7_key0E; break;
		case(0x0F): return kirk7_key0F; break;
		case(0x10): return kirk7_key10; break;
		case(0x11): return kirk7_key11; break;
		case(0x12): return kirk7_key12; break;
		case(0x38): return kirk7_key38; break;
		case(0x39): return kirk7_key39; break;
		case(0x3A): return kirk7_key3A; break;
		case(0x4B): return kirk7_key4B; break;
		case(0x53): return kirk7_key53; break;
		case(0x57): return kirk7_key57; break;
		case(0x5D): return kirk7_key5D; break;
		case(0x63): return kirk7_key63; break;
		case(0x64): return kirk7_key64; break;
		default: return (u8*)KIRK_INVALID_SIZE; break; //need to get the real error code for that, placeholder now :)
	}
}

int kirk_CMD1_ex(u8* outbuff, u8* inbuff, int size, KIRK_CMD1_HEADER* header)
{
    u8* buffer = (u8*)malloc(size);
    memcpy(buffer, header, sizeof(KIRK_CMD1_HEADER));
    memcpy(buffer+sizeof(KIRK_CMD1_HEADER), inbuff, header->data_size);
    int ret = kirk_CMD1(outbuff, buffer, size, 1);
    free(buffer);
    return ret;
}

int sceUtilsSetFuseID(u8*fuse)
{
	memcpy(fuseID, fuse, 16);
	return 0;
}

int sceUtilsBufferCopyWithRange(u8* outbuff, int outsize, u8* inbuff, int insize, int cmd)
{
    switch(cmd)
    {
		case KIRK_CMD_DECRYPT_PRIVATE: 
             if(insize % 16) return SUBCWR_NOT_16_ALGINED;
             if(kirk_CMD1(outbuff, inbuff, insize, 1) == KIRK_HEADER_HASH_INVALID) return SUBCWR_HEADER_HASH_INVALID;
             return 0;
             break;
		case KIRK_CMD_ENCRYPT_IV_0: return kirk_CMD4(outbuff, inbuff, insize); break;
		case KIRK_CMD_DECRYPT_IV_0: return kirk_CMD7(outbuff, inbuff, insize); break;
		case KIRK_CMD_PRIV_SIG_CHECK: return kirk_CMD10(inbuff, insize); break;
		case KIRK_CMD_SHA1_HASH: return kirk_CMD11(outbuff, inbuff, insize); break;
	}
	return -1;
}
