Started security.

1 files changed, 122 insertions, 23 deletions

diff --git a/client/src/injection/pe.h b/client/src/injection/pe.h
index f246108..fab41bb 100644
--- a/client/src/injection/pe.h
+++ b/client/src/injection/pe.h
@@ -3,14 +3,93 @@
 #include "../util/native.h"
 
 namespace pe {
+
+#pragma pack(push, 4)
+	struct reloc_entry_t {
+		uint16_t                    offset : 12;
+		uint16_t                    type : 4;
+	};
+
+	struct reloc_block_t {
+		uint32_t                    base_rva;
+		uint32_t                    size_block;
+		reloc_entry_t               entries[1];   // Variable length array
+
+
+		inline reloc_block_t* get_next() { return (reloc_block_t*)((char*)this + this->size_block); }
+		inline uint32_t num_entries() { return (reloc_entry_t*)get_next() - &entries[0]; }
+	};
+
+	struct image_named_import_t
+	{
+		uint16_t            hint;
+		char                name[1];
+	};
+
+#pragma pack(push, 8)
+	struct image_thunk_data_x64_t
+	{
+		union
+		{
+			uint64_t        forwarder_string;
+			uint64_t        function;
+			uint64_t        address;                   // -> image_named_import_t
+			struct
+			{
+				uint64_t    ordinal : 16;
+				uint64_t    _reserved0 : 47;
+				uint64_t    is_ordinal : 1;
+			};
+		};
+	};
+#pragma pack(pop)
+
+	struct image_thunk_data_x86_t
+	{
+		union
+		{
+			uint32_t        forwarder_string;
+			uint32_t        function;
+			uint32_t        address;                   // -> image_named_import_t
+			struct
+			{
+				uint32_t    ordinal : 16;
+				uint32_t    _reserved0 : 15;
+				uint32_t    is_ordinal : 1;
+			};
+		};
+	};
+#pragma pack(pop)
+
+	template<bool x64,
+		typename base_type = typename std::conditional<x64, image_thunk_data_x64_t, image_thunk_data_x86_t>::type>
+	struct image_thunk_data_t : base_type {};
+
+	template<bool x64, typename base_type = typename std::conditional<x64, IMAGE_NT_HEADERS64, IMAGE_NT_HEADERS32>::type>
+	struct nt_headers_t : base_type {};
+
+	struct import_t {
+		std::string name;
+		uint32_t rva;
+	};
+
+	struct section_t {
+		std::string name;
+		size_t size;
+		size_t v_size;
+		uint32_t rva;
+		uint32_t va;
+	};
+
 	class virtual_image {
 		std::unordered_map<std::string, uintptr_t> m_exports;
+		std::vector<section_t> m_sections;
 
 		IMAGE_NT_HEADERS64* m_nt;
 		uintptr_t m_base;
 	public:
 		virtual_image() : m_nt{ nullptr }, m_base{ 0 } {};
-		virtual_image(const std::string_view mod) : m_base{0}, m_nt{nullptr} {
+		virtual_image(const std::string_view mod) : m_base{ 0 }, m_nt{ nullptr } {
 			auto peb = util::peb();
 			if (!peb) return;
 
@@ -30,7 +109,7 @@ namespace pe {
 					m_base = uintptr_t(entry->DllBase);
 					auto dos = reinterpret_cast<IMAGE_DOS_HEADER*>(m_base);
 
-					m_nt = reinterpret_cast<native::nt_headers_t<true>*>(m_base + dos->e_lfanew);
+					m_nt = reinterpret_cast<nt_headers_t<true>*>(m_base + dos->e_lfanew);
 
 					parse_exports();
 					break;
@@ -38,6 +117,24 @@ namespace pe {
 			}
 		}
 
+		virtual_image(const uintptr_t base) : m_base{ base }, m_nt{ nullptr } {
+			auto dos = reinterpret_cast<IMAGE_DOS_HEADER*>(m_base);
+
+			m_nt = reinterpret_cast<nt_headers_t<true>*>(m_base + dos->e_lfanew);
+		}
+
+		void parse_sections() {
+			auto secs = IMAGE_FIRST_SECTION(m_nt);
+			const size_t n = m_nt->FileHeader.NumberOfSections;
+
+			for (size_t i = 0; i < n; i++) {
+				auto sec = secs[i];
+
+				auto name = reinterpret_cast<const char*>(sec.Name);
+				m_sections.emplace_back(section_t{ name, sec.SizeOfRawData, sec.Misc.VirtualSize, sec.PointerToRawData, sec.VirtualAddress });
+			}
+		};
+
 		void parse_exports() {
 			auto dir = m_nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
 			auto exp =
@@ -60,20 +157,11 @@ namespace pe {
 		}
 
 		auto& exports() { return m_exports; }
-		operator bool() { return m_base != 0; }
-	};
-
-	struct import_t {
-		std::string name;
-		uint32_t rva;
-	};
+		auto &base() { return m_base; }
+		auto& nt() { return m_nt; }
+		auto& sections() { return m_sections; }
 
-	struct section_t {
-		std::string name;
-		size_t size;
-		size_t v_size;
-		uint32_t rva;
-		uint32_t va;
+		operator bool() { return m_base != 0; }
 	};
 
 	template <bool x64 = false>
@@ -83,9 +171,9 @@ namespace pe {
 		std::unordered_map<std::string, std::vector<import_t>> m_imports;
 		std::vector<section_t> m_sections;
 
-		native::nt_headers_t<x64>* m_nt;
+		nt_headers_t<x64>* m_nt;
 		IMAGE_DOS_HEADER* m_dos;
-		std::vector<std::pair<uint32_t, native::reloc_entry_t>> m_relocs;
+		std::vector<std::pair<uint32_t, reloc_entry_t>> m_relocs;
 
 	public:
 		image() = default;
@@ -94,7 +182,7 @@ namespace pe {
 
 
 			m_dos = reinterpret_cast<IMAGE_DOS_HEADER*>(m_buffer.data());
-			m_nt = reinterpret_cast<native::nt_headers_t<x64>*>(m_buffer.data() + m_dos->e_lfanew);
+			m_nt = reinterpret_cast<nt_headers_t<x64>*>(m_buffer.data() + m_dos->e_lfanew);
 
 			load();
 		}
@@ -134,8 +222,8 @@ namespace pe {
 			if (!reloc_dir.Size) return;
 
 			const auto ptr = rva_to_ptr(reloc_dir.VirtualAddress);
-			auto block = reinterpret_cast<native::reloc_block_t*>(ptr);
-			
+			auto block = reinterpret_cast<reloc_block_t*>(ptr);
+
 			while (block->base_rva) {
 				for (size_t i = 0; i < block->num_entries(); ++i) {
 					auto entry = block->entries[i];
@@ -156,11 +244,11 @@ namespace pe {
 			for (uint32_t i = 0; i < table->Name; i = table->Name, ++table) {
 				auto mod_name = std::string(reinterpret_cast<char*>(rva_to_ptr(table->Name)));
 
-				auto thunk = reinterpret_cast<native::image_thunk_data_t<x64>*>(rva_to_ptr(table->OriginalFirstThunk));
+				auto thunk = reinterpret_cast<image_thunk_data_t<x64>*>(rva_to_ptr(table->OriginalFirstThunk));
 
 				auto step = x64 ? sizeof(uint64_t) : sizeof(uint32_t);
 				for (uint32_t index = 0; thunk->address; index += step, ++thunk) {
-					auto named_import = reinterpret_cast<native::image_named_import_t*>(rva_to_ptr(thunk->address));
+					auto named_import = reinterpret_cast<image_named_import_t*>(rva_to_ptr(thunk->address));
 
 					if (!thunk->is_ordinal) {
 						import_t data;
@@ -178,7 +266,7 @@ namespace pe {
 		void copy(std::vector<char>& out) {
 			out.resize(m_nt->OptionalHeader.SizeOfImage);
 
-			std::memcpy(&out[0], &m_buffer[0], 4096);
+			std::memcpy(&out[0], &m_buffer[0], m_nt->OptionalHeader.SizeOfHeaders);
 
 			for (auto& sec : m_sections) {
 				std::memcpy(&out[sec.va], &m_buffer[sec.rva], sec.size);
@@ -209,4 +297,15 @@ namespace pe {
 		auto& relocs() const { return m_relocs; }
 		auto& sections() const { return m_sections; }
 	};
+
+
+	static virtual_image &ntdll() {
+		static virtual_image img{};
+		if (!img) {
+			img = virtual_image("ntdll.dll");
+		}
+		return img;
+	}
+
+	void get_all_modules(std::unordered_map<std::string, virtual_image>& modules);
 };  // namespace pe
\ No newline at end of file