#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <arpa/inet.h>
#include <openssl/md5.h>


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

typedef struct {
    u8 zeroes[128]; // padding
    u32 imet; // "IMET"
    u8 unk[8];  // 0x0000060000000003 fixed, unknown purpose
    u32 sizes[3]; // icon.bin, banner.bin, sound.bin
    u32 flag1; // unknown
	u8 names[7][84]; // JP, EN, DE, FR, ES, IT, NL
    u8 zeroes_2[840]; // padding
    u8 crypto[16]; // MD5 of 0x40 to 0x640 in header. crypto should be all 0's when calculating final MD5
} IMET;

typedef struct {
	u32 imd5; // 'IMD5'
	u32 filesize; //size of rest of file
	u8 zeroes[8]; //padding
	u8 crypto[16]; //MD5 of rest of file
} IMD5;



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

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

u32 filesize (FILE *file)
{
	u32 curpos, endpos;
	if ( file == NULL )
	{
		return -1;
	}
	curpos = ftell ( file );
	fseek ( file, 0, 2 );
	endpos = ftell ( file );
	fseek ( file, curpos, 0 );
	return endpos;
}

void IMETSign (const char * mainfile, const char * destfile, const char * channelname, const char * iconbinfile, const char * bannerbinfile, const char * soundbinfile)
{
	IMET header;
	FILE * infile, * outfile, * bannerbin, * iconbin, * soundbin;
	u32 fsize;
	char * buffer;
	char tag [] = {'T', 'E', 'M', 'I'};
	int i, j, k;
	
	memset(header.zeroes, 0, sizeof(u8) * 128);
	
	header.imet = be32((u8 *) &tag);
	
	memset(header.unk, 0, sizeof(u8) * 8);
	header.unk[2] = 6;
	header.unk[7] = 3;
	
	iconbin = fopen(iconbinfile, "r");
	if (iconbin==NULL) 
	{
		printf ("icon.bin file error.\n");
		return;
	}
	
	bannerbin = fopen(bannerbinfile, "r");
	if (bannerbin==NULL) 
	{
		printf ("banner.bin file error.\n");
		return;
	}
	
	soundbin = fopen(soundbinfile, "r");
	if (soundbin==NULL) 
	{
		printf ("sound.bin file error.\n");
		return;
	}
	
	header.sizes[0] = htonl(filesize(bannerbin));
	header.sizes[1] = htonl(filesize(bannerbin));
	header.sizes[2] = htonl(filesize(soundbin));
	
	fclose(bannerbin);
	fclose(iconbin);
	fclose(soundbin);
	
	header.flag1 = 0x00;
	
	memset(header.names, 0, sizeof(u8) * 84 * 7);
	
	for(i=0, k=0; i<strlen(channelname); i++)
	{
		for(j=0; j<8; j++)
		{
			header.names[j][k]=0;
		}
		k++;
		for(j=0; j<8; j++)
		{
			header.names[j][k]=channelname[i];
		}
		k++;
	}
	
	memset(header.zeroes_2, 0, sizeof(u8) * 840);
	memset(header.crypto, 0, sizeof(u8) * 16);
	MD5(((u8 *) &header) + 0x40, 0x600, (u8 *) header.crypto);
	
	infile = fopen(mainfile, "r");
	if (infile==NULL) 
	{
		printf ("Input file error.\n");
		return;
	}
	fseek (infile , 0 , SEEK_END);
	fsize = ftell (infile);
	rewind (infile);
	buffer = (char *) malloc( sizeof(char) * fsize);
	if (buffer == NULL) 
	{
		printf ("Buffer error.\n"); 
		return;
	}
	fread (buffer, 1, fsize, infile);
	fclose (infile);
	
	outfile = fopen(destfile, "w");
	if (outfile==NULL) 
	{
		printf ("Output file error.\n");
		return;
	}
	
	fwrite (&header, 1, sizeof(IMET), outfile );
	fwrite (buffer, 1, fsize, outfile );
	
	fclose(outfile);
	free (buffer);
}

void IMD5Sign (const char * mainfile, const char * destfile)
{
	IMD5 header;
	FILE * infile, * outfile;
	u32 fsize;
	unsigned char * buffer;
	char tag [] = {'5', 'D', 'M', 'I'};
	
	infile = fopen(mainfile, "r");
	if (infile==NULL) 
	{
		printf ("Input file error.\n");
		return;
	}
	
	fseek (infile , 0 , SEEK_END);
	fsize = ftell (infile);
	rewind (infile);
	
	buffer = (unsigned char *) malloc( sizeof(char) * fsize);
	if (buffer == NULL) 
	{
		printf ("Buffer error.\n"); 
		return;
	}
	
	fread (buffer, 1, fsize, infile);
	printf("Read input file.\n");
	fclose (infile);
	
	header.imd5 = be32((u8 *) &tag);
	header.filesize = htonl(fsize);
	memset(header.zeroes, 0, sizeof(u8) * 8);
	MD5(buffer, fsize, (u8 *) header.crypto);
	
	outfile = fopen(destfile, "w");
	if (outfile==NULL) 
	{
		printf ("Output file error.\n");
		return;
	}
	
	fwrite (&header, 1, sizeof(IMD5), outfile );
	fwrite (buffer, 1, fsize, outfile );
	printf("Output new file with IMD5 header.\n");
	
	fclose(outfile);
	free (buffer);
}


int main (int argc, const char *argv[]) 
{
    if(argc < 4)
	{
		printf("Usage: bannersigner -IMET/IMD5 <infile> <outfile> [\"CHANNELNAME\" icon.bin banner.bin sound.bin]\n-IMET = IMET header, used in 00000000.app\n-IMD5 = IMD5 header, used in banner.bin, sound.bin, and icon.bin\n[\"CHANNELNAME\" banner.bin icon.bin sound.bin] = Required for -IMET, otherwise ignored.\n");
		return 0;
	}
	
	if(strncmp(argv[1], "-IMET", 5)==0) //IMET
	{
		IMETSign(argv[2], argv[3], argv[4], argv[5], argv[6], argv[7]);
	}
	else if(strncmp(argv[1], "-IMD5", 5)==0) //IMD5 STUFF
	{
		IMD5Sign(argv[2], argv[3]);
	}
	else
	{
		printf("Usage: bannersigner -IMET/IMD5 <infile> <outfile> [CHANNELNAME icon.bin banner.bin sound.bin]\n-IMET = IMET header, used in 00000000.app\n-IMD5 = IMD5 header, used in banner.bin, sound.bin, and icon.bin\n[CHANNELNAME banner.bin icon.bin sound.bin] = Required for -IMET, otherwise ignored.\n");
		return 0;
	}
	
    return 0;
}