// 
// this tools is going for a random chance of guessing the
// retail AES key required to decrypt/unpack retail xex files
// 
// the havege library is also required to compile this.
// http://www.irisa.fr/caps/projects/hipsor/index.php
// 

#include <stdio.h>
#include <string.h>
#include <malloc.h>
#include "aes/aes.h"
#include "sha1/sha1.h"
#include "havege/havege.h"

#define SWAP32(x) \
	((((x)&0xFF000000) >> 24) | \
	 (((x)&0x00FF0000) >>  8) | \
	 (((x)&0x0000FF00) <<  8) | \
	 (((x)&0x000000FF) << 24) )


const char title[]	=	"Xex Retail AES key guesser v1.0";

const char usage[]	=	"Usage: XexRetail <xex filename>";

const char info[]	=	"If the correct key is found it will be saved to an XexAes.key.\n"
						"It will also be printed to console.\n\n"
						"Press Ctrl+C at any time to stop guessing.";

typedef struct CompressionInfo {
    int Reserved;
    int CompressionWindow;
    int Size1;
    int Size2;
    unsigned char Hash[20];
} CompressionInfo;

typedef struct OptionalHeaderEntry {
    int ID;
    int Data;
} OptionalHeaderEntry;

void aes_dec_cbc(aes_context *ctx, unsigned char* input, int size, unsigned char* output, unsigned char* ivec);

int main(int argc, const char* argv[])
{
	unsigned char key_retail[16];
	int* key_retail_ptr = (int*)key_retail;
	unsigned char key_session_enc[16];
	unsigned char key_session[16];
	unsigned char ivec[16];
	aes_context ctx_aes;
	sha1_context ctx_sha1;
	int i;
	FILE* fd = 0;
	unsigned char xex_magic[4];
	int xex_info_offset;
	int xex_data_offset;
	int xex_comp_offset;
	int xex_num_entries;
	CompressionInfo comp_info;
	OptionalHeaderEntry entry;
	unsigned char* enc_buff = NULL;
	unsigned char* dec_buff = NULL;
	int enc_size;
	unsigned char enc_hash[20];
	unsigned char calc_hash[20];
//	unsigned __int64 key_count;
	unsigned int key_count;

	printf("%s\n", title);
	if(argc != 2)
	{
		printf("%s\n", usage);
		return 1;
	}
	
	// open file and check it is an xex
	fd = fopen(argv[1], "rb");
	if(fd == NULL)
	{
		printf("Error opening %s\n", argv[1]);
		return 2;
	}
	fread(xex_magic, 4, 1, fd);
	if( memcmp(xex_magic, "XEX2", 4) )
	{
		printf("Error: %s is not an xex file\n", argv[1]);
		fclose(fd);
		return 3;
	}
	
	// get encrypted session key from xex
	fseek(fd, 0x10, SEEK_SET);
	fread(&xex_info_offset, 1, 4, fd);
	xex_info_offset = SWAP32(xex_info_offset);
	fseek(fd, xex_info_offset + 0x150, SEEK_SET);
    fread(&key_session_enc, 1, 16, fd);
	
	// find compression info in xex
	fseek(fd, 0x14, SEEK_SET);
	fread(&xex_num_entries, 4, 1, fd);
	xex_num_entries = SWAP32(xex_num_entries);
	xex_comp_offset = 0;
	for(i=0; i<xex_num_entries; i++)
	{
		fread(&entry, sizeof(entry), 1, fd);
		entry.ID = SWAP32(entry.ID);
		entry.Data = SWAP32(entry.Data);
		if(entry.ID == 0x3FF)
		{
			xex_comp_offset = entry.Data;
			break;
		}
	}
	if(xex_comp_offset == 0)
	{
		printf("Error: no compression info found!\n");
		fclose(fd);
		return 3;
	}
	
	// check if compression info is supported
	fseek(fd, xex_comp_offset, SEEK_SET);
	fread(&comp_info, sizeof(comp_info), 1, fd);
	comp_info.Reserved = SWAP32(comp_info.Reserved);
	if(comp_info.Reserved != 0x24)
	{
		printf("Error: invalid compression info found (%X)\n", comp_info.Reserved);
		fclose(fd);
		return 4;
	}
	
	// get compression values required for decryption attempts
	memcpy(enc_hash, comp_info.Hash, 20);
	enc_size = SWAP32(comp_info.Size2);
	enc_buff = malloc(enc_size);
	dec_buff = malloc(enc_size);
	fseek(fd, 0x08, SEEK_SET);
	fread(&xex_data_offset, 1, 4, fd);
	xex_data_offset = SWAP32(xex_data_offset);
	fseek(fd, xex_data_offset, SEEK_SET);
	fread(enc_buff, enc_size, 1, fd);
	fclose(fd);
	printf("Testing keys over %d bytes of data\n", enc_size);
	
	// try to guess
	printf("\n%s\n\n", info);
	for(key_count=0; 1; key_count++)
	{
		// randomly guess retail aes key
		key_retail_ptr[0] = ndrand();
		key_retail_ptr[1] = ndrand();
		key_retail_ptr[2] = ndrand();
		key_retail_ptr[3] = ndrand();
//		printf("trying key\n");
//		for(i=0; i<16; i++) printf("%02X ", key_retail[i]); printf("\n");
		
		// uses devkit key (can be used to test that program works properly
//		memset(key_retail, 0, 16);
		
		// use retail key aes key to decrypt xex session aes key
		aes_set_key(&ctx_aes, key_retail, 128);
		aes_decrypt(&ctx_aes, key_session_enc, key_session);
		
		// attempt xex decryption using generated session key
		memset(ivec, 0, 16);
		aes_set_key(&ctx_aes, key_session, 128);
		aes_dec_cbc(&ctx_aes, enc_buff, enc_size, dec_buff, ivec);
		
		// calculate hash over decrypted data
		sha1_starts(&ctx_sha1);
		sha1_update(&ctx_sha1, dec_buff, enc_size);
		sha1_finish(&ctx_sha1, calc_hash);
		
		// check if calculated hash matches stored hash
		// if so then the guessed key was correct!
		if( memcmp(enc_hash, calc_hash, 20) == 0)
		{
			printf("Retail Xex key found!!!!!!!!!!\n");
			for(i=0; i<16; i++) printf("%02X ", key_retail[i]); printf("\n");
			
			fd = fopen("XexRetail.key", "w+");
			fprintf(fd, "Xbox360 Xex Retail AES Key\n\n");
			for(i=0; i<16; i++) fprintf(fd, "%02X ", key_retail[i]); fprintf(fd, "\n");
			fclose(fd);
			break;
		}
		
		// let the user know how many guesses they've done
//		if((key_count % 1000) == 0) printf("tested %I64d keys\n", key_count);
		if((key_count % 10000) == 0) printf(".");
	}
	
	// done
	free(enc_buff);
	free(dec_buff);
	return 0;
}


void aes_dec_cbc(aes_context *ctx, unsigned char* input, int size, unsigned char* output, unsigned char* ivec)
{
	int offset, i;
	
	for(offset=0; offset<size; offset+=16)
	{
		aes_decrypt(ctx, &input[offset], &output[offset]);
		for(i=0; i<16; i++)
			output[offset+i] ^= ivec[i];
		memcpy(ivec, &input[offset], 16);
	}
}

