#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/bn.h>
#include <openssl/sha.h>
#include <openssl/hmac.h>
#include <openssl/evp.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include "ReadMBMLoader.h"

using namespace std;

unsigned long readBufLong(unsigned char *buffer, unsigned long offset)
{
	return (unsigned long)buffer[offset] + ((unsigned long)buffer[offset+1] << 8) + 
		((unsigned long)buffer[offset+2] << 16) + ((unsigned long)buffer[offset+3] << 24);
}

void swap(unsigned char &a, unsigned char &b) {
	unsigned char t=a;
	a=b;
	b=t;
}

void Key::SetKey(unsigned char *buffer)
{
	ID = readBufLong(buffer, 0x0);
	Type = readBufLong(buffer, 0x4);
	Rights = readBufLong(buffer, 0x8);	
	ModuloLength = readBufLong(buffer, 0xc);	
	EValue = readBufLong(buffer, 0x10);
	if (Modulo) delete [] Modulo;	
	Modulo = new unsigned char [ModuloLength];
	memcpy(Modulo, buffer+0x14, ModuloLength);
	for (int i=0;i<(ModuloLength>>1);i++) swap(Modulo[i],Modulo[ModuloLength-1-i]);	

	if (rsakey) RSA_free(rsakey);	
	rsakey = RSA_new();
	BIGNUM *n = BN_bin2bn(Modulo, ModuloLength, NULL);
	unsigned char buf[4];
	buf[0] = EValue>>24;
	buf[1] = (EValue>>16) & 0xff;
	buf[2] = (EValue>>8) & 0xff;
	buf[3] = EValue & 0xff;
	BIGNUM *e = BN_bin2bn(buf, 4, NULL);
	if (rsakey->n) BN_free(n);
	rsakey->n = n;
	if (rsakey->e) BN_free(e);
	rsakey->e = e;
}

void Key::PrintKey()
{
	int i;
	fprintf(stdout, "ID=%08x\nType=%08x\nRights=%08x\nModuloLength(in Bits)=%d\nEValue=%d\n",
		(unsigned int) ID, (unsigned int) Type, (unsigned int) Rights,
		(unsigned int) ModuloLength << 3, (unsigned int)EValue);
	for (i=0; i<ModuloLength; i++) fprintf (stdout, "%02x ", (unsigned int)(Modulo[i]));
	fprintf (stdout, "\n");
}

int Key::WritePEM(const char *fname)
{
	FILE *f = fopen(fname, "w");
	PEM_write_RSA_PUBKEY(f, rsakey);
	fclose(f);
}

void CertPK::SetCertPK(unsigned char *buffer)
{
	int i;

	if (Name) delete [] Name;
	Name = new char [13];
	memcpy(Name, buffer, 12);
	Name[12]=0;

	if (strcmp(Name, "CertPK_")) {
		fprintf(stderr, "ERROR: Invalid Signature, CertPK_ expected, found %s\n", Name);
		return;
	}

	CertType = readBufLong(buffer, 0xc);
	MinVerPK = readBufLong(buffer, 0x10);
	MinVerPPA = readBufLong(buffer, 0x14);
	MinVerRD1 = readBufLong(buffer, 0x18);
	MinVerRD2 = readBufLong(buffer, 0x1c);
	MinVerISW = readBufLong(buffer, 0x20);
	MinVerKI = readBufLong(buffer, 0x24);
	MinVerPAU = readBufLong(buffer, 0x28);
	MinVerPAS = readBufLong(buffer, 0x2c);
	WatchdogParam = readBufLong(buffer, 0x30);

	NumberKeys = readBufLong(buffer, 0x148);
	Keys = new Key *[NumberKeys+1];
	for (i=0;i<NumberKeys+1;i++) Keys[i] = new Key();
	Keys[0]->SetKey(buffer+0x34);
	for (i=1;i<NumberKeys+1;i++) Keys[i]->SetKey(buffer+0x14c+(i-1)*0x114); 

	Rights = readBufLong(buffer, 0x7c4);
	MSV = readBufLong(buffer, 0x7c8);
	MSVMask = readBufLong(buffer, 0x7cc);

	if (SignerInfo) delete [] SignerInfo;
	SignerInfo = new unsigned char [0x10];
	memcpy(SignerInfo, buffer + 0x840,0x10);
	if (SignatureInfo) delete [] SignatureInfo;
	SignatureInfo = new unsigned char [0x10];
	memcpy(SignatureInfo, buffer + 0x850,0x10);	
}

void CertPK::PrintCertPK()
{
	if (!Keys) return;
	int i;
	fprintf(stdout, "Name: %s\n", Name);
	fprintf(stdout, "CertType=%08x\nMinVerPK=%08x\nMinVerPPA=%08x\nMinVerRD1=%08x\nMinVerRD2=%08x\nMinVerISW=%08x\nMinVerKI=%08x\nMinVerPAU=%08x\nMinVerPAS=%08x\nWatchdogParam=%08x\n",
		(unsigned int)CertType, (unsigned int)MinVerPK, (unsigned int)MinVerPPA, (unsigned int)MinVerRD1, (unsigned int)MinVerRD2, 
		(unsigned int)MinVerISW, (unsigned int)MinVerKI, (unsigned int)MinVerPAU, (unsigned int)MinVerPAS, (unsigned int)WatchdogParam);
	fprintf(stdout, "MasterKey:\n");
	Keys[0]->PrintKey();
	for (i=1;i<NumberKeys+1;i++) {
		fprintf(stdout, "Key%02d:\n", i);
		Keys[i]->PrintKey();
	}

	fprintf(stdout, "Rights=%08x\nMSV=%08x\nMSVMask=%08x\n", (unsigned int) Rights, (unsigned int) MSV, (unsigned int) MSVMask);
	for (i=0;i<0x10;i++) fprintf(stdout, "%02x ", (unsigned int)(SignerInfo[i]));
	fprintf(stdout, "\n");
	for (i=0;i<0x10;i++) fprintf(stdout, "%02x ", (unsigned int)(SignatureInfo[i]));
	fprintf(stdout, "\n");
}

int CertPK::WritePEM(const char *fname)
{
	char str[256];
	for (int i=0; i<NumberKeys+1;i++) {
		sprintf(str, "%s_%02d.pem", fname, i);
		Keys[i]->WritePEM(str);
	}
}

void CertPK::DecodePKSignature(unsigned char *buffer)
{
/*	int i;
	int dgst_size;
	unsigned char result[256];
	memset(result, 0, 256);
	dgst_size = RSA_public_decrypt(256, buffer + 0x860, result, Keys[0]->rsakey, RSA_PKCS1_PADDING);
	fprintf (stdout, "Decrypted signature using PKCS1_PADDING. Size=%d\n", dgst_size);
	for (i=0;i<256;i++) {
		fprintf (stdout, "%02x ", result[i]);
	}
	fprintf (stdout,"\n\n");
	memset(result, 0, 256);
	dgst_size = RSA_public_decrypt(256, buffer + 0x860, result, Keys[0]->rsakey, RSA_NO_PADDING);
	fprintf (stdout, "Decrypted signature using PKCS1_PADDING. Size=%d\n", dgst_size);
	for (i=0;i<256;i++) {
		fprintf (stdout, "%02x ", result[i]);
	}
	fprintf (stdout,"\n\n");*/
	int start, len, i;
	unsigned char *signature = new unsigned char [256];
	memcpy(signature, buffer+0x860, 256);
	for (start=0; start<0x860; start++) {
		for (len=256; len<0x861-start; len+=8) {
			if (RSA_verify(NID_sha1, buffer + start, len, signature, 0x100, Keys[0]->rsakey))
			{
				fprintf(stdout, "signature verified start=%08x len=%08x\n", start, len);
			}
		}
	}
	for (i=0;i<(256>>1);i++) swap(signature[i],signature[256-1-i]);	
	for (start=0; start<0x860; start++) {
		for (len=256; len<0x861-start; len+=8) {
			if (RSA_verify(NID_sha1, buffer + start, len, signature, 0x100, Keys[0]->rsakey))
			{
				fprintf(stdout, "reverse signature verified start=%08x len=%08x\n", start, len);
			}
		}
	}
	delete [] signature;
}

int main (int argc, char **argv)
{
	unsigned char *buffer = NULL;
	FILE *f = fopen("mbmloader_cert", "rb");
	if (f) {
		int pos;
		int end;

		pos = ftell(f);
		fseek(f, 0, SEEK_END);
		end = ftell (f);
		fseek(f, pos, SEEK_SET);

		buffer = new unsigned char[end];
		pos = fread(buffer,1,end,f);
		if (pos!=end) return -1; 
	}
	fclose (f);

	CertPK *cert = new CertPK();
	cert->SetCertPK(buffer);
	cert->PrintCertPK();
	cert->WritePEM("milekeys");
	//cert->DecodePKSignature(buffer);
	delete cert;

	return 0;
}