/* wadpacker v0.1 by BlaZe
   sorry for the horrible code, it got
   hacked together in an afternoon and 
   did not bother to clean it up yet ;)
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <arpa/inet.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <openssl/sha.h>
#include <openssl/aes.h>

#define	WAD_HEADER_TYPE	"\x49\x73\x00\x00"

#define round_up(x,n)   (-(-(x) & -(n)))
#define round64(x)      round_up(x,0x40)
#define	round16(x)	round_up(x,0x10)

typedef unsigned char u8;
typedef unsigned short u16;
typedef unsigned int u32;
typedef unsigned long long u64;

typedef struct {
        u32 hlen;
        u32 htype;
        u32 cert_len;
        u32 zeroes;
        u32 ticket_len;
        u32 tmd_len;
        u32 app_len;
        u32 opening_len;

} WadHeader;

typedef struct {
	u8 cid[8];
	u8 sha1[20];
} cid_st;

//u8 common_key[0x10];

extern unsigned char global_ticket[676];

u32 be32(u8 *d)
{
	return (d[3] | (d[2] << 8) | (d[1] << 16) | (d[0] << 24));
}

u16 be16(u8 *d)
{
	return ( d[1] | (d[0] << 8) );
}

u64 be64(u8 *d)
{
	return ((u64)be32(d) << 32 | be32(d+4));
}

void hexdump(u8 *data, u32 len)
{
	int i;

	for(i=0; i<len; i++) {
		printf("%02x ",(u8)data[i]);
		if (((i+1)%8) == 0) printf(" ");
		if (((i+1)%16) == 0) printf("\n");
	}
	printf("\n");

	return;
}


int trucha_sign_ticket(u8 *ticket, u32 ticket_len)
{
        u8 sha1[20];
        u16 i;

        for(i=0; i<65535; i++) {
                memcpy(ticket+0x1f1, &i, sizeof(u16));
                SHA1(ticket+0x140, ticket_len-0x140, sha1);
                if (sha1[0] == 0x00) break;
        }

        if (i == 65535)
                return 0;
        else    return 1;
}


int write_data(FILE *f, u8 *data, int size)
{
	int i;

	fwrite(data, size, 1, f);
	for (i=size; round64(size)-i >0; i++)
		fputc(0, f);

	return 1;
}

int file_getsize(char *f)
{
	struct stat st;

	if (stat(f, &st) < 0)
		return -1;
	
	return st.st_size;
}

int file_sha1(char *f, u8 *sha1, u32 size)
{
	FILE *fp;
	char *buf;

	fp = fopen(f, "rb");
	if (fp == NULL)
		return 0;

	buf = malloc(size);
	if (buf == NULL)
		return 0;

	fread((char*)buf, size, 1, fp);

	SHA1((const unsigned char*)buf, (unsigned long)size, (unsigned char*)sha1);
	
	free(buf);
	fclose(fp);

	return 1;
}

int get_common_key(u8 *common_key)
{
	char path[512];
	char *home;
	FILE *fk;
	
	home = getenv("HOME");
	sprintf(path, "%s/wii/common-key", home);
	fk = fopen(path, "rb");
	if (fk == NULL)
		return 0;
	
	fread((u8*)common_key, 0x10, 1, fk);
	
	fclose(fk);
	return 1;
}

int get_title_key(u8 *ticket, u8 *title_id, u8 *title_key)
{
	u8 iv[0x10];
	u8 common_key[0x10];
	AES_KEY k;

	if (get_common_key(common_key) != 1)
		return 0;
		
	//printf("Decrypting title key ...");
	memset((u8*)iv, 0, 0x10);
	memcpy((u8*)iv, title_id, 0x8);
	//printf("IV:\n");
	//hexdump(iv, 0x10);
	AES_set_decrypt_key(common_key, 128, &k);
	AES_cbc_encrypt(ticket+0x1bf, title_key, 0x10, &k, iv, AES_DECRYPT);
	//printf("DECRYPTED TITLE KEY: \n");
	//hexdump(title_key, 0x10);
	//printf(" done.\n");

	return 1;
}

FILE *fopen_shared(u8 *sha1)
{
	FILE *f, *fc;
	char path[512];
	char cid[9];
	cid_st c;
	char *home;
	
	if ((home = getenv("HOME")) == NULL) {
		return NULL;
	}

	sprintf(path, "%s/wii/shared1/content.map",home);
	fc = fopen(path, "rb");
	if (fc == NULL)
		return NULL;
	
	while(!feof(fc)) {
		fread((u8*)&c, sizeof(c), 1, fc);
		if (memcmp((u8*)c.sha1, sha1, 20) == 0) {
			fclose(fc);
			cid[8] = '\0';
			memcpy((char*)cid, (char*)c.cid, 8);
			sprintf(path, "%s/wii/shared1/%s.app",home,cid);
			fc = fopen(path, "rb");
			return fc;
		}
	}
	fclose(fc);
	return NULL;
}


int tik_patch(FILE *f, u8 *enc_title_key, u8 *title_id) 
{
	printf("Patching ticket...");
	fflush(stdout);
	memcpy((u8*)global_ticket+0x1bf, enc_title_key, 0x10);
	memcpy((u8*)global_ticket+0x1dc, title_id, 0x8);
	if (trucha_sign_ticket(global_ticket, 0x2a4) == 0)
		return 0;

	write_data(f, global_ticket, 0x2a4);

	printf(" done.\n\n");

	return 1;
}

int create_wad(FILE *f, char *tmd_path, u8 tkpatch)
{
	WadHeader header;
	AES_KEY k;
	int i;
	FILE *fw, *fr;
	u8 title_key[0x10], iv[0x10], enc_title_key[0x10];
	u8 title_id[0x8];
	u32 tmd_size;
	u8 tik[0x2c0];
	u32 tik_size;
	u8 cetk[0xa00];
	u32 cetk_size;
	u8 *tmd;
	u8 title[5];
	char tik_path[15];
	u16 nfiles;
	u32 total_app_size;
	u32 cid;
	char file[15], banner_file[15];
	char cetk_path[512];
	u32 banner_len;
	u8 *buf;
	char *home;


	tmd_size = file_getsize(tmd_path);
	tmd = malloc(tmd_size);
	if (tmd == NULL) return 0;

	fr = fopen(tmd_path, "rb");
	if (fr == NULL) {
		fprintf(stderr, "tmd error\n");
		return 0;
	}
	fread(tmd, tmd_size, 1, fr);
	fclose(fr);
	memcpy((u8*)title_id, tmd+0x18c, 0x8);
	memcpy(title, tmd+0x190, 4);
	title[4] = '\0';

	for(i=0; i<4; i++) {
		sprintf(tik_path+2*i, "%02x",title[i]);
	}
	sprintf(tik_path+2*i, "%s",".tik");

	tik_size = file_getsize(tik_path);
	fr = fopen(tik_path, "rb");
	if (fr == NULL) {
		fprintf(stderr, "tik error: %s\n",tik_path);
		return 0;
	}
	fread(tik, tik_size, 1, fr);
	fclose(fr);
	memcpy((u8*)enc_title_key, tik+0x1bf, 0x10);

	home = getenv("HOME");
	sprintf(cetk_path, "%s/wii/cetk", home);
	cetk_size = file_getsize(cetk_path);
	fr = fopen(cetk_path, "rb");
	if (fr == NULL) {
		fprintf(stderr, "cetk error: %s\n",cetk_path);
		return 0;
	}
	fread(cetk, cetk_size, 1, fr);
	fclose(fr);

	if (get_title_key(tik, tmd+0x18c, title_key) != 1) {
		fprintf(stderr, "key error\n");
		return 0;
	}
	AES_set_encrypt_key(title_key, 128, &k);

	fseek(f, 0, SEEK_SET);
	/* blank header for now */
	for(i=0; i<0x40; i++)
		fputc(0x00, f);
	
	write_data(f, cetk, cetk_size);

	// to patch or not to patch...
	if (tkpatch) {
		if (tik_patch(f,enc_title_key,title_id) == 0) {
			fprintf(stderr, "patch error\n");
			return 0;
		}
	} else {
		write_data(f, tik, tik_size);
	}
	write_data(f, tmd, tmd_size);

	total_app_size = 0;
	memcpy((u8*)&nfiles, tmd+0x1de, 2);
	nfiles = ntohs(nfiles);
	tmd_size = nfiles*0x24 + 0x1e4; // tmd size without padding, correct size for header
	printf("Packing %d files ...\n",nfiles);
	for(i=0; i<nfiles; i++) {
		u32 fsize;
		u16 type;
		memset((u8*)iv, 0, 0x10);
		memcpy((u8*)iv, (u8*)tmd+0x01e8+0x24*i, 0x2);
		memcpy((u8*)&cid, (u8*)tmd+0x1e4+0x24*i, 0x4);
		memcpy((u8*)&type, (u8*)tmd+0x01ea+0x24*i, 0x2);
		cid = ntohl(cid);
		AES_set_encrypt_key(title_key, 128, &k);
		memset((char*)file, 0, 10);
		sprintf(file, "%08x.app", cid);
		printf("\t#%d Adding %s %04x\n",i,file,ntohs(type));

		fsize = be32((u8*)tmd+0x1f0+0x24*i);
		if (i == 0) {
			sprintf(banner_file, "%08x.app", cid);
			banner_len = fsize;
		}
		
		if (ntohs(type) & 0x8000) { 
			// shared content
			fr = fopen_shared((u8*)tmd+0x1f4+0x24*i);
		} else {
			// not shared
			fr = fopen(file, "rb");
		}
		if (fr == NULL) {
			fprintf(stderr, "Error reading file %s\n",file);
			return 0;
		}
		buf = (u8*)malloc(round64(fsize));
		if (buf == NULL)
			return 0;

		memset(buf, 0, round64(fsize));
		fread((u8*)buf, fsize, 1, fr);
		AES_cbc_encrypt(buf, buf, round16(fsize), &k, iv, AES_ENCRYPT);
		fwrite((u8*)buf, round64(fsize), 1, f);
		free(buf);
		fclose(fr);
		total_app_size += round64(fsize);
	}
	fr = fopen(banner_file, "rb");
	if (fr == NULL) {
		fprintf(stderr, "Error reading banner file %s, WTf?\n",banner_file);
		return 0;
	}
	buf = (u8*)malloc(round64(banner_len));
	if (buf == NULL)
        	return 0;

	memset(buf, 0, banner_len);
	fread((u8*)buf, banner_len, 1, fr);
	fwrite((u8*)buf, banner_len, 1, f);
	free(buf);
	fclose(fr);

	fseek(f, 0, SEEK_SET);
	
	/* build wad header */
	memset((u8*)&header, 0, sizeof(header));
        header.hlen = htonl(0x20);
        memcpy((u8*)&(header.htype), (u8*)WAD_HEADER_TYPE, 4);
        header.cert_len = htonl(0xa00);
        header.zeroes = 0;
        header.ticket_len = htonl(0x2a4);
        header.tmd_len = htonl(tmd_size);
        header.app_len = htonl(total_app_size);
        header.opening_len = htonl((banner_len));
	write_data(f, (u8*)&header, sizeof(header));

	return 1;
}


int main(int argc, char **argv)
{
	FILE *f;
	u8 patch;
	extern unsigned char blzlogo[2107];

	fwrite(blzlogo, 2107, 1, stdout);
	printf("\n");
	printf(" WADPacker v0.1 by * BlaZe *\n");
	printf("\tCompile a set of files into a WAD file.\n");
	printf("\tOptionally, ticket-patch the resulting WAD.\n");
	printf("\n");

	if (argc < 3) {
		printf("usage:\t%s <WAD> <tmd-file> <-p for tik-patch\n",argv[0]);
		return 0;
	}

	f = fopen(argv[1], "wb");
	if (f == NULL) {
		fprintf(stderr, "Cannot open file %s\n",argv[1]);
		return 0;
	}

	if (argc > 3)
		patch = 1;
	else	patch = 0;

	if (create_wad(f, argv[2], patch) == 0) {
		printf("Errors occurred. Aborting...\n\n");
		fclose(f);
		return 0;
	}

	fclose(f);

	printf("Enjoy.\n\n");
	return 0;
}
