#include <iostream>
#include <fstream>
#include <iomanip>

#include <Windows.h>

#define Is2power(x) (!(x & (x - 1)))
#define ALIGN_DOWN(x, align) (x & ~(align - 1))
#define ALIGN_UP(x, align) ((x & (align - 1)) ? ALIGN_DOWN(x, align) + align : x) 

int main(int argc, const char* argv[])
{
	if(argc != 2)
	{
		std::cout << "Usage: sectons.exe pe_file" << std::endl;
		return 0;
	}

	std::ifstream pefile;
	pefile.open(argv[1], std::ios::in | std::ios::binary);
	if(!pefile.is_open())
	{
		std::cout << "Can't open file" << std::endl;
		return 0;
	}

	
	pefile.seekg(0, std::ios::end);
	std::streamoff filesize = pefile.tellg();
	pefile.seekg(0);

	IMAGE_DOS_HEADER dos_header;
	pefile.read(reinterpret_cast<char*>(&dos_header), sizeof(IMAGE_DOS_HEADER));
	if(pefile.bad() || pefile.eof())
	{
		std::cout << "Unable to read IMAGE_DOS_HEADER" << std::endl;
		return 0;
	}

	if(dos_header.e_magic != 'ZM')
	{
		std::cout << "IMAGE_DOS_HEADER signature is incorrect" << std::endl;
		return 0;
	}

	if((dos_header.e_lfanew % sizeof(DWORD)) != 0)
	{
		std::cout << "PE header is not DWORD-aligned" << std::endl;
		return 0;
	}

	pefile.seekg(dos_header.e_lfanew);
	if(pefile.bad() || pefile.fail())
	{
		std::cout << "Cannot reach IMAGE_NT_HEADERS" << std::endl;
		return 0;
	}

	IMAGE_NT_HEADERS32 nt_headers;
	pefile.read(reinterpret_cast<char*>(&nt_headers), sizeof(IMAGE_NT_HEADERS32) - sizeof(IMAGE_DATA_DIRECTORY) * 16);
	if(pefile.bad() || pefile.eof())
	{
		std::cout << "Error reading IMAGE_NT_HEADERS32" << std::endl;
		return 0;
	}

	if(nt_headers.Signature != 'EP')
	{
		std::cout << "Incorrect PE signature" << std::endl;
		return 0;
	}

	if(nt_headers.OptionalHeader.Magic != 0x10B)
	{
		std::cout << "This PE is not PE32" << std::endl;
		return 0;
	}

	DWORD first_section = dos_header.e_lfanew + nt_headers.FileHeader.SizeOfOptionalHeader + sizeof(IMAGE_FILE_HEADER) + sizeof(DWORD) /* Signature */;
	
	pefile.seekg(first_section);
	if(pefile.bad() || pefile.fail())
	{
		std::cout << "Cannot reach section headers" << std::endl;
		return 0;
	}

	std::cout << std::hex << std::showbase << std::left;

	for(int i = 0; i < nt_headers.FileHeader.NumberOfSections; i++)
	{
		IMAGE_SECTION_HEADER header;
		pefile.read(reinterpret_cast<char*>(&header), sizeof(IMAGE_SECTION_HEADER));
		if(pefile.bad() || pefile.eof())
		{
			std::cout << "Error reading section header" << std::endl;
			return 0;
		}

		if(!header.SizeOfRawData && !header.Misc.VirtualSize)
		{
			std::cout << "Virtual and Physical sizes of section can't be 0 at the same time" << std::endl;
			return 0;
		}

		if(header.SizeOfRawData != 0)
		{
			if(ALIGN_DOWN(header.PointerToRawData, nt_headers.OptionalHeader.FileAlignment) + header.SizeOfRawData > filesize)
			{
				std::cout << "Incorrect section address or size" << std::endl;
				return 0;
			}

			DWORD virtual_size_aligned;

			if(header.Misc.VirtualSize == 0)
				virtual_size_aligned = ALIGN_UP(header.SizeOfRawData, nt_headers.OptionalHeader.SectionAlignment);
			else
				virtual_size_aligned = ALIGN_UP(header.Misc.VirtualSize, nt_headers.OptionalHeader.SectionAlignment);

			if(header.VirtualAddress + virtual_size_aligned > ALIGN_UP(nt_headers.OptionalHeader.SizeOfImage, nt_headers.OptionalHeader.SectionAlignment))
			{
				std::cout << "Incorrect section address or size" << std::endl;
				return 0;
			}
		}

		char name[9] = {0};
		memcpy(name, header.Name, 8);
		std::cout << std::setw(20) << "Section: " << name << std::endl << "=======================" << std::endl;
		std::cout << std::setw(20) << "Virtual size:" << header.Misc.VirtualSize << std::endl;
		std::cout << std::setw(20) << "Raw size:" << header.SizeOfRawData << std::endl;
		std::cout << std::setw(20) << "Virtual address:" << header.VirtualAddress << std::endl;
		std::cout << std::setw(20) << "Raw address:" << header.PointerToRawData << std::endl;

		std::cout << std::setw(20) << "Characteristics: ";
		if(header.Characteristics & IMAGE_SCN_MEM_READ)
			std::cout << "R ";
		if(header.Characteristics & IMAGE_SCN_MEM_WRITE)
			std::cout << "W ";
		if(header.Characteristics & IMAGE_SCN_MEM_EXECUTE)
			std::cout << "X ";
		if(header.Characteristics & IMAGE_SCN_MEM_DISCARDABLE)
			std::cout << "discardable ";
		if(header.Characteristics & IMAGE_SCN_MEM_SHARED)
			std::cout << "shared";

		std::cout << std::endl << std::endl;
	}

	return 0;
}