#include <iostream>
#include <windows.h>
#include <fstream>
using namespace std;

int main()
{
	cout << "=== Analyzing ArrayScramble.exe ===" << endl;
	fstream inFile;
	ofstream oFile;

	// Hold DOS header
	IMAGE_DOS_HEADER img_dos_header;
	// Hold File header
	IMAGE_FILE_HEADER img_file_header;
	// "Optional" header
	IMAGE_OPTIONAL_HEADER img_opt_header;

	inFile.open("ArrayScramble.exe", ios::binary | ios::ate | ios::in | ios::out);
	if (!inFile.is_open())
	{
		cout << "Failed to open executable." << endl;
		goto quit;
	}

	inFile.seekg(0, ios::beg);
	inFile.read(reinterpret_cast<char *>(&img_dos_header), sizeof(IMAGE_DOS_HEADER));
	if (img_dos_header.e_magic != IMAGE_DOS_SIGNATURE)
	{
		cout << "PE format not recognized." << endl;
		goto quit;
	}

	inFile.seekg(img_dos_header.e_lfanew, ios::beg);
	long type;
	inFile.read(reinterpret_cast<char *>(&type), 4);
	if (type != IMAGE_NT_SIGNATURE)
	{
		cout << "PE Signature not recognized." << endl;
		goto quit;
	}

	inFile.read(reinterpret_cast<char *>(&img_file_header), sizeof(IMAGE_FILE_HEADER));

	cout << "Target CPU: ";
	if (img_file_header.Machine == IMAGE_FILE_MACHINE_I386)
		cout << "I386";
	else if (img_file_header.Machine == IMAGE_FILE_MACHINE_IA64)
		cout << "IA64";
	else
		cout << "Unknown";
	cout << endl;

	cout << "Type: ";
	if (img_file_header.Characteristics & IMAGE_FILE_DLL)
		cout << "DLL";
	else if (img_file_header.Characteristics & IMAGE_FILE_EXECUTABLE_IMAGE)
		cout << "EXE";
	else
		cout << "Unknown";
	cout << endl;

	inFile.read(reinterpret_cast<char *>(&img_opt_header), sizeof(IMAGE_OPTIONAL_HEADER));
	cout << "Subsystem: ";
	switch (img_opt_header.Subsystem)
	{
		case IMAGE_SUBSYSTEM_WINDOWS_GUI:
			cout << "Windows";
			break;
		case IMAGE_SUBSYSTEM_WINDOWS_CUI:
			cout << "Console";
			break;
		case IMAGE_SUBSYSTEM_NATIVE:
			cout << "Native";
			break;
	}
	cout << endl;

	int img_section_start = static_cast<int>(inFile.tellg());
	// Read in sections ^_^
	IMAGE_SECTION_HEADER *img_section_hdrs = new IMAGE_SECTION_HEADER[img_file_header.NumberOfSections];
	for (int i = 0; i < img_file_header.NumberOfSections; i++)
	{
		inFile.read(reinterpret_cast<char *>(&img_section_hdrs[i]), sizeof(IMAGE_SECTION_HEADER));
	}

	char *codeData;

	cout << "Sections: " << endl;
	for (int i = 0; i < img_file_header.NumberOfSections; i++)
	{
		cout << "\t" << img_section_hdrs[i].Name << endl;
		if (img_section_hdrs[i].Characteristics & IMAGE_SCN_CNT_CODE)
		{
			cout << "\t\t" << "CODE SECTION! Length: " << img_section_hdrs[i].SizeOfRawData << endl;

			codeData = new char[img_section_hdrs[i].SizeOfRawData];
			inFile.seekg(img_section_hdrs[i].PointerToRawData, ios::beg);
			inFile.read(codeData, img_section_hdrs[i].SizeOfRawData);

			oFile.open("test.txt", ios::binary | ios::out | ios::trunc);
			oFile.write(codeData, img_section_hdrs[i].SizeOfRawData);
			oFile.close();
		}
	}

	delete[] img_section_hdrs;
	delete[] codeData;

	inFile.close();

quit:
	system("pause");
	return 0;
}