add_new_section.cpp

// addscn.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include <iostream>
#include <windows.h>
#include <string>
#include <memory>
#include <system_error>
#include <sstream>
#include <iomanip>

namespace pe_utils
{

	// Modern RAII wrapper for Windows handles
	class ScopedHandle
	{
		HANDLE handle_ = INVALID_HANDLE_VALUE;

	public:
		explicit ScopedHandle(HANDLE h = INVALID_HANDLE_VALUE) : handle_(h) {}
		~ScopedHandle()
		{
			if (isValid())
			{
				CloseHandle(handle_);
			}
		}

		ScopedHandle(const ScopedHandle &) = delete;
		ScopedHandle &operator=(const ScopedHandle &) = delete;

		ScopedHandle(ScopedHandle &&other) noexcept : handle_(other.handle_)
		{
			other.handle_ = INVALID_HANDLE_VALUE;
		}

		ScopedHandle &operator=(ScopedHandle &&other) noexcept
		{
			if (this != &other)
			{
				if (isValid())
				{
					CloseHandle(handle_);
				}
				handle_ = other.handle_;
				other.handle_ = INVALID_HANDLE_VALUE;
			}
			return *this;
		}

		bool isValid() const { return handle_ != INVALID_HANDLE_VALUE && handle_ != nullptr; }
		HANDLE get() const { return handle_; }

		HANDLE release()
		{
			HANDLE temp = handle_;
			handle_ = INVALID_HANDLE_VALUE;
			return temp;
		}
	};

	// RAII wrapper for mapped files
	class MappedFile
	{
		ScopedHandle file_;
		ScopedHandle mapping_;
		PBYTE view_ = nullptr;
		DWORD size_low_ = 0;
		DWORD size_high_ = 0;

	public:
		explicit MappedFile(const std::wstring &path, bool write_access = false)
		{
			file_ = ScopedHandle(CreateFileW(path.c_str(),
																			 write_access ? (GENERIC_READ | GENERIC_WRITE) : GENERIC_READ,
																			 0, nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr));
			if (!file_.isValid())
			{
				throw std::system_error(GetLastError(), std::system_category(), "Failed to open file");
			}

			size_low_ = GetFileSize(file_.get(), &size_high_);
			if (size_high_ != 0)
			{
				throw std::runtime_error("Files larger than 4GB are not supported");
			}
		}

		~MappedFile()
		{
			if (view_)
			{
				UnmapViewOfFile(view_);
			}
		}

		void mapReadOnly()
		{
			mapping_ = ScopedHandle(CreateFileMapping(file_.get(), nullptr, PAGE_READONLY, 0, 0, nullptr));
			if (!mapping_.isValid())
			{
				throw std::system_error(GetLastError(), std::system_category(), "Failed to create file mapping");
			}

			view_ = static_cast<PBYTE>(MapViewOfFile(mapping_.get(), FILE_MAP_READ, 0, 0, 0));
			if (!view_)
			{
				throw std::system_error(GetLastError(), std::system_category(), "Failed to map view of file");
			}
		}

		void mapReadWrite(DWORD new_size)
		{
			mapping_ = ScopedHandle(CreateFileMapping(file_.get(), nullptr, PAGE_READWRITE, 0, new_size, nullptr));
			if (!mapping_.isValid())
			{
				throw std::system_error(GetLastError(), std::system_category(), "Failed to create file mapping");
			}

			view_ = static_cast<PBYTE>(MapViewOfFile(mapping_.get(), FILE_MAP_READ | FILE_MAP_WRITE, 0, 0, 0));
			if (!view_)
			{
				throw std::system_error(GetLastError(), std::system_category(), "Failed to map view of file");
			}
		}

		PBYTE getView() const { return view_; }
		DWORD getSize() const { return size_low_; }
	};

	// Alignment utilities
	constexpr DWORD alignDown(DWORD x, DWORD align)
	{
		// Precondition: align must be a power of 2
		return (x / align) * align; // Integer division rounds down automatically
	}

	constexpr DWORD alignUp(DWORD x, DWORD align)
	{
		// Precondition: align must be a power of 2
		return ((x + align - 1) / align) * align;
	}

	class PEFile
	{
		std::unique_ptr<MappedFile> mapped_file_;
		WORD number_of_sections_;
		DWORD section_alignment_;
		DWORD file_alignment_;

	public:
		explicit PEFile(const std::wstring &path)
		{
			mapped_file_ = std::make_unique<MappedFile>(path, true);
			validatePEFile();
		}

		void appendSection(const std::string &name, DWORD virtual_size, DWORD characteristics)
		{
			auto *dos_header = reinterpret_cast<PIMAGE_DOS_HEADER>(mapped_file_->getView());
			auto *nt_headers = reinterpret_cast<PIMAGE_NT_HEADERS>(mapped_file_->getView() + dos_header->e_lfanew);

			WORD size_of_optional_header = nt_headers->FileHeader.SizeOfOptionalHeader;
			auto *file_header = &(nt_headers->FileHeader);
			auto *first_section = reinterpret_cast<PIMAGE_SECTION_HEADER>(
					reinterpret_cast<UINT_PTR>(file_header) + sizeof(IMAGE_FILE_HEADER) + size_of_optional_header);

			// Check space for new section
			auto *first_byte_of_section_data = mapped_file_->getView() + first_section->PointerToRawData;
			size_t available_space = first_byte_of_section_data -
															 (reinterpret_cast<BYTE *>(&first_section[number_of_sections_]));

			if (available_space < sizeof(IMAGE_SECTION_HEADER))
			{
				throw std::runtime_error("No room for new section header");
			}

			// Calculate new size and create writable mapping first
			DWORD new_size = alignUp(mapped_file_->getSize() + virtual_size, file_alignment_);
			mapped_file_->mapReadWrite(new_size);

			// After mapping for write access, we need to get the pointers again
			dos_header = reinterpret_cast<PIMAGE_DOS_HEADER>(mapped_file_->getView());
			nt_headers = reinterpret_cast<PIMAGE_NT_HEADERS>(mapped_file_->getView() + dos_header->e_lfanew);
			file_header = &(nt_headers->FileHeader);
			first_section = reinterpret_cast<PIMAGE_SECTION_HEADER>(
					reinterpret_cast<UINT_PTR>(file_header) + sizeof(IMAGE_FILE_HEADER) + size_of_optional_header);

			// Add new section header
			auto *new_section = &first_section[number_of_sections_];
			auto *last_section = &first_section[number_of_sections_ - 1];

			// Now we can safely write to the memory
			std::memset(new_section, 0, sizeof(IMAGE_SECTION_HEADER));
			std::memcpy(&new_section->Name, name.c_str(), std::min<size_t>(name.length(), 8));

			new_section->Misc.VirtualSize = virtual_size;
			new_section->VirtualAddress = alignUp(
					last_section->VirtualAddress + last_section->Misc.VirtualSize,
					section_alignment_);
			new_section->SizeOfRawData = alignUp(virtual_size, file_alignment_);
			new_section->PointerToRawData = mapped_file_->getSize();
			new_section->Characteristics = characteristics;

			// Update headers
			number_of_sections_++;
			nt_headers->FileHeader.NumberOfSections = number_of_sections_;
			nt_headers->OptionalHeader.SizeOfImage = alignUp(
					new_section->VirtualAddress + new_section->Misc.VirtualSize,
					section_alignment_);

			// Zero initialize new section
			std::memset(mapped_file_->getView() + new_section->PointerToRawData,
									0, new_section->SizeOfRawData);

			std::stringstream ss;
			ss << "Section added successfully:\n"
				 << "  Raw data offset: 0x" << std::hex << new_section->PointerToRawData << "\n"
				 << "  Virtual size: 0x" << std::hex << virtual_size << "\n"
				 << "  RVA: 0x" << std::hex << new_section->VirtualAddress;
			std::cout << ss.str() << std::endl;
		}

	private:
		void validatePEFile()
		{
			mapped_file_->mapReadOnly();

			auto *dos_header = reinterpret_cast<PIMAGE_DOS_HEADER>(mapped_file_->getView());
			if (dos_header->e_magic != IMAGE_DOS_SIGNATURE)
			{
				throw std::runtime_error("Invalid PE file: DOS signature mismatch");
			}

			auto *nt_headers = reinterpret_cast<PIMAGE_NT_HEADERS>(mapped_file_->getView() + dos_header->e_lfanew);
#ifdef _WIN64
			constexpr WORD EXPECTED_MACHINE = IMAGE_FILE_MACHINE_AMD64;
#else
			constexpr WORD EXPECTED_MACHINE = IMAGE_FILE_MACHINE_I386;
#endif

			if (nt_headers->Signature != IMAGE_NT_SIGNATURE ||
					nt_headers->FileHeader.Machine != EXPECTED_MACHINE)
			{
				throw std::runtime_error("Invalid PE file: NT headers invalid");
			}

			number_of_sections_ = nt_headers->FileHeader.NumberOfSections;
			section_alignment_ = nt_headers->OptionalHeader.SectionAlignment;
			file_alignment_ = nt_headers->OptionalHeader.FileAlignment;
		}
	};

} // namespace pe_utils

int wmain(int argc, wchar_t *argv[])
{
	try
	{
		if (argc < 5)
		{
			std::wcout << L"USAGE: " << argv[0]
								 << L" <path to PE file> <section name> <VirtualSize> <Characteristics>\n\n"
								 << L"VirtualSize: decimal (5021) or hex (0x12c)\n"
								 << L"Characteristics: hex (0xC0000040) or 'text'/'data'/'rdata'\n"
								 << L"  text:  0x60000020 (CODE|EXECUTE|READ)\n"
								 << L"  data:  0xC0000040 (INITIALIZED_DATA|READ|WRITE)\n"
								 << L"  rdata: 0x40000040 (INITIALIZED_DATA|READ)\n"
#ifdef _WIN64
								 << L"\nNote: 64-bit executables only\n";
#else
								 << L"\nNote: 32-bit executables only\n";
#endif
			return EXIT_SUCCESS;
		}

		// Parse parameters
		std::wstring path = argv[1];
		std::string section_name;
		int required_size = WideCharToMultiByte(CP_UTF8, 0, argv[2], -1, nullptr, 0, nullptr, nullptr);
		section_name.resize(required_size);
		WideCharToMultiByte(CP_UTF8, 0, argv[2], -1, &section_name[0], required_size, nullptr, nullptr);

		// Parse virtual size
		DWORD virtual_size = 0;
		std::wstring size_str = argv[3];
		if (size_str.substr(0, 2) == L"0x")
		{
			virtual_size = wcstoul(size_str.c_str() + 2, nullptr, 16);
		}
		else
		{
			virtual_size = wcstoul(size_str.c_str(), nullptr, 10);
		}

		// Parse characteristics
		DWORD characteristics = 0;
		std::wstring char_str = argv[4];
		if (char_str.substr(0, 2) == L"0x")
		{
			characteristics = wcstoul(char_str.c_str() + 2, nullptr, 16);
		}
		else if (char_str == L"text")
		{
			characteristics = IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ;
		}
		else if (char_str == L"data")
		{
			characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
		}
		else if (char_str == L"rdata")
		{
			characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ;
		}

		// Process the PE file
		pe_utils::PEFile pe_file(path);
		pe_file.appendSection(section_name, virtual_size, characteristics);

		return EXIT_SUCCESS;
	}
	catch (const std::exception &e)
	{
		std::cerr << "Error: " << e.what() << std::endl;
		return EXIT_FAILURE;
	}
}