///////////////////////////////////////////////////////////
// BDP-S390 Firmware Modification Code:
//    SquashFS Sub-file Extraction from Sony File
//--------------------------------------------------------
// Copyright (C) 2013 Malcolm Stagg
///////////////////////////////////////////////////////////

#include <stdio.h>
#include <iostream>
#include <fstream>
#include <cstdlib>
#include <string.h>

//for FILE_4

using namespace std;

int findmagic(int magicnum, int countneeded, ifstream &infile) {
	int curint;
	infile.read((char *)&curint, 4);
	int i=0;
	int magiccount = 0;
	int addr1 = 0;
	while (infile) {
		if (curint == magicnum)
			magiccount++;
		else
			magiccount=0;
		if (magiccount==countneeded) {
			addr1 = infile.tellg();
			return addr1;
		}
		infile.read((char *)&curint, 4);
	}
	return -1;
}

void printBIT(int addr1, int addr2, ifstream &infile) {
	cout << "BINARY INFO TABLE FOUND AT ADDRESS " << hex << addr1 << endl;
	infile.seekg(addr1, ios::beg);
	addr2-=20;
	int i=0;
	while (addr1<addr2) {
		int partid, offsetfile, binsize, offsetpart, bininfo;
		infile.read((char *)&partid, 4);
		infile.read((char *)&offsetfile, 4);
		infile.read((char *)&binsize, 4);
		infile.read((char *)&offsetpart, 4);
		infile.read((char *)&bininfo, 4);
		cout << i << ": pid = " << hex << partid << " offsf = " << offsetfile
		<< " size = " << binsize << " offsp = " << offsetpart << " info = " << bininfo << endl;
		int curaddr = infile.tellg();
		infile.seekg(offsetfile, ios::beg);
		int begword;
		infile.read((char *)&begword, 4);
		if (begword == 0x73717368) {
			unsigned char * squashdata = (unsigned char *)malloc(binsize);
			memcpy(squashdata, (unsigned char *)&begword, 4);
			infile.read((char *)squashdata+4, binsize-4);
			char filename[256];
			sprintf(filename, "squash_%d", partid);
			ofstream squashout;
			squashout.open(filename, ofstream::binary);
			if (squashout) {
				squashout.write((char *)squashdata, binsize);
				squashout.close();
			}
			free(squashdata);
			cout << "\nSQUASHFS filesystem saved to " << filename << "\n\n";
		} else
			cout << " : " << hex << begword << endl;
		infile.seekg(curaddr, ios::beg);
		addr1 += 20;
		i++;
	}
}

int main(int argc, char * argv[]) {
	ifstream infile;
	
	if (argc != 2) {
		cout << "Usage: extractsubfiles [FILE_4]" << endl;
		exit(1);
	}
	char * infilename = argv[1];
	
	infile.open(infilename, ifstream::binary);
	
	if (!infile) {
		cout << "cannot open input file" << endl;
		exit(1);
	}
	
	int addr1, addr2;
	
	addr1 = findmagic(0x8530ABCD, 5, infile);
	
	if (addr1>=0) {
		addr2 = findmagic(0x8530EFEF, 5, infile);
		if (addr2>=0) {
			printBIT(addr1, addr2, infile);
		} else {
			cout << "BINARY INFO TABLE NOT FOUND!" << endl;
			exit(1);
		}
	} else {
		cout << "BINARY INFO TABLE NOT FOUND!" << endl;
		exit(1);
	}
	
	infile.seekg(0, ios::beg);
	addr1 = findmagic(0x50495469, 2, infile);
	
	if (addr1>=0) {
		addr2 = findmagic(0x69544950, 4, infile);
		if (addr2>=0) {
			cout << "PARTITION INFO TABLE FOUND AT ADDRESS " << hex << addr1 << endl;
			infile.seekg(addr1, ios::beg);
			addr2-=16;
			int i=0;
			int unk1, unk2;
			infile.read((char *)&unk1, 4);
			infile.read((char *)&unk2, 4);
			addr1+=8;
			cout << "unknown1 = " << hex << unk1 << " unknown2 = " << unk2 << endl;
			while (addr1<addr2) {
				int nandsize, binoffs, priv1, priv2;
				
				//nand offset?
				//partition table offset
				//part table length
				//?
				//binary info table offset //(all BITs are equal)
				//binary info table length
				//?
				//?
				
				infile.read((char *)&nandsize, 4);
				infile.read((char *)&binoffs, 4);
				infile.read((char *)&priv1, 4);
				infile.read((char *)&priv2, 4);
				cout << i << ": nandsize = " << hex << nandsize << " binoffs = " << binoffs
				<< " priv1 = " << priv1 << " priv2 = " << priv2 << endl;
				addr1 += 16;
				i++;
			}
		} else {
			cout << "PARTITION INFO TABLE NOT FOUND!" << endl;
			exit(1);
		}
	} else {
		cout << "PARTITION INFO TABLE NOT FOUND!" << endl;
		exit(1);
	}
	
	infile.seekg(0, ios::beg);
	int curint1, curint2;
	do {
		addr1 = findmagic(0x53526F58, 1, infile);
		infile.read((char *)&curint1, 4);
		infile.read((char *)&curint2, 4);
	} while (infile && curint2 != 0x586F5253);
	
	if (infile)
		cout << "Current checksum is " << hex << curint1 << endl;

	//compare to calculated checksum
	unsigned char chksum[4];
	for (int i=0;i<4;i++)
		chksum[i]=0;
	
	unsigned char c1;
	infile.seekg(0, ios::beg);
	c1 = infile.get();
	int i=0;
	while (infile) {
		i%=4;
		chksum[i] ^= c1;
		c1 = infile.get();
		i++;
	}
	infile.close();

	chksum[0] ^= 0x53^0x58;
	chksum[1] ^= 0x52^0x6F;
	chksum[2] ^= 0x6F^0x52;
	chksum[3] ^= 0x58^0x53;
	
	if (chksum[0] == 0xFF && chksum[1] == 0xFF && chksum[2] == 0xFF && chksum[3] == 0xFF)
		cout << "Checksums Match!\n";
	else
		cout << "CHECKSUM ERROR!!!!!\n";
	
	//sub_5C7884 is where partition info table is read
	
	
	//a_country_table_num2cfg
	
}