diff --git a/.clang-tidy b/.clang-tidy index cc82bf8..c2198a6 100644 --- a/.clang-tidy +++ b/.clang-tidy @@ -36,6 +36,7 @@ Checks: > -cppcoreguidelines-pro-type-union-access, -cppcoreguidelines-special-member-functions, -hicpp-avoid-c-arrays, + -hicpp-signed-bitwise, -hicpp-special-member-functions, -misc-include-cleaner, -misc-no-recursion, diff --git a/.github/workflows/clang-tidy.yml b/.github/workflows/clang-tidy.yml index 392ded1..378be2a 100644 --- a/.github/workflows/clang-tidy.yml +++ b/.github/workflows/clang-tidy.yml @@ -51,6 +51,7 @@ jobs: run: | echo "Checks: '-*'" > build/.clang-tidy echo "Checks: '-*'" > vendor/zasm/.clang-tidy + echo "Checks: '-*'" > vendor/unicorn/.clang-tidy - name: Run clang tidy run: | diff --git a/.gitmodules b/.gitmodules index 1cedb45..5bfb1dc 100644 --- a/.gitmodules +++ b/.gitmodules @@ -16,3 +16,6 @@ [submodule "vendor/common"] path = vendor/common url = https://github.com/es3n1n/common +[submodule "vendor/unicorn"] + path = vendor/unicorn + url = https://github.com/unicorn-engine/unicorn.git diff --git a/CMakeLists.txt b/CMakeLists.txt index 121c0fd..c6ad85d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -30,6 +30,11 @@ endif() # Options option(OBFUSCATOR_BUILD_TESTS "" ON) +option(UNICORN_LEGACY_STATIC_ARCHIVE "" OFF) +option(BUILD_SHARED_LIBS "" OFF) + +# Variables +set(UNICORN_ARCH aarch64) project(obfuscator LANGUAGES diff --git a/cmake.toml b/cmake.toml index 56496c6..afed966 100644 --- a/cmake.toml +++ b/cmake.toml @@ -9,6 +9,11 @@ msvc-runtime = "static" [options] OBFUSCATOR_BUILD_TESTS = true +UNICORN_LEGACY_STATIC_ARCHIVE = false +BUILD_SHARED_LIBS = false + +[variables] +UNICORN_ARCH = "aarch64" [conditions] build-tests = "OBFUSCATOR_BUILD_TESTS" diff --git a/scripts/adjust_compile_commands.py b/scripts/adjust_compile_commands.py old mode 100644 new mode 100755 index c260752..fa26652 --- a/scripts/adjust_compile_commands.py +++ b/scripts/adjust_compile_commands.py @@ -1,3 +1,4 @@ +#!/usr/bin/env python3 import json from sys import argv diff --git a/scripts/bytes_to_array.py b/scripts/bytes_to_array.py new file mode 100755 index 0000000..d43b874 --- /dev/null +++ b/scripts/bytes_to_array.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python3 + +buf = '' +while True: + try: + buf += input() + except (KeyboardInterrupt, EOFError): + break + +data_str = ', '.join([f'0x{byte:02X}' for byte in bytes.fromhex(buf)]) +print('\nconstexpr auto kData = std::to_array({%data%});'.replace('%data%', data_str)) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 9ba10f6..d86a5f8 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -19,7 +19,6 @@ target_compile_features(obfuscator-project INTERFACE if(MSVC) # msvc target_compile_options(obfuscator-project INTERFACE - "/wd4661" "/MP" ) endif() @@ -31,6 +30,16 @@ set(obfuscator-lib_SOURCES "lib/analysis/bb_decomp/jumptables.cpp" "lib/analysis/passes/misc/bb_insn_passes.cpp" "lib/config_parser/config_parser.cpp" + "lib/cont/elf/elf_image.cpp" + "lib/cont/elf/rebuilder/detail/elf_copy_sections.cpp" + "lib/cont/elf/rebuilder/detail/elf_init_header.cpp" + "lib/cont/elf/rebuilder/detail/elf_update_relocations.cpp" + "lib/cont/pe/image.cpp" + "lib/cont/pe/rebuilder/detail/copy_sections.cpp" + "lib/cont/pe/rebuilder/detail/erase_metadata.cpp" + "lib/cont/pe/rebuilder/detail/init_header.cpp" + "lib/cont/pe/rebuilder/detail/update_checksum.cpp" + "lib/cont/pe/rebuilder/detail/update_relocations.cpp" "lib/easm/assembler/assembler.cpp" "lib/easm/disassembler/disassembler.cpp" "lib/func_parser/map/map.cpp" @@ -46,16 +55,11 @@ set(obfuscator-lib_SOURCES "lib/mathop/operations/impl/xor.cpp" "lib/obfuscator/obfuscator.cpp" "lib/obfuscator/transforms/startup.cpp" - "lib/pe/pe.cpp" - "lib/pe/rebuilder/detail/copy_sections.cpp" - "lib/pe/rebuilder/detail/erase_metadata.cpp" - "lib/pe/rebuilder/detail/init_header.cpp" - "lib/pe/rebuilder/detail/update_checksum.cpp" - "lib/pe/rebuilder/detail/update_relocations.cpp" "lib/analysis/analysis.hpp" "lib/analysis/bb_decomp/bb_decomp.hpp" "lib/analysis/common/common.hpp" "lib/analysis/common/debug.hpp" + "lib/analysis/common/pass_context.hpp" "lib/analysis/common/provider.hpp" "lib/analysis/lru_reg/lru_reg.hpp" "lib/analysis/observer/observer.hpp" @@ -69,6 +73,12 @@ set(obfuscator-lib_SOURCES "lib/cli/cli.hpp" "lib/config_parser/config_parser.hpp" "lib/config_parser/structs.hpp" + "lib/cont/base.hpp" + "lib/cont/cont.hpp" + "lib/cont/elf/image.hpp" + "lib/cont/elf/rebuilder/rebuilder.hpp" + "lib/cont/pe/image.hpp" + "lib/cont/pe/rebuilder/rebuilder.hpp" "lib/easm/assembler/assembler.hpp" "lib/easm/cursor/cursor.hpp" "lib/easm/debug/debug.hpp" @@ -103,16 +113,8 @@ set(obfuscator-lib_SOURCES "lib/obfuscator/transforms/transforms/util/bcf.hpp" "lib/obfuscator/transforms/transforms/util/opaque_predicates.hpp" "lib/obfuscator/transforms/types.hpp" - "lib/pe/arch/arch.hpp" - "lib/pe/common/common.hpp" - "lib/pe/common/types.hpp" - "lib/pe/debug/debug.hpp" - "lib/pe/pe.hpp" - "lib/pe/rebuilder/detail/common.hpp" - "lib/pe/rebuilder/rebuilder.hpp" "lib/util/format.hpp" "lib/util/iterators.hpp" - "lib/util/passes.hpp" "lib/util/sections.hpp" "lib/util/structs.hpp" "lib/util/types.hpp" @@ -133,6 +135,7 @@ target_link_libraries(obfuscator-lib PUBLIC obfuscator-project zasm linux-pe + elf magic_enum LLVMDemangle es3n1n::common @@ -145,6 +148,7 @@ set(obfuscator_SOURCES "lib/analysis/bb_decomp/bb_decomp.hpp" "lib/analysis/common/common.hpp" "lib/analysis/common/debug.hpp" + "lib/analysis/common/pass_context.hpp" "lib/analysis/common/provider.hpp" "lib/analysis/lru_reg/lru_reg.hpp" "lib/analysis/observer/observer.hpp" @@ -158,6 +162,12 @@ set(obfuscator_SOURCES "lib/cli/cli.hpp" "lib/config_parser/config_parser.hpp" "lib/config_parser/structs.hpp" + "lib/cont/base.hpp" + "lib/cont/cont.hpp" + "lib/cont/elf/image.hpp" + "lib/cont/elf/rebuilder/rebuilder.hpp" + "lib/cont/pe/image.hpp" + "lib/cont/pe/rebuilder/rebuilder.hpp" "lib/easm/assembler/assembler.hpp" "lib/easm/cursor/cursor.hpp" "lib/easm/debug/debug.hpp" @@ -192,16 +202,8 @@ set(obfuscator_SOURCES "lib/obfuscator/transforms/transforms/util/bcf.hpp" "lib/obfuscator/transforms/transforms/util/opaque_predicates.hpp" "lib/obfuscator/transforms/types.hpp" - "lib/pe/arch/arch.hpp" - "lib/pe/common/common.hpp" - "lib/pe/common/types.hpp" - "lib/pe/debug/debug.hpp" - "lib/pe/pe.hpp" - "lib/pe/rebuilder/detail/common.hpp" - "lib/pe/rebuilder/rebuilder.hpp" "lib/util/format.hpp" "lib/util/iterators.hpp" - "lib/util/passes.hpp" "lib/util/sections.hpp" "lib/util/structs.hpp" "lib/util/types.hpp" @@ -233,10 +235,12 @@ if(OBFUSCATOR_BUILD_TESTS) # build-tests "tests/func_parser/map/test.map.compilers.cpp" "tests/func_parser/map/test.map.ida.cpp" "tests/func_parser/test.pdb.compilers.cpp" + "tests/transforms/test.decomp_break.cpp" "lib/analysis/analysis.hpp" "lib/analysis/bb_decomp/bb_decomp.hpp" "lib/analysis/common/common.hpp" "lib/analysis/common/debug.hpp" + "lib/analysis/common/pass_context.hpp" "lib/analysis/common/provider.hpp" "lib/analysis/lru_reg/lru_reg.hpp" "lib/analysis/observer/observer.hpp" @@ -250,6 +254,12 @@ if(OBFUSCATOR_BUILD_TESTS) # build-tests "lib/cli/cli.hpp" "lib/config_parser/config_parser.hpp" "lib/config_parser/structs.hpp" + "lib/cont/base.hpp" + "lib/cont/cont.hpp" + "lib/cont/elf/image.hpp" + "lib/cont/elf/rebuilder/rebuilder.hpp" + "lib/cont/pe/image.hpp" + "lib/cont/pe/rebuilder/rebuilder.hpp" "lib/easm/assembler/assembler.hpp" "lib/easm/cursor/cursor.hpp" "lib/easm/debug/debug.hpp" @@ -284,16 +294,8 @@ if(OBFUSCATOR_BUILD_TESTS) # build-tests "lib/obfuscator/transforms/transforms/util/bcf.hpp" "lib/obfuscator/transforms/transforms/util/opaque_predicates.hpp" "lib/obfuscator/transforms/types.hpp" - "lib/pe/arch/arch.hpp" - "lib/pe/common/common.hpp" - "lib/pe/common/types.hpp" - "lib/pe/debug/debug.hpp" - "lib/pe/pe.hpp" - "lib/pe/rebuilder/detail/common.hpp" - "lib/pe/rebuilder/rebuilder.hpp" "lib/util/format.hpp" "lib/util/iterators.hpp" - "lib/util/passes.hpp" "lib/util/sections.hpp" "lib/util/structs.hpp" "lib/util/types.hpp" @@ -313,6 +315,7 @@ if(OBFUSCATOR_BUILD_TESTS) # build-tests obfuscator-project obfuscator::lib GTest::gtest_main + unicorn ) set_target_properties(obfuscator-tests PROPERTIES diff --git a/src/bin/entry.cpp b/src/bin/entry.cpp index 9d217e3..5d2cb25 100644 --- a/src/bin/entry.cpp +++ b/src/bin/entry.cpp @@ -1,24 +1,13 @@ #include "config_parser/config_parser.hpp" +#include "cont/cont.hpp" #include "obfuscator/obfuscator.hpp" #include "obfuscator/transforms/scheduler.hpp" -#include "pe/arch/arch.hpp" -#include "pe/common/common.hpp" #include #include #include namespace { - template - void bootstrap(Img* raw_image, config_parser::Config& config) { - pe::Image image(raw_image); - - obfuscator::Instance inst(&image, config); - inst.run(); - - logger::info("startup: bye-bye"); - } - - int startup(config_parser::Config& config) try { + int startup(config_parser::Config& config) { rnd::detail::seed(config.obfuscator_config().seed); const auto& binary_path = config.obfuscator_config().binary_path; @@ -28,23 +17,25 @@ namespace { throw std::runtime_error("Got empty binary"); } - auto* img_x64 = reinterpret_cast(file->data()); - auto* img_x86 = reinterpret_cast(img_x64); - - if (!pe::common::is_valid(img_x64)) { - throw std::runtime_error("Invalid pe header"); + std::unique_ptr image; + switch (cont::get_image_type(*file)) { + case cont::ContImageType::PE: + image = std::make_unique(file->data()); + logger::info("main: PE image loaded"); + break; + case cont::ContImageType::ELF: + image = std::make_unique(file->data()); + logger::info("main: ELF image loaded"); + break; + default: + throw std::runtime_error("Got unsupported image type"); } - if (pe::arch::is_x64(img_x64)) { - bootstrap(img_x64, config); - } else { - bootstrap(img_x86, config); - } + obfuscator::Instance inst(image.get(), config); + inst.run(); + logger::info("startup: bye-bye"); return 0; - } catch (std::runtime_error& err) { - logger::critical("RUNTIME ERROR: {}", err.what()); - return 1; } } // namespace @@ -53,6 +44,9 @@ int main(const int argc, const char* argv[]) try { auto config = config_parser::from_argv(argc, argv); return startup(config); +} catch (std::exception& err) { + logger::critical("RUNTIME ERROR: {}", err.what()); + return 1; } catch (...) { logger::critical("Unknown runtime error"); return 1; diff --git a/src/cmake.toml b/src/cmake.toml index 42bd2b7..5ea3663 100644 --- a/src/cmake.toml +++ b/src/cmake.toml @@ -2,7 +2,7 @@ type = "interface" compile-features = ["cxx_std_23"] compile-definitions = ["NOMINMAX"] -msvc.compile-options = ["/wd4661", "/MP"] +msvc.compile-options = ["/MP"] [target.obfuscator-lib] alias = "obfuscator::lib" @@ -13,6 +13,7 @@ link-libraries = [ "obfuscator-project", "zasm", "linux-pe", + "elf", "magic_enum", "LLVMDemangle", "es3n1n::common", @@ -29,7 +30,7 @@ condition = "build-tests" type = "executable" sources = ["tests/**.cpp", "lib/**.hpp"] include-directories = ["tests/"] -link-libraries = ["obfuscator-project", "obfuscator::lib", "GTest::gtest_main"] +link-libraries = ["obfuscator-project", "obfuscator::lib", "GTest::gtest_main", "unicorn"] cmake-after = """ FetchContent_MakeAvailable(resources) target_compile_definitions(obfuscator-tests PRIVATE OBFUSCATOR_RESOURCES_PATH="${resources_SOURCE_DIR}") diff --git a/src/lib/analysis/analysis.cpp b/src/lib/analysis/analysis.cpp index c22bc2f..5d94e6c 100644 --- a/src/lib/analysis/analysis.cpp +++ b/src/lib/analysis/analysis.cpp @@ -1,27 +1,31 @@ #include "analysis/analysis.hpp" -#include "util/passes.hpp" +#include "analysis/common/pass_context.hpp" #include "analysis/passes/label_references.hpp" #include "analysis/passes/misc/bb_insn_passes.hpp" namespace analysis { - template - void Function::apply_passes(Img* image) { - // \note: @es3n1n: for the apply_bb/apply_insn callbacks please check out the file - // `analysis/transforms/misc/bb_insn_passes.hpp`, pass that would need to iter bb/insns - // by themselves should be inserted here - // - ::passes::apply< // - passes::bb_insn_passes_t, // - passes::label_references_t // - >(this, image); + void Function::apply_passes(std::optional image) { + /// \note: @es3n1n: + /// for the apply_bb/apply_insn callbacks please check out the file + /// `analysis/transforms/misc/bb_insn_passes.hpp`, + /// passes that would need to iter bb/insns by themselves should be inserted here + + /// Constructing the pass context + PassContext ctx = { + .image_mode = image_mode, + .image = image, + .function = this, + }; + + passes::bb_insn_passes_t::apply(ctx); + passes::label_references_t::apply(ctx); } - template - void Function::calc_range() { + void Function::calc_range() { // Reset state // - range.start = ULLONG_MAX; + range.start = std::numeric_limits::max(); range.end = nullptr; // Iterating over instructions and updating range @@ -41,5 +45,72 @@ namespace analysis { }); } - PE_DECL_TEMPLATE_CLASSES(Function); + void Function::setup(bb_decomp::Instance& decomp, std::optional image) { + bb_storage = decomp.export_blocks(); + program = decomp.export_program(); + calc_range(); + + /// Init the bb provider + bb_provider = std::make_shared(); + + /// Set RVA finder + bb_provider->set_rva_finder([storage = bb_storage.get()](const rva_t rva, bb_t*) -> std::optional> { + /// Find by RVA + auto it = std::ranges::find_if(storage->basic_blocks, [rva](auto&& bb) -> bool { + return bb->start_rva.has_value() && bb->start_rva.value() == rva; // + }); + + /// Return wrapped in optional + return it == std::end(storage->basic_blocks) ? std::nullopt : std::make_optional(*it); + }); + + /// Set VA finder + if (image.has_value()) { + bb_provider->set_va_finder([img_base = (*image)->get_image_base(), + provider = bb_provider.get()](const rva_t va, bb_t* callee) -> std::optional> { + /// Substract base and find by RVA + return provider->find_by_start_rva(va - img_base, callee); // + }); + } else { + bb_provider->set_va_finder([](const rva_t, bb_t*) -> std::optional> { + assert(false); /// Nameless functions does not have VAs + return std::nullopt; + }); + } + + /// Set Label finder + bb_provider->set_label_finder([storage = bb_storage.get()](const zasm::Label* label, bb_t*) -> std::optional> { + for (auto& bb : storage->basic_blocks) { + /// Continue if bb doesn't contain this label + if (!bb->contains_label(label->getId())) { + continue; + } + + return bb; + } + + return std::nullopt; + }); + + /// Set reference acquire callback + bb_provider->set_ref_acquire([storage = bb_storage.get()](const bb_t* bb) -> std::optional> { + /// Try to find by ptr + auto it = std::ranges::find_if(storage->basic_blocks, [bb](const auto& p) -> bool { + return p.get() == bb; // + }); + + /// Not found + if (it == std::end(storage->basic_blocks)) { + return std::nullopt; + } + + /// Found + return std::make_optional(*it); + }); + + assembler = std::make_shared(*program); + observer = std::make_shared(program, bb_storage, bb_provider); + + apply_passes(image); + } } // namespace analysis diff --git a/src/lib/analysis/analysis.hpp b/src/lib/analysis/analysis.hpp index 0a82500..93e6e1d 100644 --- a/src/lib/analysis/analysis.hpp +++ b/src/lib/analysis/analysis.hpp @@ -10,81 +10,30 @@ #include namespace analysis { - template class Function { public: - Function(Img* image, const func_parser::function_t& func): parsed_func(func) { - bb_decomp::Instance bb_decomp_inst(image, func.rva, func.size); - bb_storage = bb_decomp_inst.export_blocks(); - program = bb_decomp_inst.export_program(); - - calc_range(); - - /// Init the bb provider - bb_provider = std::make_shared(); - - /// Set RVA finder - bb_provider->set_rva_finder([storage = bb_storage.get()](const rva_t rva, bb_t*) -> std::optional> { - /// Find by RVA - auto it = std::ranges::find_if(storage->basic_blocks, [rva](auto&& bb) -> bool { - return bb->start_rva.has_value() && bb->start_rva.value() == rva; // - }); - - /// Return wrapped in optional - return it == std::end(storage->basic_blocks) ? std::nullopt : std::make_optional(*it); - }); - - /// Set VA finder - bb_provider->set_va_finder([img_base = image->raw_image->get_nt_headers()->optional_header.image_base, - provider = bb_provider.get()](const rva_t va, bb_t* callee) -> std::optional> { - /// Substract base and find by RVA - return provider->find_by_start_rva(va - img_base, callee); // - }); - - /// Set Label finder - bb_provider->set_label_finder([storage = bb_storage.get()](const zasm::Label* label, bb_t*) -> std::optional> { - for (auto& bb : storage->basic_blocks) { - /// Continue if bb doesn't contain this label - if (!bb->contains_label(label->getId())) { - continue; - } - - return bb; - } - - return std::nullopt; - }); - - /// Set reference acquire callback - bb_provider->set_ref_acquire([storage = bb_storage.get()](const bb_t* bb) -> std::optional> { - /// Try to find by ptr - auto it = std::ranges::find_if(storage->basic_blocks, [bb](const auto& p) -> bool { - return p.get() == bb; // - }); - - /// Not found - if (it == std::end(storage->basic_blocks)) { - return std::nullopt; - } - - /// Found - return std::make_optional(*it); - }); - - assembler = std::make_shared(*program); - observer = std::make_shared(program, bb_storage, bb_provider); - - apply_passes(image); + Function(cont::ImageBase* image, const func_parser::function_t& func): image_mode(image->mode()), parsed_func(func), lru_reg(LRUReg(image_mode)) { + bb_decomp::Instance bb_decomp_inst(image, func.rva, func.size); + setup(bb_decomp_inst, image); + } + + Function(const cont::ImageMode image_mode, const std::span raw_data) + : image_mode(image_mode), parsed_func(std::nullopt), lru_reg(LRUReg(image_mode)) { + /// \fixme @es3n1n: this is wrong + bb_decomp::Instance bb_decomp_inst(image_mode == cont::ImageMode::X64 ? zasm::MachineMode::AMD64 : zasm::MachineMode::I386, raw_data); + setup(bb_decomp_inst); } ~Function() = default; Function(const Function& instance) - : program(instance.program), assembler(instance.assembler), observer(instance.observer), bb_storage(instance.bb_storage), - parsed_func(instance.parsed_func), range(instance.range), lru_reg(instance.lru_reg), bb_provider(instance.bb_provider) { } + : program(instance.program), assembler(instance.assembler), observer(instance.observer), image_mode(instance.image_mode), + bb_storage(instance.bb_storage), parsed_func(instance.parsed_func), range(instance.range), lru_reg(instance.lru_reg), + bb_provider(instance.bb_provider) { } private: - void apply_passes(Img* image); + void apply_passes(std::optional image = std::nullopt); void calc_range(); + void setup(bb_decomp::Instance& decomp, std::optional image = std::nullopt); public: // A zasm program instance that contains all of our instructions @@ -93,13 +42,16 @@ namespace analysis { std::shared_ptr assembler; std::shared_ptr observer; + /// + cont::ImageMode image_mode; + // A list of split basic blocks // std::shared_ptr bb_storage; // Info about the function from the .map/.pdb files // - func_parser::function_t parsed_func; + std::optional parsed_func; // A start/end range of function // @@ -107,7 +59,7 @@ namespace analysis { // Least recently used register info // - LRUReg lru_reg; + LRUReg lru_reg; // A list of references within the image, key is the instruction and value is RVA // it referenced @@ -124,13 +76,16 @@ namespace analysis { std::shared_ptr bb_provider; }; - template - Function analyse(Img* image, const func_parser::function_t& function) { - auto result = Function(image, function); - logger::debug("analysis: analysed function {}", function); + inline Function analyse(cont::ImageBase* image, const func_parser::function_t& function) { + auto result = Function(image, function); + logger::debug("analysis: analysed function {} (range: {:#x}:{:#x})", function, result.range.start, result.range.end); if (auto size = result.range.size(); size < easm::kMaxEntryInstructionSize) { throw std::runtime_error(std::format("analysis: Minimal function size is {} bytes, got {}", easm::kMaxEntryInstructionSize, size)); } return result; } + + inline auto analyse(const cont::ImageMode image_mode, const std::span function_data) { + return Function(image_mode, function_data); + } } // namespace analysis diff --git a/src/lib/analysis/bb_decomp/bb_decomp.cpp b/src/lib/analysis/bb_decomp/bb_decomp.cpp index bd03f6c..f8c1583 100644 --- a/src/lib/analysis/bb_decomp/bb_decomp.cpp +++ b/src/lib/analysis/bb_decomp/bb_decomp.cpp @@ -3,19 +3,19 @@ #include namespace analysis::bb_decomp { - template - void Instance::collect() { + void Instance::collect() { // First of all, we should clear the previous results just in case // clear(); - // Setup the bb provider - // - const auto img_base = image_->raw_image->get_nt_headers()->optional_header.image_base; - - // Make successor proxy - bb_provider_->set_va_finder([this, img_base](const rva_t virt_addr, const bb_t* callee) { - return make_successor(virt_addr - img_base, callee); // + // Setup va finder for bb provider + /// \note @es3n1n: Base address for nameless stuff will be 0x0 + std::uintptr_t base_address = 0x0; + if (image_.has_value()) { + base_address = (*image_)->get_image_base(); + } + bb_provider_->set_va_finder([this, base_address](const rva_t virt_addr, const bb_t* callee) { + return make_successor(virt_addr - base_address, callee); // }); // Make successor proxy @@ -72,13 +72,16 @@ namespace analysis::bb_decomp { // Starting with the first basic block, and it will process others automatically // logger::info("bb_decomp: running phase 1"); - process_bb(function_start_); + /// \note @es3n1n: Nameless functions should always start at 0 + process_bb(function_start_.value_or(0)); // Expand jumptables // logger::info("bb_decomp: running phase 2"); - collect_jumptables(); - collect_jumptable_entries(); + if (image_.has_value()) { + collect_jumptables(); + collect_jumptable_entries(); + } // Splitting basic blocks (pt.1) // @@ -108,13 +111,21 @@ namespace analysis::bb_decomp { // dump(); } - template - std::shared_ptr Instance::process_bb(const rva_t rva) { + std::shared_ptr Instance::process_bb(const rva_t rva) { // Initialising stuff - // \fixme: @es3n1n: override `get_nt_headers` in `pe::Image` class - const std::uint64_t image_base = image_->raw_image->get_nt_headers()->optional_header.image_base; + // \fixme: @es3n1n: make a `get_nt_headers` func in `pe::Image` class + + /// \note @es3n1n: We always assume base address as 0 for nameless functions + /// since we use the addresses to access function's data span. Do not change. + std::uint64_t image_base = 0; + if (image_.has_value()) { + image_base = (*image_)->get_image_base(); + } const memory::address virtual_address = rva + image_base; - const std::uint8_t* data_start = image_->rva_to_ptr(static_cast(rva.inner())); + + /// Get either the pointer to function within PE, or start of the bytes passed + const std::uint8_t* data_start = + image_.has_value() ? (*image_)->rva_to_ptr(static_cast(rva.inner())) : (function_code_.value().data() + rva.inner()); // Init basic block info // @@ -159,7 +170,7 @@ namespace analysis::bb_decomp { // Ignoring anything that wouldn't affect IP // - if (!insn_desc->is_jump() && !(insn_desc->flags & UNABLE_TO_ESTIMATE_JCC)) { + if (!insn_desc->is_jump() && (insn_desc->flags & UNABLE_TO_ESTIMATE_JCC) == 0) { continue; } @@ -174,8 +185,7 @@ namespace analysis::bb_decomp { return result; } - template - void Instance::update_refs() { + void Instance::update_refs() { /// Remove stuff that was marked as to be deleted sanitize(); @@ -189,7 +199,7 @@ namespace analysis::bb_decomp { } } - /// Sum stats + /// Some stats std::size_t weird_nodes = 0; for (auto* node = program_->getHead(); node != nullptr; node = node->getNext()) { @@ -204,9 +214,10 @@ namespace analysis::bb_decomp { continue; } - // This is kinda unsafe but whatever.. + /// Store the node and instruction pinsn->node_ref = node; pinsn->ref = node->getIf(); + assert(pinsn->ref != nullptr); // This node is up2date, we can remove it as we checked it if (insns.contains(pinsn)) { @@ -230,8 +241,7 @@ namespace analysis::bb_decomp { } } - template - void Instance::split() { + void Instance::split() { /// \note @es3n1n: Looks kinda scary, but splitting 2k+ basic blocks took me ~350ms so /// i guess we'll keep it as it is (PR welcome), perhaps an interval/segment tree could be used here logger::debug("analysis: splitting BBs.."); @@ -330,8 +340,7 @@ namespace analysis::bb_decomp { } while (split_something); } - template - void Instance::sanitize() { + void Instance::sanitize() { const auto erased_bbs = std::erase_if(basic_blocks_, [](auto& basic_block) -> bool { const auto nodes_erased = std::erase_if(basic_block.second->instructions, [](auto& insn) -> bool { return insn->flags & TO_BE_REMOVED; // @@ -349,8 +358,7 @@ namespace analysis::bb_decomp { } } - template - void Instance::insert_jmps() { + void Instance::insert_jmps() { logger::debug("bb_decomp: veryfing BB intersections.."); /// Lookup for the basic blocks that for some reason aren't jumping to their successor(s) for (auto& bb : std::views::values(basic_blocks_)) { @@ -440,8 +448,7 @@ namespace analysis::bb_decomp { } } - template - void Instance::update_rescheduled_cf() { + void Instance::update_rescheduled_cf() { logger::debug("bb_decomp: updating rescheduled CF.."); /// Iterating over the all basic blocks @@ -490,8 +497,7 @@ namespace analysis::bb_decomp { } } - template - void Instance::update_tree() { + void Instance::update_tree() { logger::debug("bb_decomp: updating the BB tree.. (this could take some time)"); /// Since we splitted/merged some basic blocks, there could be some @@ -571,9 +577,8 @@ namespace analysis::bb_decomp { } } - template - void Instance::dump() { - logger::info("-- Basic blocks for function {:#x}", function_start_); + void Instance::dump() { + logger::info("-- Basic blocks for function {:#x}", function_start_.value_or(0x0)); for (auto& v : std::views::values(basic_blocks_)) { debug::dump_bb(*v); @@ -582,8 +587,7 @@ namespace analysis::bb_decomp { logger::info("-- EOF"); } - template - void Instance::dump_to_visualizer() { + void Instance::dump_to_visualizer() { const auto path = std::filesystem::path(R"(E:\local-projects\obfuscator\scripts\bb_preview\data\)"); int iter = 0; @@ -592,6 +596,4 @@ namespace analysis::bb_decomp { iter += 1; } } - - PE_DECL_TEMPLATE_CLASSES(Instance); } // namespace analysis::bb_decomp diff --git a/src/lib/analysis/bb_decomp/bb_decomp.hpp b/src/lib/analysis/bb_decomp/bb_decomp.hpp index 5e561b0..c59cdba 100644 --- a/src/lib/analysis/bb_decomp/bb_decomp.hpp +++ b/src/lib/analysis/bb_decomp/bb_decomp.hpp @@ -1,7 +1,7 @@ #pragma once #include "analysis/common/common.hpp" #include "analysis/common/provider.hpp" -#include "pe/pe.hpp" +#include "cont/base.hpp" #include #include @@ -10,21 +10,30 @@ namespace analysis::bb_decomp { /// \brief BB Decomposition instance /// \tparam Img Image - template class Instance { public: - Instance(Img* image, const rva_t rva, const std::optional function_size = std::nullopt) - : image_(image), function_start_(rva), function_size_(function_size), program_(std::make_shared(image->guess_machine_mode())), - assembler_(std::make_shared(*program_)), decoder_(easm::Decoder(image_->guess_machine_mode())), + /// Non-nameless function + explicit Instance(cont::ImageBase* image, const rva_t rva, const std::optional function_size = std::nullopt) + : machine_mode_(image->machine_mode()), image_(image), function_start_(rva), function_size_(function_size), + program_(std::make_shared(machine_mode_)), assembler_(std::make_shared(*program_)), + decoder_(easm::Decoder(machine_mode_)), bb_provider_(std::make_shared()) { + collect(); + } + + /// Nameless function + Instance(const zasm::MachineMode machine_mode, std::span function_data) + : machine_mode_(machine_mode), function_code_(function_data), program_(std::make_shared(machine_mode_)), + assembler_(std::make_shared(*program_)), decoder_(easm::Decoder(machine_mode_)), bb_provider_(std::make_shared()) { collect(); } + ~Instance() = default; Instance(const Instance& instance) - : image_(instance.image_), function_start_(instance.function_start_), function_size_(instance.function_size_), - basic_blocks_(instance.basic_blocks_), program_(instance.program_), assembler_(instance.assembler_), decoder_(instance.decoder_), - jump_tables_(instance.jump_tables_), bb_provider_(instance.bb_provider_) { } + : machine_mode_(instance.machine_mode_), image_(instance.image_), function_start_(instance.function_start_), + function_size_(instance.function_size_), basic_blocks_(instance.basic_blocks_), program_(instance.program_), assembler_(instance.assembler_), + decoder_(instance.decoder_), jump_tables_(instance.jump_tables_), bb_provider_(instance.bb_provider_) { } void collect(); void split(); @@ -58,7 +67,7 @@ namespace analysis::bb_decomp { void update_tree(); void update_rescheduled_cf(); - // jumptables shenainigans + // jumptables shenanigans void collect_jumptables(); void collect_jumptable_entries(); void expand_jumptables(); @@ -94,7 +103,7 @@ namespace analysis::bb_decomp { } [[nodiscard]] std::shared_ptr make_virtual_bb() { - auto result = std::make_shared(image_->guess_machine_mode()); + auto result = std::make_shared(machine_mode_); virtual_basic_blocks_.emplace_back(result); return result; } @@ -108,7 +117,7 @@ namespace analysis::bb_decomp { return false; } - return (rva - function_start_) >= function_size_; + return (rva - function_start_.value()) >= function_size_; } [[nodiscard]] std::shared_ptr at(const rva_t rva) { @@ -116,7 +125,7 @@ namespace analysis::bb_decomp { return it->second; } - basic_blocks_[rva] = std::make_shared(image_->guess_machine_mode()); + basic_blocks_[rva] = std::make_shared(machine_mode_); return basic_blocks_[rva]; } @@ -138,10 +147,15 @@ namespace analysis::bb_decomp { return basic_block->push_label(assembler_->getCursor(), bb_provider_.get()); } - const Img* image_ = nullptr; - rva_t function_start_ = nullptr; + /// Not set for nameless functions + zasm::MachineMode machine_mode_; + std::optional image_ = std::nullopt; + std::optional function_start_ = std::nullopt; std::optional function_size_ = std::nullopt; + /// Not set for non-nameless functions + std::optional> function_code_ = std::nullopt; + std::unordered_map> basic_blocks_; std::vector> virtual_basic_blocks_; // = without the rva @@ -154,9 +168,8 @@ namespace analysis::bb_decomp { std::shared_ptr bb_provider_; }; - template - std::vector collect(Img* image, const rva_t rva, std::optional size = std::nullopt) { - const auto inst = Instance(image, rva, size); + inline std::shared_ptr collect(cont::ImageBase* image, const rva_t rva, std::optional size = std::nullopt) { + const auto inst = Instance(image, rva, size); return inst.export_blocks(); } } // namespace analysis::bb_decomp diff --git a/src/lib/analysis/bb_decomp/jumptables.cpp b/src/lib/analysis/bb_decomp/jumptables.cpp index 3a8145b..e894881 100644 --- a/src/lib/analysis/bb_decomp/jumptables.cpp +++ b/src/lib/analysis/bb_decomp/jumptables.cpp @@ -2,13 +2,11 @@ #include "analysis/bb_decomp/bb_decomp.hpp" #include "analysis/common/debug.hpp" #include +#include /// \todo @es3n1n: Notify the linker somehow that it should erase jumptable pointers too namespace analysis::bb_decomp { - template - void Instance::collect_jumptables() { - const auto machine_mode = image_->guess_machine_mode(); - + void Instance::collect_jumptables() { for (auto& basic_block : std::views::values(basic_blocks_)) { for (std::size_t i = 0; i < basic_block->size(); ++i) { const auto& insn = basic_block->instructions.at(i); @@ -47,7 +45,7 @@ namespace analysis::bb_decomp { /// Match the dst reg if (const auto* dst_reg = prev_insn->ref->getOperandIf(0); - dst_reg == nullptr || dst_reg->getRoot(machine_mode).getId() != jmp_reg->getRoot(machine_mode).getId()) { + dst_reg == nullptr || dst_reg->getRoot(machine_mode_).getId() != jmp_reg->getRoot(machine_mode_).getId()) { return; } @@ -88,7 +86,7 @@ namespace analysis::bb_decomp { /// Get the dst reg operand const auto* const dst_op = prev_insn->ref->getOperandIf(0); - if (auto* const src_op = prev_insn->ref->getOperandIf(1); // + if (const auto* const src_op = prev_insn->ref->getOperandIf(1); // dst_op == nullptr || src_op == nullptr) { return; } @@ -98,7 +96,7 @@ namespace analysis::bb_decomp { assert(mem_index_op != nullptr); /// If matches, then yeah we found it - if (mem_index_op->getBase().getId() != dst_op->getId()) { + if (mem_index_op == nullptr || mem_index_op->getBase().getId() != dst_op->getId()) { return; } @@ -125,8 +123,11 @@ namespace analysis::bb_decomp { } } - template - void Instance::collect_jumptable_entries() { + void Instance::collect_jumptable_entries() { + if (!image_.has_value()) { + throw std::runtime_error("analysis: (jt) no image specified"); + } + /// Now we have to bruteforce the number of entries per table. /// I know, i know, it's not the proper solution; however, parsing /// the jumptables isn't that trivial of a task and it requires @@ -143,7 +144,7 @@ namespace analysis::bb_decomp { /// colliding with entries from different jump tables. for (auto& [rva, info] : jump_tables_) { /// Get the table start - auto* table = image_->template rva_to_ptr(rva); + const auto* table = (*image_)->rva_to_ptr(rva); if (table == nullptr) { throw std::runtime_error("analysis: unable to find the jump table, huh?"); } @@ -159,13 +160,12 @@ namespace analysis::bb_decomp { } /// Get the entry ptr - auto ptr = image_->template rva_to_ptr(entry); - if (!ptr) { + if (const auto* ptr = (*image_)->rva_to_ptr(entry); ptr == nullptr) { break; } /// Get the section and check if its executable - if (const auto* section = image_->rva_to_section(entry); // + if (const auto* section = (*image_)->rva_to_section(entry); // !section->characteristics.cnt_code) { break; } @@ -179,8 +179,7 @@ namespace analysis::bb_decomp { } } - template - void Instance::expand_jumptables() { + void Instance::expand_jumptables() { for (auto& [rva, info] : jump_tables_) { /// First, we should replace the /// `mov reg, [bla+bla*bla+0x1337]` with `lea reg, [bla+bla*bla]` @@ -198,11 +197,11 @@ namespace analysis::bb_decomp { /// Copy auto mem_op = *pmem_op; - auto jmp_reg = zasm::x86::Gp(pjmp_reg->getRoot(image_->guess_machine_mode()).getId()); // eax->rax (just in case) + auto jmp_reg = zasm::x86::Gp(pjmp_reg->getRoot(machine_mode_).getId()); // eax->rax (just in case) /// Remove the imm part (that points to the jump table) assert(mem_op.getDisplacement() == rva.as()); - mem_op.setBitSize(jmp_reg.getBitSize(image_->guess_machine_mode())); + mem_op.setBitSize(jmp_reg.getBitSize(machine_mode_)); mem_op.setDisplacement(0); /// Remove the base @@ -310,6 +309,4 @@ namespace analysis::bb_decomp { } } } - - PE_DECL_TEMPLATE_CLASSES(Instance); } // namespace analysis::bb_decomp \ No newline at end of file diff --git a/src/lib/analysis/common/common.hpp b/src/lib/analysis/common/common.hpp index 1e0831e..8b874e8 100644 --- a/src/lib/analysis/common/common.hpp +++ b/src/lib/analysis/common/common.hpp @@ -470,8 +470,8 @@ namespace analysis { std::reverse(iter, instructions.end()); } - [[nodiscard]] std::shared_ptr last_non_jmp_insn(zasm::Program* program = nullptr, const bool destroy_jmps = false, - const bool include_conditional_jmps = false) const { + [[nodiscard]] std::optional> last_non_jmp_insn(zasm::Program* program = nullptr, const bool destroy_jmps = false, + const bool include_conditional_jmps = false) const { auto insns = temp_insns_copy(); for (auto it = insns.rbegin(); it != insns.rend(); std::advance(it, 1)) { auto insn = *it; @@ -497,7 +497,7 @@ namespace analysis { return insn; } - throw std::runtime_error(std::format("last_non_jmp_insn: unable to query ({})", static_cast(include_conditional_jmps))); + return std::nullopt; } [[nodiscard]] bool contains_label(const zasm::Label::Id label_id) const { diff --git a/src/lib/analysis/common/pass_context.hpp b/src/lib/analysis/common/pass_context.hpp new file mode 100644 index 0000000..1fff7a1 --- /dev/null +++ b/src/lib/analysis/common/pass_context.hpp @@ -0,0 +1,11 @@ +#pragma once +#include "analysis/analysis.hpp" +#include "cont/base.hpp" + +namespace analysis { + struct PassContext { + cont::ImageMode image_mode; + std::optional image; + Function* function; + }; +} // namespace analysis \ No newline at end of file diff --git a/src/lib/analysis/lru_reg/lru_reg.hpp b/src/lib/analysis/lru_reg/lru_reg.hpp index 390f996..cb1a1d9 100644 --- a/src/lib/analysis/lru_reg/lru_reg.hpp +++ b/src/lib/analysis/lru_reg/lru_reg.hpp @@ -1,6 +1,6 @@ #pragma once +#include "cont/base.hpp" #include "easm/easm.hpp" -#include "pe/pe.hpp" #include #include @@ -146,23 +146,21 @@ namespace analysis { /// \brief An lru cache for all types of GP registers /// \tparam Img X64 or X86 image, depending on this image, the gp64 lru cache could be available - template class LRUReg { - constexpr static bool IsX64 = pe::is_x64_v; using RegTy = zasm::x86::Gp; public: DEFAULT_DTOR(LRUReg); DEFAULT_COPY(LRUReg); - LRUReg() { + explicit LRUReg(const cont::ImageMode image_mode): image_mode_(image_mode) { /// Push X86 registers for (const auto reg_id : detail::kRegistersX86) { push(reg_id); } /// Push x64 registers if needed - if constexpr (IsX64) { + if (image_mode_ == cont::ImageMode::X64) { for (const auto reg_id : detail::kRegistersX64) { push(reg_id); } @@ -227,7 +225,7 @@ namespace analysis { /// \param random should we choose a random register across least recently used registers? /// \return Register [[nodiscard]] RegTy get_gp64(const bool random = false) { - assert(IsX64); // no gp64 registers on x86 + assert(image_mode_ == cont::ImageMode::X64); // no gp64 registers on x86 return RegTy{to_gp64_if_needed(storage_.get(random))}; } @@ -260,14 +258,15 @@ namespace analysis { /// \brief Get machine mode based on the image tparam type /// \return machine_mode - [[nodiscard]] static zasm::MachineMode machine_mode() noexcept { - return IsX64 ? zasm::MachineMode::AMD64 : zasm::MachineMode::I386; + [[nodiscard]] zasm::MachineMode machine_mode() const noexcept { + /// \todo @es3n1n: get this from the image + return image_mode_ == cont::ImageMode::X64 ? zasm::MachineMode::AMD64 : zasm::MachineMode::I386; } /// \brief Converts any gp register to its base register, gp64 for x64 and gp32 for x86 /// \param reg_id register that it should convert /// \return Converted register id - [[nodiscard]] static RegID to_gp_ptr(const RegID reg_id) noexcept { + [[nodiscard]] RegID to_gp_ptr(const RegID reg_id) const noexcept { const auto reg = zasm::Reg{reg_id}; return reg.getRoot(machine_mode()).getId(); } @@ -275,13 +274,15 @@ namespace analysis { /// \brief Convert to GP64, if needed /// \param reg_id register that it should convert /// \return Converted GP64 reg id - [[nodiscard]] static RegID to_gp64_if_needed(const RegID reg_id) noexcept { + [[nodiscard]] RegID to_gp64_if_needed(const RegID reg_id) const noexcept { const auto gp_ptr = to_gp_ptr(reg_id); - return IsX64 ? gp_ptr : easm::reg_convert::gp32_to_gp64(gp_ptr); + return image_mode_ == cont::ImageMode::X64 ? gp_ptr : easm::reg_convert::gp32_to_gp64(gp_ptr); } private: /// \brief gp uptr lru container LRURegContainer storage_; + /// \brief image mode + cont::ImageMode image_mode_; }; } // namespace analysis \ No newline at end of file diff --git a/src/lib/analysis/passes/collect_img_references.hpp b/src/lib/analysis/passes/collect_img_references.hpp index 3ee8296..201dda6 100644 --- a/src/lib/analysis/passes/collect_img_references.hpp +++ b/src/lib/analysis/passes/collect_img_references.hpp @@ -3,12 +3,16 @@ #include "util/structs.hpp" namespace analysis::passes { - template struct collect_img_references_t { DEFAULT_CT_CTOR_DTOR(collect_img_references_t); NON_COPYABLE(collect_img_references_t); - static bool apply_insn(Function* function, insn_t& instruction, Img* image) { + static bool apply_insn(const PassContext& ctx, insn_t& instruction) { + // This is a weird pass, since it checks by image base we can't get it to work with nameless functions + if (!ctx.image.has_value()) { + return false; + } + // Looking for IMMs in the insn // const auto* imm = instruction.find_operand_if(); @@ -18,8 +22,9 @@ namespace analysis::passes { // Obtaining IMM value and image base // + auto* image = *ctx.image; const auto imm_value = imm->value(); - const auto base_address = image->raw_image->get_nt_headers()->optional_header.image_base; + const auto base_address = image->get_image_base(); // Skip instruction if imm isn't in the range of image // @@ -29,7 +34,7 @@ namespace analysis::passes { // Remembering reference // - function->image_references[imm_value - base_address].emplace_back(&instruction); + ctx.function->image_references[imm_value - base_address].emplace_back(&instruction); return true; } }; diff --git a/src/lib/analysis/passes/collect_lookup_table.hpp b/src/lib/analysis/passes/collect_lookup_table.hpp index 2d33865..1d27781 100644 --- a/src/lib/analysis/passes/collect_lookup_table.hpp +++ b/src/lib/analysis/passes/collect_lookup_table.hpp @@ -3,12 +3,11 @@ #include "util/structs.hpp" namespace analysis::passes { - template struct collect_lookup_table_t { DEFAULT_CT_CTOR_DTOR(collect_lookup_table_t); NON_COPYABLE(collect_lookup_table_t); - static bool apply_insn(Function* function, insn_t& instruction, Img* /*image*/) { + static bool apply_insn(const PassContext& ctx, insn_t& instruction) { if (!instruction.rva.has_value()) { /// \fixme @es3n1n: this could be pretty bad that we don't push newly /// created insns to the lookup table, although we should be just fine without them @@ -17,7 +16,7 @@ namespace analysis::passes { // Building rva<->insn lookup table // - function->instructions_lookup[*instruction.rva] = &instruction; + ctx.function->instructions_lookup[*instruction.rva] = &instruction; return true; } }; diff --git a/src/lib/analysis/passes/label_references.hpp b/src/lib/analysis/passes/label_references.hpp index 9a4d4d8..e3999fb 100644 --- a/src/lib/analysis/passes/label_references.hpp +++ b/src/lib/analysis/passes/label_references.hpp @@ -4,18 +4,17 @@ #include "util/structs.hpp" namespace analysis::passes { - template struct label_references_t { DEFAULT_CT_CTOR_DTOR(label_references_t); NON_COPYABLE(label_references_t); - static bool apply(Function* function, Img* /*image*/) { + static bool apply(const PassContext& ctx) { // Iterating over referenced RVAs within the function // - for (const auto& [referenced_insn_rva, insn_ptrs] : function->image_references) { + for (const auto& [referenced_insn_rva, insn_ptrs] : ctx.function->image_references) { // Skip if reference is not within the function // - if (!(referenced_insn_rva >= function->range.start && referenced_insn_rva <= function->range.end)) { + if (!(referenced_insn_rva >= ctx.function->range.start && referenced_insn_rva <= ctx.function->range.end)) { continue; } @@ -31,27 +30,27 @@ namespace analysis::passes { // Creating label for the referenced loc // - const auto referenced_loc_label = function->program.get()->createLabel(referenced_loc_name.c_str()); - const auto referenced_loc_label_node = function->program.get()->bindLabel(referenced_loc_label); + const auto referenced_loc_label = ctx.function->program->createLabel(referenced_loc_name.c_str()); + const auto referenced_loc_label_node = ctx.function->program->bindLabel(referenced_loc_label); if (!referenced_loc_label_node) [[unlikely]] { throw std::runtime_error("analysis: Unable to bind the label"); } // Moving label node to the referenced instruction // - const auto referenced_insn = function->instructions_lookup.find(referenced_insn_rva); - if (referenced_insn == function->instructions_lookup.end()) [[unlikely]] { + const auto referenced_insn = ctx.function->instructions_lookup.find(referenced_insn_rva); + if (referenced_insn == ctx.function->instructions_lookup.end()) [[unlikely]] { throw std::runtime_error("analysis: Unable to find referenced insn"); } - function->program.get()->moveBefore(referenced_insn->second->node_ref, *referenced_loc_label_node); - referenced_insn->second->bb_ref->push_label(*referenced_loc_label_node, function->bb_provider.get()); + ctx.function->program->moveBefore(referenced_insn->second->node_ref, *referenced_loc_label_node); + referenced_insn->second->bb_ref->push_label(*referenced_loc_label_node, ctx.function->bb_provider.get()); // Iterating over instructions that referenced this RVA // - for (auto* referrer_ptr : insn_ptrs) { + for (const auto* referrer_ptr : insn_ptrs) { // Looking for the imm in this instruction // - const auto imm_operand_index = referrer_ptr->template find_operand_index_if(); + const auto imm_operand_index = referrer_ptr->find_operand_index_if(); // No imm, huh? // diff --git a/src/lib/analysis/passes/lru_reg.hpp b/src/lib/analysis/passes/lru_reg.hpp index cf90eb1..b8d039a 100644 --- a/src/lib/analysis/passes/lru_reg.hpp +++ b/src/lib/analysis/passes/lru_reg.hpp @@ -3,16 +3,15 @@ #include "util/structs.hpp" namespace analysis::passes { - template struct lru_reg_t { DEFAULT_CT_CTOR_DTOR(lru_reg_t); NON_COPYABLE(lru_reg_t); - static bool apply_insn(Function* function, const insn_t& instruction, Img* /*image*/) { - /// Collect all the registers and push them to LRU + static bool apply_insn(const PassContext& ctx, const insn_t& instruction) { + /// Collect all the registers and push them to LRU. + /// \todo @es3n1n: We should track life time of registers for (auto& reg : easm::get_all_registers(*instruction.ref)) { - /// \note @es3n1n: We are pushing only known registers to avoid xmm/ymm/zmm stuff, we only need GP32/64 - function->lru_reg.push_known(reg.getId()); + ctx.function->lru_reg.push_known(reg.getId()); } return true; diff --git a/src/lib/analysis/passes/misc/bb_insn_passes.cpp b/src/lib/analysis/passes/misc/bb_insn_passes.cpp index defeb5f..34c140b 100644 --- a/src/lib/analysis/passes/misc/bb_insn_passes.cpp +++ b/src/lib/analysis/passes/misc/bb_insn_passes.cpp @@ -1,3 +1,5 @@ +#include + #include "analysis/passes/misc/bb_insn_passes.hpp" #include "analysis/common/common.hpp" @@ -8,36 +10,32 @@ namespace analysis::passes { namespace { - template - bool on_insn(Function* function, insn_t& instruction, Img* image) { + bool on_insn(PassContext& ctx, insn_t& instruction) { bool result = false; - result |= reloc_marker_t::apply_insn(function, instruction, image); - result |= collect_img_references_t::apply_insn(function, instruction, image); - result |= collect_lookup_table_t::apply_insn(function, instruction, image); - result |= lru_reg_t::apply_insn(function, instruction, image); + result |= reloc_marker_t::apply_insn(ctx, instruction); + result |= collect_img_references_t::apply_insn(ctx, instruction); + result |= collect_lookup_table_t::apply_insn(ctx, instruction); + result |= lru_reg_t::apply_insn(ctx, instruction); return result; } } // namespace - template - bool bb_insn_passes_t::apply(Function* function, Img* image) { + bool bb_insn_passes_t::apply(PassContext& pass_context) { bool result = false; // Iterating over BB and invoking callbacks // // - function->bb_storage->iter_bbs([&](bb_t& basic_block) -> void { + pass_context.function->bb_storage->iter_bbs([&](bb_t& basic_block) -> void { // Iterating over instructions and invoking callbacks // - std::for_each(basic_block.instructions.begin(), basic_block.instructions.end(), [&function, &image, &result](auto& instruction) -> void { // - result |= on_insn(function, *instruction, image); + std::ranges::for_each(basic_block.instructions, [&pass_context, &result](auto& instruction) -> void { // + result |= on_insn(pass_context, *instruction); }); }); return result; } - - PE_DECL_TEMPLATE_STRUCTS(bb_insn_passes_t); } // namespace analysis::passes diff --git a/src/lib/analysis/passes/misc/bb_insn_passes.hpp b/src/lib/analysis/passes/misc/bb_insn_passes.hpp index 7cbedbd..0b0faa6 100644 --- a/src/lib/analysis/passes/misc/bb_insn_passes.hpp +++ b/src/lib/analysis/passes/misc/bb_insn_passes.hpp @@ -1,5 +1,5 @@ #pragma once -#include "analysis/analysis.hpp" +#include "analysis/common/pass_context.hpp" #include "util/structs.hpp" // @@ -9,11 +9,10 @@ // namespace analysis::passes { - template struct bb_insn_passes_t { DEFAULT_CT_CTOR_DTOR(bb_insn_passes_t); NON_COPYABLE(bb_insn_passes_t); - static bool apply(Function* function, Img* image); + static bool apply(PassContext& pass_context); }; } // namespace analysis::passes diff --git a/src/lib/analysis/passes/reloc_marker.hpp b/src/lib/analysis/passes/reloc_marker.hpp index 3197a80..ccae7ea 100644 --- a/src/lib/analysis/passes/reloc_marker.hpp +++ b/src/lib/analysis/passes/reloc_marker.hpp @@ -1,14 +1,14 @@ #pragma once #include "analysis/analysis.hpp" +#include "analysis/common/pass_context.hpp" #include "util/structs.hpp" namespace analysis::passes { - template struct reloc_marker_t { DEFAULT_CT_CTOR_DTOR(reloc_marker_t); NON_COPYABLE(reloc_marker_t); - static bool apply_insn(Function* function [[maybe_unused]], insn_t& instruction, Img* image) { + static bool apply_insn(PassContext& ctx, insn_t& instruction) { // Would be set to true if instruction contains imm/ip operands // const zasm::Imm* imm = instruction.find_operand_if(); @@ -20,9 +20,13 @@ namespace analysis::passes { return false; } - /// Obtain needed stuff from pe - const auto image_base = image->raw_image->get_nt_headers()->optional_header.image_base; - const auto ptr_size = image->get_ptr_size(); + /// \fixme @es3n1n: we always assume base at 0x0 for nameless functions + std::uintptr_t image_base = 0; + if (ctx.image.has_value()) { + image_base = (*ctx.image)->get_image_base(); + } + /// \fixme @es3n1n: move to util + const auto ptr_size = ctx.image_mode == cont::ImageMode::X64 ? sizeof(std::uint64_t) : sizeof(std::uint32_t); // Force reloc mem if (mem != nullptr && mem->getBase().isIP()) { @@ -41,9 +45,15 @@ namespace analysis::passes { return true; } - // At this point, we are 100% sure that imm is set to something, so we can ignore the `imm != 0` check. + // For nameless functions there is no image, + // we should omit header checks as we don't have any header-relocations to proceed with. + if (!ctx.image.has_value()) { + return false; + } + auto* image = *ctx.image; + + // At this point, we are 100% sure that imm is set to something. // If there's an imm with the size of uintptr_t, we should check maybe it's present in the .reloc dir - // if (imm != nullptr && getBitSize(imm->getBitSize()) == (ptr_size * CHAR_BIT) && instruction.length >= ptr_size) { // Trying to find relocation from PE header within the instruction // \todo @es3n1n: check segments instead of just bruteforcing diff --git a/src/lib/analysis/var_alloc/var_alloc.hpp b/src/lib/analysis/var_alloc/var_alloc.hpp index cf3f314..a409d72 100644 --- a/src/lib/analysis/var_alloc/var_alloc.hpp +++ b/src/lib/analysis/var_alloc/var_alloc.hpp @@ -1,6 +1,5 @@ #pragma once #include "analysis/lru_reg/lru_reg.hpp" -#include "pe/pe.hpp" namespace analysis { struct SymVar { @@ -35,12 +34,11 @@ namespace analysis { } }; - template class VarAlloc { public: - DEFAULT_CT_CTOR_DTOR(VarAlloc); + DEFAULT_CTOR_DTOR(VarAlloc); DEFAULT_COPY(VarAlloc); - explicit VarAlloc(LRUReg* lru_reg): lru_reg_(lru_reg) { } + explicit VarAlloc(LRUReg* lru_reg): lru_reg_(lru_reg) { } /// \brief Get least recently used register as Gp8 /// \param random should we choose a random register across least recently used registers? @@ -87,9 +85,9 @@ namespace analysis { /// \brief Push flags to stack /// \param assembler zasm assembler ptr - static void push_flags(zasm::x86::Assembler* assembler) { + void push_flags(zasm::x86::Assembler* assembler) const { /// \fixme @es3n1n: we should track the instructions that affect CF instead - if constexpr (pe::is_x64_v) { + if (lru_reg_->machine_mode() == zasm::MachineMode::AMD64) { assembler->pushfq(); } else { assembler->pushfd(); @@ -106,8 +104,8 @@ namespace analysis { /// \brief Pop flags from stack /// \param assembler zasm assembler ptr - static void pop_flags(zasm::x86::Assembler* assembler) { - if constexpr (pe::is_x64_v) { + void pop_flags(zasm::x86::Assembler* assembler) const { + if (lru_reg_->machine_mode() == zasm::MachineMode::AMD64) { assembler->popfq(); } else { assembler->popfd(); @@ -134,6 +132,11 @@ namespace analysis { return stack_space_used_; } + /// \fixme @es3n1n: this is wrong + [[nodiscard]] zasm::MachineMode machine_mode() const noexcept { + return lru_reg_->machine_mode(); + } + private: /// \brief Filter out registers that are already in use /// \param callback callback that should return allocated reg @@ -168,6 +171,6 @@ namespace analysis { /// \brief How many bytes would we need for storing all allocated vars std::size_t stack_space_used_ = 0; /// \brief LRU registers storage - LRUReg* lru_reg_ = nullptr; + LRUReg* lru_reg_ = nullptr; }; } // namespace analysis \ No newline at end of file diff --git a/src/lib/cli/cli.hpp b/src/lib/cli/cli.hpp index 319110c..20246c7 100644 --- a/src/lib/cli/cli.hpp +++ b/src/lib/cli/cli.hpp @@ -11,21 +11,21 @@ namespace cli { {{"-pdb", "[path]", ""}, "Set custom .pdb file location"}, {{"-map", "[path]", ""}, "Set custom .map file location"}, {{"-f", "[name]", ""}, "Start new function configuration"}, + {{"-r", "[rva]", ""}, "Start new function configuration with RVA"}, {{"-t", "[name]", ""}, "Start new transform configuration"}, {{"-g", "[name]", ""}, "Start new transform global configuration"}, {{"-v", "[name]", "[value]"}, "Push value"}, }); - template - void dump_transforms(const std::string_view platform_name) { + inline void dump_transforms() { /// Header - logger::info("Available {} transforms:", platform_name); + logger::info("Available transforms:"); /// Iterate over the transforms - for (auto& scheduler = obfuscator::TransformScheduler::get().for_arch(); // - auto& [tag, transform] : scheduler.transforms) { + for (auto& scheduler = obfuscator::TransformScheduler::get(); // + auto& [tag, transform] : scheduler.container.transforms) { /// Get the shared cfg - auto& shared_cfg = obfuscator::TransformSharedConfigStorage::get().get_for(tag); + const auto& shared_cfg = obfuscator::TransformSharedConfigStorage::get().get_for(tag); /// Dump transform name logger::info<1>("{}", shared_cfg.name); @@ -51,7 +51,7 @@ namespace cli { logger::info("Shared transform variables (e.g could be set for every transform):"); /// Get some scheduler and any transform + shared config for it - const auto& scheduler = obfuscator::TransformScheduler::get().for_arch(); + const auto& scheduler = obfuscator::TransformScheduler::get().container; const auto& shared_cfg = obfuscator::TransformSharedConfigStorage::get().get_for(scheduler.transforms.begin()->first); /// Dump all vars + their defaults @@ -86,10 +86,7 @@ namespace cli { "-v SomeGlobalName 1337"); pad(); - detail::dump_transforms("x64"); - pad(); - - detail::dump_transforms("x86"); + detail::dump_transforms(); pad(); detail::dump_shared_vars(); diff --git a/src/lib/config_parser/config_parser.cpp b/src/lib/config_parser/config_parser.cpp index 6fc75eb..9721cb1 100644 --- a/src/lib/config_parser/config_parser.cpp +++ b/src/lib/config_parser/config_parser.cpp @@ -92,6 +92,14 @@ namespace config_parser { continue; } + /// Function RVA + if (arg_ == "-r" && next_arg_.has_value()) { + state.current_function = &result.create_function_config(); + state.current_function->rva = string_parser::parse_int64(next_arg_.value()); + skip(1); + continue; + } + /// Transform configuration start if (arg_ == "-t" && next_arg_.has_value() && state.busy()) { state.current_transform = &state.current_function->transform_configurations.emplace_back(); @@ -115,21 +123,12 @@ namespace config_parser { } if (arg_ == "-seed" && next_arg_.has_value()) { - /// \todo @sovissa: - /// [1] add correct processing of literals to string_parser - /// [2] add support for negative values to arg parser - std::size_t seed_value_base = 10; - std::string_view next_arg_view{next_arg_.value()}; if (next_arg_view.starts_with('-')) { next_arg_view = next_arg_view.substr(1); } - if (next_arg_view.length() >= 2 && // - (next_arg_view.starts_with("0x") || next_arg_view.starts_with("0X"))) { - seed_value_base = 16; - } - seed = string_parser::parse_uint64(next_arg_.value(), seed_value_base); + seed = string_parser::parse_uint64(next_arg_view); skip(1); continue; } diff --git a/src/lib/config_parser/config_parser.hpp b/src/lib/config_parser/config_parser.hpp index b0df5dd..9ceed12 100644 --- a/src/lib/config_parser/config_parser.hpp +++ b/src/lib/config_parser/config_parser.hpp @@ -30,6 +30,10 @@ namespace config_parser { return global_transform_configurations_; } + [[nodiscard]] std::vector function_configurations() { + return function_configurations_; + } + [[nodiscard]] auto begin() { return function_configurations_.begin(); } diff --git a/src/lib/config_parser/structs.hpp b/src/lib/config_parser/structs.hpp index 09bf1c9..61cc472 100644 --- a/src/lib/config_parser/structs.hpp +++ b/src/lib/config_parser/structs.hpp @@ -1,5 +1,6 @@ #pragma once #include +#include #include #include #include @@ -10,9 +11,27 @@ namespace config_parser { std::unordered_map values; }; + using transform_configurations_t = std::vector; + + struct nameless_function_configuration_t { + transform_configurations_t transform_configurations; + }; + struct function_configuration_t { - std::string function_name; - std::vector transform_configurations; + std::optional function_name; + std::optional rva; + transform_configurations_t transform_configurations; + + [[nodiscard]] std::string name() const { + if (function_name.has_value()) { + return *function_name; + } + if (rva.has_value()) { + return std::format("sub_{:x}", *rva); + } + /// \todo @es3n1n: Swap nameless funcs to use this struct too + return "nameless"; + } }; struct obfuscator_config_t { diff --git a/src/lib/cont/base.hpp b/src/lib/cont/base.hpp new file mode 100644 index 0000000..0b54d74 --- /dev/null +++ b/src/lib/cont/base.hpp @@ -0,0 +1,488 @@ +#pragma once +#include +#include +#include +#include + +#include "es3n1n/common/memory/address.hpp" +#include "util/structs.hpp" + +#include +#include +#include +#include + +namespace cont { + enum struct ContImageType : std::uint8_t { + PE = 0, + ELF = 1, + }; + + enum struct ImageMode : std::uint8_t { + X64 = 0, + X86 = 1, + }; + + enum struct RelocationType : std::uint8_t { + Absolute = 0, + High = 1, + Low = 2, + HighLow = 3, + HighAdj = 4, + Ia64Imm64 = 5, + Dir64 = 6, + GlobDat64 = 7, + GlobDat32 = 8, + JumpSlot = 9, + Copy = 10, + }; + + enum struct RelocationSource : std::uint8_t { + REL = 0, + RELA = 1, + }; + + struct Relocation { + memory::address rva; + std::uint8_t size = 0; // in bytes + RelocationType type = RelocationType::Absolute; + std::optional sym_index = std::nullopt; + std::optional addend = std::nullopt; + std::optional info_raw = std::nullopt; + std::optional source = std::nullopt; + }; + + enum struct DirectoryType : std::uint8_t { + Export = 0, + Import = 1, + Resource = 2, + Exception = 3, + Security = 4, + Reloc = 5, + Debug = 6, + Copyright = 7, + Architecture = 8, + GlobalPtr = 9, + Tls = 10, + LoadConfig = 11, + BoundImport = 12, + Iat = 13, + DelayImport = 14, + ComDescriptor = 15, + MAX_LENGTH = 16, + }; + + struct DirectoryProperties { + std::size_t offset; // from the section start + std::size_t size; // in bytes + }; + + struct Section { + std::optional name = std::nullopt; + std::size_t virtual_size = 0U; + std::size_t virtual_address = 0U; + std::size_t size_raw_data = 0U; + std::size_t ptr_raw_data = 0U; + + std::vector raw_data; + + bool symbolic = false; // ignore while linking + std::optional elf_alignment = std::nullopt; + std::optional elf_type = std::nullopt; + + std::array, std::to_underlying(DirectoryType::MAX_LENGTH)> contains_directories = {}; + union Characteristics { + uint32_t flags; + struct { + uint32_t cnt_code:1; // Section contains code. + uint32_t cnt_init_data:1; // Section contains initialized data. + uint32_t cnt_uninit_data:1; // Section contains uninitialized data. + uint32_t lnk_info:1; // Section contains comments or some other type of information. + uint32_t lnk_remove:1; // Section contents will not become part of image. + uint32_t lnk_comdat:1; // Section contents comdat. + uint32_t no_defer_spec_exc:1; // Reset speculative exceptions handling bits in the TLB entries for this section. + uint32_t mem_far:1; + uint32_t mem_purgeable:1; + uint32_t mem_locked:1; + uint32_t mem_preload:1; + uint32_t alignment:4; // Alignment calculated as: n ? 1 << ( n - 1 ) : 16 + uint32_t lnk_nreloc_ovfl:1; // Section contains extended relocations. + uint32_t mem_discardable:1; // Section can be discarded. + uint32_t mem_not_cached:1; // Section is not cachable. + uint32_t mem_not_paged:1; // Section is not pageable. + uint32_t mem_shared:1; // Section is shareable. + uint32_t mem_execute:1; // Section is executable. + uint32_t mem_read:1; // Section is readable. + uint32_t mem_write:1; // Section is writeable. + }; + } characteristics = {}; + + [[nodiscard]] std::optional directory_info(const DirectoryType dir_type) const { + return contains_directories[std::to_underlying(dir_type)]; + } + + [[nodiscard]] bool has_directory(const DirectoryType type) const { + return contains_directories[std::to_underlying(type)].has_value(); + } + + void set_contained_dir(const DirectoryType type, const std::size_t offset, const std::size_t size) { + contains_directories[std::to_underlying(type)] = DirectoryProperties{.offset = offset, .size = size}; + } + }; + + class ImageBase { + public: + explicit ImageBase(const ContImageType image_type, const memory::address raw_image): image_type_(image_type), raw_image_(raw_image) { } + virtual ~ImageBase() = default; + DEFAULT_COPY(ImageBase); + + [[nodiscard]] virtual ImageMode mode() const = 0; + [[nodiscard]] virtual bool verify_integrity() const = 0; + + [[nodiscard]] virtual std::size_t get_image_base() const = 0; + [[nodiscard]] virtual std::size_t get_section_alignment() const = 0; + [[nodiscard]] virtual std::size_t get_file_alignment() const = 0; + + [[nodiscard]] virtual std::vector rebuild_image() = 0; + + void realign_sections() { + /// Nothing to realign + if (sections.size() <= 1) { + return; + } + + /// Making sure that all section virtual sizes are aligned + for (std::size_t i = 0; i < sections.size() - 1; ++i) { + auto& sec = sections.at(i); + const auto& next_sec = sections.at(i + 1); + sec.virtual_size = next_sec.virtual_address - sec.virtual_address; + } + } + + void erase_section_if(const std::function& pred) { + if (const auto iter = std::ranges::find_if(sections, pred); iter != sections.end()) { + sections.erase(iter); + } else { + throw std::runtime_error("cont: Unable to erase section by predicate"); + } + } + + [[nodiscard]] Section& new_section(Section& object) { + const auto section_alignment = get_section_alignment(); + const auto file_alignment = get_file_alignment(); + + auto& new_sec = sections.emplace_back(object); + assert(new_sec.size_raw_data > 0); + + const auto last_va_section = std::ranges::max_element(sections, [](const Section& sec, const Section& sec2) -> auto { //) + return (sec.virtual_address + sec.virtual_size) < (sec2.virtual_address + sec2.virtual_size); + }); + const auto last_raw_section = std::ranges::max_element(sections, [](const Section& sec, const Section& sec2) -> auto { // + return (sec.ptr_raw_data + sec.size_raw_data) < (sec2.ptr_raw_data + sec2.size_raw_data); + }); + + assert(last_va_section != std::end(sections) && last_raw_section != std::end(sections)); + + new_sec.virtual_size = new_sec.size_raw_data = memory::address{new_sec.size_raw_data} // + .align_up(section_alignment) + .as(); + + new_sec.virtual_address = memory::address{last_va_section->virtual_address + last_va_section->virtual_size} // + .align_up(section_alignment) + .as(); + + new_sec.ptr_raw_data = memory::address{last_raw_section->ptr_raw_data + last_raw_section->size_raw_data} // + .align_up(file_alignment) + .as(); + + new_sec.raw_data.resize(new_sec.size_raw_data); + + return new_sec; + } + + [[nodiscard]] Section& find_section_if(const std::function& pred) { + const auto iter = std::ranges::find_if(sections, pred); + if (iter == sections.end()) { + throw std::runtime_error("cont: Unable to find section by predicate"); + } + return *iter; + } + + [[nodiscard]] Section* find_section_if_ptr(const std::function& pred) { + const auto iter = std::ranges::find_if(sections, pred); + if (iter == sections.end()) { + return nullptr; + } + return &*iter; + } + + [[nodiscard]] std::vector
find_sections_if(const std::function& pred) const { + std::vector
result; + for (const auto& section : sections) { + if (pred(section)) { + result.push_back(section); + } + } + return result; + } + + [[nodiscard]] Section& find_last_section() { + const auto result = std::ranges::max_element(sections, [](const Section& lhs, const Section& rhs) -> bool { // + return lhs.virtual_address < rhs.virtual_address; + }); + + if (result == std::end(sections)) { + throw std::runtime_error("cont: Unable to find last section"); + } + + return *result; + } + + [[nodiscard]] virtual Section* rva_to_section(std::uint32_t rva) { + const auto iter = std::ranges::find_if(sections, [rva](const Section& sec) -> bool { // + return rva >= sec.virtual_address && rva <= (sec.virtual_address + sec.virtual_size); + }); + + if (iter == sections.end()) { + return nullptr; + } + + return &*iter; + } + + template + [[nodiscard]] Ty* rva_to_ptr(const memory::address rva) { + const auto* section = rva_to_section(rva.as()); + if (section == nullptr) { + return nullptr; + } + + const auto offset = rva.inner() - section->virtual_address; + return memory::address{section->raw_data.data()}.offset(static_cast(offset)).as>(); + } + + [[nodiscard]] zasm::MachineMode machine_mode() const { + switch (mode()) { + case ImageMode::X64: + return zasm::MachineMode::AMD64; + case ImageMode::X86: + return zasm::MachineMode::I386; + default: + throw std::out_of_range("cont::ImageBase::machine_mode: Unsupported image mode"); + } + } + + [[nodiscard]] std::size_t ptr_size() const { + switch (mode()) { + case ImageMode::X64: + return sizeof(std::uint64_t); + case ImageMode::X86: + return sizeof(std::uint32_t); + default: + throw std::out_of_range("cont::ImageBase::ptr_size: Unsupported image mode"); + } + } + + [[nodiscard]] memory::address raw_image() const noexcept { + return raw_image_; + } + + [[nodiscard]] ContImageType image_type() const noexcept { + return image_type_; + } + + protected: + virtual void update_sections() = 0; + virtual void update_relocations() = 0; + + void initialize() { + update_sections(); + update_relocations(); + } + + memory::address raw_image_; + ContImageType image_type_; + + public: + /// An unordered map that consists of {rva: reloc_info} + std::unordered_map relocations; + std::vector
sections; + }; +} // namespace cont + +constexpr cont::RelocationType to_cont(const win::reloc_type_id type) { + switch (type) { + case win::reloc_type_id::rel_based_absolute: + return cont::RelocationType::Absolute; + case win::reloc_type_id::rel_based_high: + return cont::RelocationType::High; + case win::reloc_type_id::rel_based_low: + return cont::RelocationType::Low; + case win::reloc_type_id::rel_based_high_low: + return cont::RelocationType::HighLow; + case win::reloc_type_id::rel_based_high_adj: + return cont::RelocationType::HighAdj; + case win::reloc_type_id::rel_based_ia64_imm64: + return cont::RelocationType::Ia64Imm64; + case win::reloc_type_id::rel_based_dir64: + return cont::RelocationType::Dir64; + } + throw std::runtime_error("pe: unsupported relocation type"); +} + +constexpr cont::DirectoryType to_cont(const win::directory_id type) { + switch (type) { + case win::directory_id::directory_entry_export: + return cont::DirectoryType::Export; + case win::directory_id::directory_entry_import: + return cont::DirectoryType::Import; + case win::directory_id::directory_entry_resource: + return cont::DirectoryType::Resource; + case win::directory_id::directory_entry_exception: + return cont::DirectoryType::Exception; + case win::directory_id::directory_entry_security: + return cont::DirectoryType::Security; + case win::directory_id::directory_entry_basereloc: + return cont::DirectoryType::Reloc; + case win::directory_id::directory_entry_debug: + return cont::DirectoryType::Debug; + // no copyright + case win::directory_id::directory_entry_architecture: + return cont::DirectoryType::Architecture; + case win::directory_id::directory_entry_globalptr: + return cont::DirectoryType::GlobalPtr; + case win::directory_id::directory_entry_tls: + return cont::DirectoryType::Tls; + case win::directory_id::directory_entry_load_config: + return cont::DirectoryType::LoadConfig; + case win::directory_id::directory_entry_bound_import: + return cont::DirectoryType::BoundImport; + case win::directory_id::directory_entry_iat: + return cont::DirectoryType::Iat; + case win::directory_id::directory_entry_delay_import: + return cont::DirectoryType::DelayImport; + case win::directory_id::directory_entry_com_descriptor: + return cont::DirectoryType::ComDescriptor; + default: + throw std::runtime_error("pe: unsupported directory type"); + } +} + +constexpr cont::Section to_cont(const win::section_header_t section) { + cont::Section result = {}; + result.name = section.name.to_string(); + result.virtual_size = static_cast(section.virtual_size); + result.virtual_address = static_cast(section.virtual_address); + result.size_raw_data = static_cast(section.size_raw_data); + result.ptr_raw_data = static_cast(section.ptr_raw_data); + result.characteristics.cnt_code = section.characteristics.cnt_code; + result.characteristics.cnt_init_data = section.characteristics.cnt_init_data; + result.characteristics.cnt_uninit_data = section.characteristics.cnt_uninit_data; + result.characteristics.lnk_info = section.characteristics.lnk_info; + result.characteristics.lnk_remove = section.characteristics.lnk_remove; + result.characteristics.lnk_comdat = section.characteristics.lnk_comdat; + result.characteristics.no_defer_spec_exc = section.characteristics.no_defer_spec_exc; + result.characteristics.mem_far = section.characteristics.mem_far; + result.characteristics.mem_purgeable = section.characteristics.mem_purgeable; + result.characteristics.mem_locked = section.characteristics.mem_locked; + result.characteristics.mem_preload = section.characteristics.mem_preload; + result.characteristics.alignment = section.characteristics.alignment; + result.characteristics.lnk_nreloc_ovfl = section.characteristics.lnk_nreloc_ovfl; + result.characteristics.mem_discardable = section.characteristics.mem_discardable; + result.characteristics.mem_not_cached = section.characteristics.mem_not_cached; + result.characteristics.mem_not_paged = section.characteristics.mem_not_paged; + result.characteristics.mem_shared = section.characteristics.mem_shared; + result.characteristics.mem_execute = section.characteristics.mem_execute; + result.characteristics.mem_read = section.characteristics.mem_read; + result.characteristics.mem_write = section.characteristics.mem_write; + return result; +} + +template + requires(traits::is_any_of_v) +[[nodiscard]] cont::Section to_cont(const Phdr& phdr) { + cont::Section s{}; + s.name = std::nullopt; + s.virtual_size = static_cast(phdr.p_memsz); + s.virtual_address = static_cast(phdr.p_vaddr); + s.size_raw_data = static_cast(phdr.p_filesz); + s.ptr_raw_data = static_cast(phdr.p_offset); + s.characteristics.mem_execute = !!(phdr.p_flags & PF_X); + s.characteristics.mem_write = !!(phdr.p_flags & PF_W); + s.characteristics.mem_read = !!(phdr.p_flags & PF_R); + s.elf_type = phdr.p_type; + s.elf_alignment = phdr.p_align; + return s; +} + +template + requires(traits::is_any_of_v) +[[nodiscard]] cont::Section to_cont(const Shdr& shdr, const std::string_view name) { + cont::Section s{}; + s.name = std::string{name}; + s.virtual_size = static_cast(shdr.sh_size); + s.virtual_address = static_cast(shdr.sh_addr); + s.size_raw_data = shdr.sh_type == SHT_NOBITS ? 0U : static_cast(shdr.sh_size); + s.ptr_raw_data = static_cast(shdr.sh_offset); + s.characteristics.mem_execute = !!(shdr.sh_flags & SHF_EXECINSTR); + s.characteristics.mem_write = !!(shdr.sh_flags & SHF_WRITE); + s.characteristics.mem_read = !!(shdr.sh_flags & SHF_ALLOC); + return s; +} + +[[nodiscard]] constexpr cont::RelocationType to_cont(const std::uint64_t elf_type, const cont::ImageMode mode) { + switch (mode) { + case cont::ImageMode::X64: { + switch (ELF64_R_TYPE(elf_type)) { + case R_X86_64_64: + case R_X86_64_RELATIVE: + return cont::RelocationType::Dir64; + case R_X86_64_32: + case R_X86_64_32S: + return cont::RelocationType::HighLow; + case R_X86_64_GLOB_DAT: + return cont::RelocationType::GlobDat64; + case R_X86_64_JUMP_SLOT: + return cont::RelocationType::JumpSlot; + case R_X86_64_COPY: + return cont::RelocationType::Copy; + default: + throw std::out_of_range(std::format("cont::to_cont: Unsupported relocation type for x64 mode {}", ELF64_R_TYPE(elf_type))); + } + } + case cont::ImageMode::X86: { + switch (ELF32_R_TYPE(elf_type)) { + case R_386_32: + case R_386_RELATIVE: + return cont::RelocationType::HighLow; + case R_386_GLOB_DAT: + return cont::RelocationType::GlobDat32; + default: + throw std::out_of_range("cont::to_cont: Unsupported relocation type for x86 mode"); + } + } + default: + throw std::out_of_range("cont::to_cont: Unsupported image mode"); + } +} + +constexpr win::reloc_type_id to_win(const cont::RelocationType type) { + switch (type) { + case cont::RelocationType::Absolute: + return win::reloc_type_id::rel_based_absolute; + case cont::RelocationType::High: + return win::reloc_type_id::rel_based_high; + case cont::RelocationType::Low: + return win::reloc_type_id::rel_based_low; + case cont::RelocationType::HighLow: + return win::reloc_type_id::rel_based_high_low; + case cont::RelocationType::HighAdj: + return win::reloc_type_id::rel_based_high_adj; + case cont::RelocationType::Ia64Imm64: + return win::reloc_type_id::rel_based_ia64_imm64; + case cont::RelocationType::Dir64: + return win::reloc_type_id::rel_based_dir64; + default: + throw std::runtime_error("pe: unsupported relocation type"); + } +} diff --git a/src/lib/cont/cont.hpp b/src/lib/cont/cont.hpp new file mode 100644 index 0000000..2a60f8a --- /dev/null +++ b/src/lib/cont/cont.hpp @@ -0,0 +1,16 @@ +#pragma once +#include "elf/image.hpp" +#include "pe/image.hpp" + +namespace cont { + [[nodiscard]] inline ContImageType get_image_type(std::span image_data) { + if (image_data.size() >= sizeof(win::DOS_HDR_MAGIC) && + memory::address(image_data.data()).read>() == win::DOS_HDR_MAGIC) { + return ContImageType::PE; + } + if (image_data.size() >= SELFMAG && std::memcmp(image_data.data(), ELFMAG, SELFMAG) == 0) { + return ContImageType::ELF; + } + throw std::runtime_error("cont: get_image_type: Unsupported image type"); + } +} // namespace cont diff --git a/src/lib/cont/elf/elf_image.cpp b/src/lib/cont/elf/elf_image.cpp new file mode 100644 index 0000000..2a7723c --- /dev/null +++ b/src/lib/cont/elf/elf_image.cpp @@ -0,0 +1,199 @@ +#include "cont/elf/image.hpp" +#include "cont/elf/rebuilder/rebuilder.hpp" + +#include +#include + +namespace cont::elf { + [[nodiscard]] ImageMode Image::mode() const { + switch (x64()->e_ident[EI_CLASS]) { + case ELFCLASS64: + return ImageMode::X64; + case ELFCLASS32: + return ImageMode::X86; + default: + throw std::runtime_error("cont::elf: Unsupported ELF class"); + } + } + + [[nodiscard]] bool Image::verify_integrity() const { + return std::memcmp(static_cast(x64()->e_ident), ELFMAG, SELFMAG) == 0; + } + + [[nodiscard]] std::size_t Image::get_image_base() const { + static auto result = [this]() -> std::size_t { + switch (mode()) { + case ImageMode::X64: { + auto phdrs = phdr_span(x64()); + const auto it = std::ranges::find_if(phdrs, [](const Elf64_Phdr& ph) { return ph.p_type == PT_LOAD; }); + if (it == phdrs.end()) { + throw std::runtime_error("cont::elf: No PT_LOAD segment found"); + } + + return it->p_vaddr; + } + + case ImageMode::X86: { + auto phdrs = phdr_span(x86()); + const auto it = std::ranges::find_if(phdrs, [](const Elf32_Phdr& ph) { return ph.p_type == PT_LOAD; }); + if (it == phdrs.end()) { + throw std::runtime_error("cont::elf: No PT_LOAD segment found"); + } + + return it->p_vaddr; + } + default: + throw std::runtime_error("cont::elf: Unsupported ELF mode"); + } + }(); + return result; + } + + [[nodiscard]] std::size_t Image::get_section_alignment() const { + static auto result = [this]() -> std::size_t { + switch (mode()) { + case ImageMode::X64: { + auto phdrs = phdr_span(x64()); + const auto it = std::ranges::find_if(phdrs, [](const Elf64_Phdr& ph) { return ph.p_type == PT_LOAD; }); + if (it == phdrs.end()) { + throw std::runtime_error("cont::elf: No PT_LOAD segment found"); + } + + return it->p_align; + } + + case ImageMode::X86: { + auto phdrs = phdr_span(x86()); + const auto it = std::ranges::find_if(phdrs, [](const Elf32_Phdr& ph) { return ph.p_type == PT_LOAD; }); + if (it == phdrs.end()) { + throw std::runtime_error("cont::elf: No PT_LOAD segment found"); + } + + return it->p_align; + } + default: + throw std::runtime_error("cont::elf: Unsupported ELF mode"); + } + }(); + return result; + } + + [[nodiscard]] std::size_t Image::get_file_alignment() const { + return get_section_alignment(); + } + + [[nodiscard]] std::vector Image::rebuild_image() { + const auto ctx = RebuilderContext{.image = this}; + return rebuild_elf(ctx); + } + + void Image::update_sections() { + auto build_from_segments = [&](const ElfEhdr* ehdr) -> void { + auto phdrs = phdr_span(ehdr); + + sections.clear(); + for (const ElfPhdr& phdr : phdrs) { + auto& sec = sections.emplace_back(to_cont(phdr)); + + sec.raw_data.resize(sec.size_raw_data, 0); + if (sec.size_raw_data != 0) { + const auto* src = raw_image_.offset(sec.ptr_raw_data).template as(); + std::memcpy(sec.raw_data.data(), src, sec.size_raw_data); + } + + if (sec.elf_type == PT_DYNAMIC) { + dynamics.clear(); + for (auto* iter = memory::address(sec.raw_data.data()).ptr(); iter != nullptr && iter->d_tag != DT_NULL; ++iter) { + auto& [tag, value] = dynamics.emplace_back(); + tag = iter->d_tag; + value = iter->d_un.d_val; + } + } + } + + if (const auto* const plt_got = find_dynamic(DT_PLTGOT); plt_got != nullptr) { + sections.emplace_back(Section{ + .name = ".plt", + .virtual_size = 0x1000, + .virtual_address = plt_got->value, + .size_raw_data = 0, + .ptr_raw_data = 0, + .symbolic = true, + .elf_type = PT_DYNAMIC, + }); + } + + std::ranges::sort(sections, [](const Section& lhs, const Section& rhs) -> bool { + /// PHDR should always be first + if (lhs.elf_type.value() == PT_PHDR) { + return true; + } + if (rhs.elf_type.value() == PT_PHDR) { + return false; + } + return lhs.virtual_address < rhs.virtual_address; + }); + }; + + switch (mode()) { + case ImageMode::X64: { + const auto* const ehdr = x64(); + build_from_segments.operator()(ehdr); + break; + } + case ImageMode::X86: { + const auto* const ehdr = x86(); + build_from_segments.operator()(ehdr); + break; + } + default: + throw std::out_of_range("cont::elf::update_sections: unsupported image mode"); + } + + logger::debug("elf: parsed {} sections", sections.size()); + } + + void Image::update_relocations() { + const auto img_mode = mode(); + + const auto proceed = [this, img_mode]() -> void { + if (const auto* const rela = find_dynamic(DT_RELA); rela != nullptr) { + const auto* const rela_sz = find_dynamic(DT_RELASZ); + if (rela_sz == nullptr || rela->value == 0) { + throw std::runtime_error("cont::elf: DT_RELA or DT_RELASZ not found or invalid"); + } + + const auto* rela_ptr = rva_to_ptr(rela->value); + for (const auto* const rela_end = rela_ptr + (rela_sz->value / sizeof(Rela)); rela_ptr < rela_end; ++rela_ptr) { + const auto converted_type = to_cont(rela_ptr->r_info, img_mode); + relocations[rela_ptr->r_offset] = + Relocation{.rva = rela_ptr->r_offset, .type = converted_type, .addend = rela_ptr->r_addend, .info_raw = rela_ptr->r_info}; + } + } + + if (const auto* const jmprel = find_dynamic(DT_JMPREL); jmprel != nullptr) { + const auto* rela_ptr = raw_image_.offset(static_cast(jmprel->value)).as(); + for (; rela_ptr->r_info != 0; ++rela_ptr) { + const auto converted_type = to_cont(rela_ptr->r_info, img_mode); + jmprel_relocations[rela_ptr->r_offset] = + Relocation{.rva = rela_ptr->r_offset, .type = converted_type, .addend = rela_ptr->r_addend, .info_raw = rela_ptr->r_info}; + } + } + }; + + switch (img_mode) { + case ImageMode::X64: { + proceed.operator()(); + break; + } + case ImageMode::X86: { + proceed.operator()(); + break; + } + default: + throw std::out_of_range("cont::elf::update_relocations: unsupported image mode"); + } + + logger::debug("elf: parsed {}(+{}) relocations", relocations.size(), jmprel_relocations.size()); + } +} // namespace cont::elf diff --git a/src/lib/cont/elf/image.hpp b/src/lib/cont/elf/image.hpp new file mode 100644 index 0000000..b2e4019 --- /dev/null +++ b/src/lib/cont/elf/image.hpp @@ -0,0 +1,98 @@ +#pragma once +#include +#include +#include + +namespace cont::elf { + struct Dynamic { + std::uint64_t tag = 0; + std::uint64_t value = 0; + }; + + class Image final : public ImageBase { + public: + explicit Image(const memory::address raw_image): ImageBase(ContImageType::ELF, raw_image) { + initialize(); + } + [[nodiscard]] ImageMode mode() const override; + [[nodiscard]] bool verify_integrity() const override; + + [[nodiscard]] std::size_t get_image_base() const override; + [[nodiscard]] std::size_t get_section_alignment() const override; + [[nodiscard]] std::size_t get_file_alignment() const override; + + [[nodiscard]] std::vector rebuild_image() override; + + void update_sections() override; + void update_relocations() override; + + [[nodiscard]] std::size_t non_symbolic_segments_count() const { + return std::accumulate(sections.begin(), sections.end(), static_cast(0), + [](const auto acc, const Section& sec) -> std::size_t { return acc + (sec.symbolic ? 0 : 1); }); + } + + [[nodiscard]] Section* find_segment_with_type(const std::size_t type) { + const auto iter = std::ranges::find_if(sections, [type](const Section& sec) -> bool { // + return sec.elf_type.value_or(PT_LOAD) == type; + }); + if (iter == sections.end()) { + return nullptr; + } + + return &*iter; + } + + [[nodiscard]] Section* rva_to_section(std::uint32_t rva) override { + const auto iter = std::ranges::find_if(sections, [rva](const Section& sec) -> bool { // + return sec.elf_type.value_or(PT_LOAD) == PT_LOAD && // + rva >= sec.virtual_address && rva <= (sec.virtual_address + sec.virtual_size); + }); + + if (iter == sections.end()) { + return nullptr; + } + + return &*iter; + } + + [[nodiscard]] Dynamic* find_dynamic(const std::size_t tag) { + const auto it = std::ranges::find_if(dynamics, [tag](const Dynamic& dyn) { return dyn.tag == tag; }); + if (it == dynamics.end()) { + return nullptr; + } + + return &*it; + } + + void delete_dynamic(const std::size_t tag) { + std::erase_if(dynamics, [tag](const Dynamic& dyn) { return dyn.tag == tag; }); + } + + std::vector dynamics; + + private: + std::unordered_map jmprel_relocations; + + [[nodiscard]] Elf64_Ehdr* x64() const { + return raw_image_.as(); + } + + [[nodiscard]] Elf32_Ehdr* x86() const { + return raw_image_.as(); + } + + template + [[nodiscard]] static std::span phdr_span(const ElfEhdr* ehdr) { + const auto* first = reinterpret_cast(reinterpret_cast(ehdr) + ehdr->e_phoff); + return {first, static_cast(ehdr->e_phnum)}; + } + + template + [[nodiscard]] static std::span shdr_span(const ElfEhdr* ehdr) { + const auto* first = reinterpret_cast(reinterpret_cast(ehdr) + ehdr->e_shoff); + return {first, static_cast(ehdr->e_shnum)}; + } + }; + + template concept AnyRawImage = traits::is_any_of_v; +} // namespace cont::elf diff --git a/src/lib/cont/elf/rebuilder/detail/elf_copy_sections.cpp b/src/lib/cont/elf/rebuilder/detail/elf_copy_sections.cpp new file mode 100644 index 0000000..4e8b5aa --- /dev/null +++ b/src/lib/cont/elf/rebuilder/detail/elf_copy_sections.cpp @@ -0,0 +1,17 @@ +#include "cont/elf/rebuilder/rebuilder.hpp" + +namespace cont::elf::detail { + void copy_sections(Image* image, std::vector& result) { + for (auto& section : image->sections) { + if (section.size_raw_data == 0 || section.symbolic) { + continue; + } + + /// Copy the section raw data to the result + const auto src = memory::address(section.raw_data.data()); + const auto dst = memory::address(result.data()).offset(static_cast(section.ptr_raw_data)); + + std::memcpy(dst.ptr(), src.ptr(), section.size_raw_data); + } + } +} // namespace cont::elf::detail diff --git a/src/lib/cont/elf/rebuilder/detail/elf_init_header.cpp b/src/lib/cont/elf/rebuilder/detail/elf_init_header.cpp new file mode 100644 index 0000000..b560504 --- /dev/null +++ b/src/lib/cont/elf/rebuilder/detail/elf_init_header.cpp @@ -0,0 +1,107 @@ +#include "cont/elf/rebuilder/rebuilder.hpp" + +namespace cont::elf::detail { + void init_header(Image* image, std::vector& result) { + const auto img_mode = image->mode(); + + /// We will be rewriting PHT entries, let's alloc a segment for that + auto pht_sec_hdr = Section{}; + pht_sec_hdr.name = "pht"; + pht_sec_hdr.size_raw_data = (img_mode == ImageMode::X64 ? sizeof(Elf64_Phdr) : sizeof(Elf32_Phdr)) * image->non_symbolic_segments_count(); + pht_sec_hdr.characteristics.mem_read = true; + pht_sec_hdr.elf_type = PT_LOAD; + auto& pht_seg = image->new_section(pht_sec_hdr); + + /// Update phdr segment + auto* old_pht_seg = + image->find_section_if_ptr([](const Section& section) -> bool { return section.elf_type.has_value() && *section.elf_type == PT_PHDR; }); + if (old_pht_seg != nullptr) { + old_pht_seg->ptr_raw_data = pht_seg.ptr_raw_data; + old_pht_seg->size_raw_data = pht_seg.size_raw_data; + old_pht_seg->virtual_address = pht_seg.virtual_address; + old_pht_seg->virtual_size = pht_seg.virtual_size; + } + + /// Serializing PHT entries + const auto serialize_pht = [&image, &pht_seg]() -> void { + auto* hdr = memory::address(pht_seg.raw_data.data()).ptr(); + + for (auto& section : image->sections) { + if (section.symbolic) { + continue; + } + const auto elf_type = section.elf_type.value_or(PT_LOAD); + const auto elf_alignment = section.elf_alignment.value_or(image->get_section_alignment()); + + decltype(ElfPhdr::p_flags) flags = 0; + if (section.characteristics.mem_read) { + flags |= PF_R; + } + if (section.characteristics.mem_write) { + flags |= PF_W; + } + if (section.characteristics.mem_execute) { + flags |= PF_X; + } + + hdr->p_type = static_cast(elf_type); + hdr->p_flags = flags; + hdr->p_offset = static_cast(section.ptr_raw_data); + hdr->p_vaddr = static_cast(section.virtual_address); + hdr->p_paddr = static_cast(section.virtual_address); + hdr->p_filesz = static_cast(section.size_raw_data); + hdr->p_memsz = static_cast(section.virtual_size); + hdr->p_align = static_cast(elf_alignment); + ++hdr; + } + }; + + /// Preparing header, it will be copied with all segments later + const auto prepare_header = [&result, &image, &pht_seg]() -> void { + const auto& last_sec = image->find_last_section(); + result.resize(last_sec.ptr_raw_data + last_sec.size_raw_data); + + auto* original_ehdr = image->rva_to_ptr(image->get_image_base()); + assert(original_ehdr != nullptr); + + /// Update PHT file offset + original_ehdr->e_phoff = static_cast(pht_seg.ptr_raw_data); + original_ehdr->e_phnum = static_cast(image->non_symbolic_segments_count()); + }; + + /// Updating .dynamic + const auto update_dynamic = [&image]() -> void { + auto* dyn = image->find_segment_with_type(PT_DYNAMIC); + if (dyn == nullptr) { + throw std::runtime_error("no dynamic?"); + } + + auto* edyn = memory::address(dyn->raw_data.data()).ptr(); + for (const auto& dynamic : image->dynamics) { + edyn->d_tag = static_cast(dynamic.tag); + edyn->d_un.d_val = static_cast(dynamic.value); + ++edyn; + } + + edyn->d_tag = DT_NULL; + edyn->d_un.d_val = DT_NULL; + }; + + switch (img_mode) { + case ImageMode::X64: { + serialize_pht.operator()(); + update_dynamic.operator()(); + prepare_header.operator()(); + break; + } + case ImageMode::X86: { + serialize_pht.operator()(); + update_dynamic.operator()(); + prepare_header.operator()(); + break; + } + default: + throw std::out_of_range("cont::elf::detail::init_header: Unsupported image mode"); + } + } +} // namespace cont::elf::detail diff --git a/src/lib/cont/elf/rebuilder/detail/elf_update_relocations.cpp b/src/lib/cont/elf/rebuilder/detail/elf_update_relocations.cpp new file mode 100644 index 0000000..167a2ef --- /dev/null +++ b/src/lib/cont/elf/rebuilder/detail/elf_update_relocations.cpp @@ -0,0 +1,118 @@ +#include "cont/elf/rebuilder/rebuilder.hpp" +#include "util/sections.hpp" + +#include + +namespace cont::elf::detail { + void update_relocations(Image* image) { + const auto img_mode = image->mode(); + + /// \todo @es3n1n: Relocations need to be sorted + /// R_X86_64_RELATIVE first, then everything else. + /// count of first R_X86_64_RELATIVE should also be updated in relascount. + + /// Wiping old relocations first + auto* rela = image->find_dynamic(DT_RELA); + auto* rela_sz = image->find_dynamic(DT_RELASZ); + if (rela != nullptr) { + if (rela_sz == nullptr || rela->value == 0) { + throw std::runtime_error("cont::elf: DT_RELA or DT_RELASZ not found or invalid"); + } + + auto* const ptr = image->rva_to_ptr(rela->value); + std::memset(ptr, 0, rela_sz->value); + } + + if (image->relocations.empty()) { + image->delete_dynamic(DT_RELA); + image->delete_dynamic(DT_RELASZ); + return; + } + + assert(rela != nullptr); + + /// We are not adding jmprel relocations, so dont erase them + + /// Allocate new segment for our relocations + auto sec_hdr = sections::get(sections::e_section_t::RELOC); + /// +1 for last DT_NULL tag + sec_hdr.size_raw_data = (image->relocations.size() + 1) * (img_mode == ImageMode::X64 ? sizeof(Elf64_Rela) : sizeof(Elf32_Rela)); + auto& sec = image->new_section(sec_hdr); + + /// Assemble the relocation structures + const auto assemble_structs = [&sec, &image, &img_mode]() -> void { + auto* out_ptr = memory::address(sec.raw_data.data()).ptr(); + + /// Separate them by type + std::unordered_multimap relocations; + for (auto& reloc : image->relocations | std::views::values) { + auto& info = reloc.info_raw; + + if (!info.has_value()) { + if (reloc.type == RelocationType::Absolute) { + continue; + } + + switch (reloc.type) { + case RelocationType::HighLow: { + // NOLINTNEXTLINE(bugprone-branch-clone) + info.emplace(static_cast((img_mode == ImageMode::X64) ? R_X86_64_RELATIVE : R_386_RELATIVE)); + break; + default: + throw std::out_of_range("cont::elf::detail::update_relocations: Unsupported relocation type for ELF"); + } + } + } + + assert(info.has_value()); + relocations.emplace(img_mode == ImageMode::X64 ? ELF64_R_TYPE(*info) : ELF32_R_TYPE(*info), &reloc); + } + + /// Assemble R_X86_64_RELATIVE first + for (auto& [rva, reloc] : relocations | std::views::filter([img_mode](const auto& pair) -> bool { + return pair.first == + (img_mode == ImageMode::X64 ? R_X86_64_RELATIVE : R_386_RELATIVE); // NOLINT(bugprone-branch-clone) + })) { + assert(reloc->info_raw.has_value()); + *out_ptr = Rela{ + .r_offset = reloc->rva.template as(), + .r_info = static_cast(*reloc->info_raw), + .r_addend = static_cast(reloc->addend.value_or(0)), + }; + ++out_ptr; + } + + /// Assemble the rest of relocations + for (auto& [rva, reloc] : relocations | std::views::filter([img_mode](const auto& pair) -> bool { + return pair.first != + (img_mode == ImageMode::X64 ? R_X86_64_RELATIVE : R_386_RELATIVE); // NOLINT(bugprone-branch-clone) + })) { + assert(reloc->info_raw.has_value()); + *out_ptr = Rela{ + .r_offset = reloc->rva.template as(), + .r_info = static_cast(*reloc->info_raw), + .r_addend = static_cast(reloc->addend.value_or(0)), + }; + ++out_ptr; + } + }; + switch (img_mode) { + case ImageMode::X64: { + assemble_structs.operator()(); + break; + } + case ImageMode::X86: { + assemble_structs.operator()(); + break; + } + default: + throw std::out_of_range("cont::elf::detail::update_relocations: unsupported image mode"); + } + + /// Update the rela entry + assert(rela != nullptr && rela_sz != nullptr); + rela->value = sec.virtual_address; + rela_sz->value = sec_hdr.size_raw_data; + logger::debug("elf: updated relocations: {} entries (now at {:#x})", image->relocations.size(), rela->value); + } +} // namespace cont::elf::detail diff --git a/src/lib/cont/elf/rebuilder/rebuilder.hpp b/src/lib/cont/elf/rebuilder/rebuilder.hpp new file mode 100644 index 0000000..d52a46a --- /dev/null +++ b/src/lib/cont/elf/rebuilder/rebuilder.hpp @@ -0,0 +1,35 @@ +#pragma once +#include "cont/elf/image.hpp" +#include + +namespace cont::elf { + namespace detail { + void update_relocations(Image* image); + void init_header(Image* image, std::vector& result); + void copy_sections(Image* image, std::vector& result); + } // namespace detail + + struct RebuilderContext { + Image* image; + }; + + [[nodiscard]] inline std::vector rebuild_elf(const RebuilderContext& ctx) { + // Init result data + // + std::vector result = {}; + auto progress = progress::Progress("elf: rebuilding", 3); + + detail::update_relocations(ctx.image); + progress.step(); + + detail::init_header(ctx.image, result); + progress.step(); + + detail::copy_sections(ctx.image, result); + progress.step(); + + // We are done here + // + return result; + } +} // namespace cont::elf diff --git a/src/lib/cont/pe/image.cpp b/src/lib/cont/pe/image.cpp new file mode 100644 index 0000000..1ed21f8 --- /dev/null +++ b/src/lib/cont/pe/image.cpp @@ -0,0 +1,235 @@ +#include "cont/pe/image.hpp" + +#include "cont/pe/rebuilder/rebuilder.hpp" +#include "magic_enum.hpp" + +namespace cont::pe { + [[nodiscard]] ImageMode Image::mode() const { + switch (x86()->get_file_header()->machine) { + case win::machine_id::amd64: + return ImageMode::X64; + case win::machine_id::i386: + return ImageMode::X86; + default: + throw std::runtime_error("cont::pe: Unsupported machine type"); + } + } + + [[nodiscard]] bool Image::verify_integrity() const { + return x86()->dos_header.e_magic == win::DOS_HDR_MAGIC; + } + + [[nodiscard]] std::size_t Image::get_image_base() const { + switch (mode()) { + case ImageMode::X64: { + return x64()->get_nt_headers()->optional_header.image_base; + } + case ImageMode::X86: { + return x86()->get_nt_headers()->optional_header.image_base; + } + default: + throw std::out_of_range("cont::pe::image_base: Unsupported image mode"); + } + } + + [[nodiscard]] std::size_t Image::get_section_alignment() const { + switch (mode()) { + case ImageMode::X64: { + return x64()->get_nt_headers()->optional_header.section_alignment; + } + case ImageMode::X86: { + return x86()->get_nt_headers()->optional_header.section_alignment; + } + default: + throw std::out_of_range("cont::pe::get_section_alignment: Unsupported image mode"); + } + } + + [[nodiscard]] std::size_t Image::get_file_alignment() const { + switch (mode()) { + case ImageMode::X64: { + return x64()->get_nt_headers()->optional_header.file_alignment; + } + case ImageMode::X86: { + return x86()->get_nt_headers()->optional_header.file_alignment; + } + default: + throw std::out_of_range("cont::pe::file_alignment: Unsupported image mode"); + } + } + + [[nodiscard]] std::size_t Image::get_base_of_code() const { + switch (mode()) { + case ImageMode::X64: { + return x64()->get_nt_headers()->optional_header.base_of_code; + } + case ImageMode::X86: { + return x86()->get_nt_headers()->optional_header.base_of_code; + } + default: + throw std::out_of_range("cont::pe::base_of_code: Unsupported image mode"); + } + } + + [[nodiscard]] std::vector Image::rebuild_image() { + auto ctx = RebuilderContext{.image = this}; + return rebuild_pe(ctx); + } + + void Image::update_sections() { + const auto proceed = [this](const Ty* raw_image) -> void { + // Obtaining stuff that would be needed + // + const auto* nt_hdr = raw_image->get_nt_headers(); + + // Reserving the num of sections + // + sections.clear(); + sections.reserve(nt_hdr->file_header.num_sections); + + for (std::size_t i = 0; i < nt_hdr->file_header.num_sections; ++i) { + // Getting a section and validating it + // + const win::section_header_t* section = nt_hdr->get_section(i); + if (section == nullptr) { + continue; + } + + // Inserting a section to the result array + // + auto& new_elem = sections.emplace_back(to_cont(*section)); + new_elem.raw_data.resize(new_elem.size_raw_data, 0); + + // NOLINTNEXTLINE + std::memcpy(new_elem.raw_data.data(), reinterpret_cast(reinterpret_cast(raw_image) + new_elem.ptr_raw_data), // + new_elem.size_raw_data); + } + + // Signalising that we successfully parsed sections + // + logger::debug("pe: parsed {} sections", sections.size()); + + // Sort by virtual address + // + const struct { + bool operator()(const Section& lhs, const Section& rhs) const { + return lhs.virtual_address < rhs.virtual_address; + } + } comp; + std::sort(sections.begin(), sections.end(), comp); + + // Marking directories + // + for (auto dir_id : magic_enum::enum_values()) { + // Trying to find the header of the directory, skipping if not present + // + auto dir_hdr = raw_image->get_directory(dir_id); + if (!dir_hdr || !dir_hdr->present()) { + continue; + } + + // Looking up the section by rva and changing flag + // + auto* sec = rva_to_section(dir_hdr->rva); + sec->set_contained_dir(to_cont(dir_id), dir_hdr->rva - sec->virtual_address, dir_hdr->size); + } + }; + + switch (mode()) { + case ImageMode::X64: + proceed(x64()); + break; + case ImageMode::X86: + proceed(x86()); + break; + default: + throw std::out_of_range("cont::pe::update_sections: Unsupported image mode"); + } + } + + void Image::update_relocations() { + const auto proceed = [this](const Ty* raw_image) -> void { + // Obtaining a pointer to reloc directory header + // + const win::data_directory_t* reloc_hdr = raw_image->get_directory(win::directory_id::directory_entry_basereloc); + if ((reloc_hdr == nullptr) || !reloc_hdr->present()) [[unlikely]] { + logger::warn("pe: relocation directory header does not present?"); + return; + } + + // Obtaining a pointer to reloc directory + // + const win::reloc_directory_t* base_reloc = rva_to_ptr(reloc_hdr->rva); + const auto reloc_size = ptr_size(); + + // Iterating over reloc blocks + // + for (const auto* reloc_block = &base_reloc->first_block; // + (reloc_block != nullptr) && (reloc_block->size_block != 0U) && (reloc_block->base_rva != 0U); // + reloc_block = reloc_block->next()) { + // Iterating over reloc entries + // + for (const auto& [offset, type] : *reloc_block) { + // Skip ignored relocations + // \todo @es3n1n: remove me + if (type == win::reloc_type_id::rel_based_absolute) { + continue; + } + + // Inserting parsed reloc data + // + auto rva = static_cast(reloc_block->base_rva) + memory::address(offset); + + // Just to be sure + // + if (relocations.contains(rva)) [[unlikely]] { + throw std::runtime_error(std::format("pe: duplicated {:#x} rva entry", rva)); + } + + // Inserting relocation info + // + relocations[rva] = Relocation{ + .rva = rva, // + .size = static_cast(reloc_size), // + .type = static_cast(type), // + }; + } + } + + logger::debug("pe: parsed total number of {} relocations", relocations.size()); + }; + + switch (mode()) { + case ImageMode::X64: + proceed(x64()); + break; + case ImageMode::X86: + proceed(x86()); + break; + default: + throw std::out_of_range("cont::pe::update_relocations: Unsupported image mode"); + } + } + + [[nodiscard]] win::data_directory_t* Image::get_directory(win::directory_id dir_id) { + const auto proceed = [dir_id](Ty* raw_image) -> win::data_directory_t* { + auto nt_hdrs = raw_image->get_nt_headers(); + if (nt_hdrs->optional_header.num_data_directories <= dir_id) { + return nullptr; + } + + return &nt_hdrs->optional_header.data_directories.entries[dir_id]; + }; + + switch (mode()) { + case ImageMode::X64: { + return proceed(x64()); + } + case ImageMode::X86: { + return proceed(x86()); + } + default: + throw std::out_of_range("cont::pe::get_directory: Unsupported image mode"); + } + } +} // namespace cont::pe \ No newline at end of file diff --git a/src/lib/cont/pe/image.hpp b/src/lib/cont/pe/image.hpp new file mode 100644 index 0000000..cbb69c2 --- /dev/null +++ b/src/lib/cont/pe/image.hpp @@ -0,0 +1,37 @@ +#pragma once +#include +#include + +namespace cont::pe { + class Image final : public ImageBase { + public: + explicit Image(const memory::address raw_image): ImageBase(ContImageType::PE, raw_image) { + initialize(); + } + [[nodiscard]] ImageMode mode() const override; + [[nodiscard]] bool verify_integrity() const override; + + [[nodiscard]] std::size_t get_image_base() const override; + [[nodiscard]] std::size_t get_section_alignment() const override; + [[nodiscard]] std::size_t get_file_alignment() const override; + [[nodiscard]] std::size_t get_base_of_code() const; + + [[nodiscard]] std::vector rebuild_image() override; + + void update_sections() override; + void update_relocations() override; + + [[nodiscard]] win::data_directory_t* get_directory(win::directory_id dir_id); + + private: + [[nodiscard]] win::image_x64_t* x64() const { + return raw_image_.as(); + } + + [[nodiscard]] win::image_x86_t* x86() const { + return raw_image_.as(); + } + }; + + template concept AnyRawImage = traits::is_any_of_v; +} // namespace cont::pe diff --git a/src/lib/cont/pe/rebuilder/detail/copy_sections.cpp b/src/lib/cont/pe/rebuilder/detail/copy_sections.cpp new file mode 100644 index 0000000..28c461d --- /dev/null +++ b/src/lib/cont/pe/rebuilder/detail/copy_sections.cpp @@ -0,0 +1,186 @@ +#include "cont/pe/rebuilder/rebuilder.hpp" + +#include + +namespace cont::pe::detail { + namespace { + win::section_header_t assemble_section_header(const Section& section) { + win::section_header_t result{}; + + std::ranges::copy(section.name.value(), reinterpret_cast(result.name.short_name)); + result.virtual_size = static_cast(section.virtual_size); + result.virtual_address = static_cast(section.virtual_address); + result.size_raw_data = static_cast(section.size_raw_data); + result.ptr_raw_data = static_cast(section.ptr_raw_data); + + /// \fixme @es3n1n: This looks ugly + result.characteristics.cnt_code = section.characteristics.cnt_code; + result.characteristics.cnt_init_data = section.characteristics.cnt_init_data; + result.characteristics.cnt_uninit_data = section.characteristics.cnt_uninit_data; + result.characteristics.lnk_info = section.characteristics.lnk_info; + result.characteristics.lnk_remove = section.characteristics.lnk_remove; + result.characteristics.lnk_comdat = section.characteristics.lnk_comdat; + result.characteristics.no_defer_spec_exc = section.characteristics.no_defer_spec_exc; + result.characteristics.mem_far = section.characteristics.mem_far; + result.characteristics.mem_purgeable = section.characteristics.mem_purgeable; + result.characteristics.mem_locked = section.characteristics.mem_locked; + result.characteristics.mem_preload = section.characteristics.mem_preload; + result.characteristics.alignment = section.characteristics.alignment; + result.characteristics.lnk_nreloc_ovfl = section.characteristics.lnk_nreloc_ovfl; + result.characteristics.mem_discardable = section.characteristics.mem_discardable; + result.characteristics.mem_not_cached = section.characteristics.mem_not_cached; + result.characteristics.mem_not_paged = section.characteristics.mem_not_paged; + result.characteristics.mem_shared = section.characteristics.mem_shared; + result.characteristics.mem_execute = section.characteristics.mem_execute; + result.characteristics.mem_read = section.characteristics.mem_read; + result.characteristics.mem_write = section.characteristics.mem_write; + return result; + } + + template + void copy_sections_(Image* image, std::vector& data) { + /// Casting our buffer as the raw image + auto* out_img = reinterpret_cast(data.data()); + auto* nt_headers = out_img->get_nt_headers(); + auto* file_header = &nt_headers->file_header; + auto* optional_header = &nt_headers->optional_header; + + /// Obtaining sections pointer within the nt headers + auto* sections = out_img->get_nt_headers()->get_sections(); + + /// Raligning sections + image->realign_sections(); + + /// Validating that there's enough space for our sections + const auto sections_start = memory::address{sections}; + auto header_end = memory::address{out_img}.offset(optional_header->size_headers); + if (sections_start + (sizeof(win::section_header_t) * (image->sections.size() + 1)) > header_end) [[unlikely]] { + throw std::runtime_error("pe: rebuilder: unable to fit new sections"); + } + + /// Erasing previous data directories info + std::memset(&optional_header->data_directories, 0, sizeof(optional_header->data_directories)); + + /// Iterating over the max value of sections + /// If the amount of our sections is less than the number of sections + /// within the PE that we're rebuilding, we would essentially need to + /// erase all the other sections that presents within the PE + for (std::size_t i = 0; i < std::max(static_cast(file_header->num_sections), image->sections.size()); ++i) { + /// Erasing previous data first + std::memset(§ions[i], 0, sizeof(win::section_header_t)); + + /// If we only need to erase the prev section data + if (i >= image->sections.size()) { + continue; + } + + /// Obtaining our section info + auto& section = image->sections.at(i); + + /// Assembling the new section header and copying it + auto sec_header = assemble_section_header(section); + if (auto write_res = memory::address{§ions[i]}.write(memory::address{&sec_header}.as(), sizeof(win::section_header_t)); + !write_res.has_value()) { + throw std::runtime_error("pe: rebuilder: unable to write section header"); + } + + /// Copying section data, if needed + if (!section.raw_data.empty()) { + auto sec_ptr = memory::address{data.data()}.offset(static_cast(section.ptr_raw_data)); + if (auto write_res = sec_ptr.write(section.raw_data.data(), section.raw_data.size()); !write_res.has_value()) { + throw std::runtime_error("pe: rebuilder: unable to write section raw data"); + } + } + + /// Updating the data directories + for (const auto dir_type : magic_enum::enum_values()) { + if (dir_type == DirectoryType::MAX_LENGTH) { + continue; + } + const auto props = section.directory_info(dir_type); + if (!props.has_value()) { + continue; + } + + const win::data_directory_t dir_hdr = { + .rva = static_cast(section.virtual_address + props->offset), + .size = static_cast(props->size), + }; + + switch (dir_type) { + case DirectoryType::Export: + optional_header->data_directories.export_directory = dir_hdr; + break; + case DirectoryType::Import: + optional_header->data_directories.import_directory = dir_hdr; + break; + case DirectoryType::Resource: + optional_header->data_directories.resource_directory = dir_hdr; + break; + case DirectoryType::Exception: + optional_header->data_directories.exception_directory = dir_hdr; + break; + case DirectoryType::Security: + optional_header->data_directories.security_directory = win::raw_data_directory_t{ + .ptr_raw_data = dir_hdr.rva, + .size = dir_hdr.size, + }; + break; + case DirectoryType::Reloc: + optional_header->data_directories.basereloc_directory = dir_hdr; + break; + case DirectoryType::Debug: + optional_header->data_directories.debug_directory = dir_hdr; + break; + case DirectoryType::Architecture: + if constexpr (!std::is_same_v) { + optional_header->data_directories.architecture_directory = dir_hdr; + } + break; + case DirectoryType::GlobalPtr: + optional_header->data_directories.globalptr_directory = dir_hdr; + break; + case DirectoryType::Tls: + optional_header->data_directories.tls_directory = dir_hdr; + break; + case DirectoryType::LoadConfig: + optional_header->data_directories.load_config_directory = dir_hdr; + break; + case DirectoryType::BoundImport: + optional_header->data_directories.bound_import_directory = dir_hdr; + break; + case DirectoryType::Iat: + optional_header->data_directories.iat_directory = dir_hdr; + break; + case DirectoryType::DelayImport: + optional_header->data_directories.delay_import_directory = dir_hdr; + break; + case DirectoryType::ComDescriptor: + optional_header->data_directories.com_descriptor_directory = dir_hdr; + break; + default: + throw std::out_of_range("pe: rebuilder: unsupported directory type"); + } + } + } + + /// Updating the sections count + file_header->num_sections = static_cast(image->sections.size()); + } + } // namespace + + void copy_sections(Image* image, std::vector& data) { + switch (image->mode()) { + case ImageMode::X64: { + copy_sections_(image, data); + break; + } + case ImageMode::X86: { + copy_sections_(image, data); + break; + } + default: + throw std::out_of_range("cont::pe::detail::copy_sections: Unsupported image mode"); + } + } +} // namespace cont::pe::detail \ No newline at end of file diff --git a/src/lib/pe/rebuilder/detail/erase_metadata.cpp b/src/lib/cont/pe/rebuilder/detail/erase_metadata.cpp similarity index 50% rename from src/lib/pe/rebuilder/detail/erase_metadata.cpp rename to src/lib/cont/pe/rebuilder/detail/erase_metadata.cpp index 673671e..ec86020 100644 --- a/src/lib/pe/rebuilder/detail/erase_metadata.cpp +++ b/src/lib/cont/pe/rebuilder/detail/erase_metadata.cpp @@ -1,10 +1,11 @@ -#include "pe/rebuilder/rebuilder.hpp" +#include "cont/pe/image.hpp" +#include "cont/pe/rebuilder/rebuilder.hpp" -namespace pe::detail { +namespace cont::pe::detail { namespace { - template - void erase_metadata_(Img* image, std::vector& data) { - auto* out_img = detail::buffer_pointer>(data); + template + void erase_metadata_(std::vector& data) { + auto* out_img = reinterpret_cast(data.data()); auto* nt_headers = out_img->get_nt_headers(); auto* file_header = &nt_headers->file_header; auto* optional_header = &nt_headers->optional_header; @@ -23,20 +24,29 @@ namespace pe::detail { /// Wipe debug directory /// \todo: wipe pdb path - std::memset(&optional_header->data_directories.debug_directory, 0, optional_header->data_directories.debug_directory.size); + if (auto& dbg_dir = optional_header->data_directories.debug_directory; dbg_dir.size > 0) { + std::memset(out_img->rva_to_ptr(dbg_dir.rva), 0, dbg_dir.size); + dbg_dir.rva = 0; + dbg_dir.size = 0; + } /// Wipe checksum std::memset(&optional_header->checksum, 0, sizeof(optional_header->checksum)); - - /// Wipe section names - auto* sections = nt_headers->get_sections(); - for (std::size_t i = 0; i < std::max(static_cast(file_header->num_sections), image->sections.size()); ++i) { - std::memset(§ions[i].name, 0, sizeof(sections[i].name)); - } } } // namespace - void erase_metadata(const ImgWrapped image, std::vector& data) { - UNWRAP_IMAGE(void, erase_metadata_); + void erase_metadata(const Image* image, std::vector& data) { + switch (image->mode()) { + case ImageMode::X64: { + erase_metadata_(data); + break; + } + case ImageMode::X86: { + erase_metadata_(data); + break; + } + default: + throw std::out_of_range("cont::pe::detail::erase_metadata: Unsupported image mode"); + } } -} // namespace pe::detail +} // namespace cont::pe::detail diff --git a/src/lib/pe/rebuilder/detail/init_header.cpp b/src/lib/cont/pe/rebuilder/detail/init_header.cpp similarity index 51% rename from src/lib/pe/rebuilder/detail/init_header.cpp rename to src/lib/cont/pe/rebuilder/detail/init_header.cpp index 39834b0..b4fe944 100644 --- a/src/lib/pe/rebuilder/detail/init_header.cpp +++ b/src/lib/cont/pe/rebuilder/detail/init_header.cpp @@ -1,12 +1,12 @@ -#include "pe/rebuilder/rebuilder.hpp" +#include "cont/pe/rebuilder/rebuilder.hpp" -namespace pe::detail { +namespace cont::pe::detail { namespace { - template - void init_header_(Img* image, std::vector& data) { + template + void init_header_(Image* image, std::vector& data) { // Obtaining header structs // - auto* nt_headers = image->raw_image->get_nt_headers(); + auto* nt_headers = image->raw_image().ptr()->get_nt_headers(); auto* optional_header = &nt_headers->optional_header; // Obtaining some other stuff from the header @@ -22,7 +22,7 @@ namespace pe::detail { // \todo: @es3n1n: size_code, size_init_data, size_uninit_data, base_of_code, num_rva_sizes // \note: @es3n1n: the sections count field is updated within the `copy_sections` pass! // - optional_header->size_image = virtual_image_size; + optional_header->size_image = static_cast(virtual_image_size); // Reserving header size // @@ -30,11 +30,22 @@ namespace pe::detail { // Copying the original header // NOLINTNEXTLINE - std::memcpy(data.data(), image->raw_image, optional_header->size_headers); + std::memcpy(data.data(), image->raw_image().ptr(), optional_header->size_headers); } } // namespace - void init_header(const ImgWrapped image, std::vector& data) { - UNWRAP_IMAGE(void, init_header_); + void init_header(Image* image, std::vector& data) { + switch (image->mode()) { + case ImageMode::X64: { + init_header_(image, data); + break; + } + case ImageMode::X86: { + init_header_(image, data); + break; + } + default: + throw std::out_of_range("cont::pe::detail::init_header: Unsupported image mode"); + } } -} // namespace pe::detail +} // namespace cont::pe::detail diff --git a/src/lib/cont/pe/rebuilder/detail/update_checksum.cpp b/src/lib/cont/pe/rebuilder/detail/update_checksum.cpp new file mode 100644 index 0000000..13eff15 --- /dev/null +++ b/src/lib/cont/pe/rebuilder/detail/update_checksum.cpp @@ -0,0 +1,29 @@ +#include "cont/pe/rebuilder/rebuilder.hpp" + +namespace cont::pe::detail { + namespace { + template + void update_checksum_(std::vector& data) { + /// Get the headers + auto* out_img = reinterpret_cast(data.data()); + + /// Update checksum + out_img->update_checksum(data.size()); + } + } // namespace + + void update_checksum(const Image* image, std::vector& data) { + switch (image->mode()) { + case ImageMode::X64: { + update_checksum_(data); + break; + } + case ImageMode::X86: { + update_checksum_(data); + break; + } + default: + throw std::out_of_range("cont::pe::detail::update_checksum: Unsupported image mode"); + } + } +} // namespace cont::pe::detail diff --git a/src/lib/pe/rebuilder/detail/update_relocations.cpp b/src/lib/cont/pe/rebuilder/detail/update_relocations.cpp similarity index 70% rename from src/lib/pe/rebuilder/detail/update_relocations.cpp rename to src/lib/cont/pe/rebuilder/detail/update_relocations.cpp index 836e930..aa8f013 100644 --- a/src/lib/pe/rebuilder/detail/update_relocations.cpp +++ b/src/lib/cont/pe/rebuilder/detail/update_relocations.cpp @@ -1,20 +1,20 @@ -#include "pe/rebuilder/rebuilder.hpp" +#include "cont/pe/rebuilder/rebuilder.hpp" #include "util/format.hpp" #include -namespace pe::detail { +namespace cont::pe::detail { namespace { constexpr std::size_t kRelocBlockAlignment = 0x1000; // Erasing previous relocations from the binary // - template - void erase_relocations(Img* image) { + template + void erase_relocations(Image* image) { // Looking for the section that contains relocations // - auto reloc_section = std::ranges::find_if(image->sections, [](const section_t& sec) -> bool { // - return sec.contains_dir.reloc.has_value(); + auto reloc_section = std::ranges::find_if(image->sections, [](const Section& sec) -> bool { // + return sec.directory_info(DirectoryType::Reloc).has_value(); }); // No relocation dir? @@ -25,8 +25,8 @@ namespace pe::detail { // Obtaining reloc entry offset from the base of section // - auto reloc_offset = image->raw_image->get_directory(win::directory_id::directory_entry_basereloc)->rva; - reloc_offset -= reloc_section->virtual_address; + auto reloc_offset = image->raw_image().ptr()->get_directory(win::directory_id::directory_entry_basereloc)->rva; + reloc_offset -= static_cast(reloc_section->virtual_address); // Obtaining reloc directory and iterating over blocks in order to get the last block // @@ -61,8 +61,7 @@ namespace pe::detail { // Assembling the new reloc section // - template - void assemble_relocations(Img* image) { + void assemble_relocations(Image* image) { // No relocations? // if (image->relocations.empty()) [[unlikely]] { @@ -79,8 +78,8 @@ namespace pe::detail { // relocation {rva=0x3FFFF} // etc // - std::unordered_map> blocks; - std::size_t section_size = 0ULL; + std::unordered_map> blocks; + auto section_header = sections::get(sections::e_section_t::RELOC); // Iterating over relocations and obtaining start RVAs, // Estimating section size @@ -91,30 +90,29 @@ namespace pe::detail { // Accounting new block header if we're creating one // if (!blocks.contains(aligned_rva)) { - section_size += sizeof(win::reloc_block_t); + section_header.size_raw_data += sizeof(win::reloc_block_t); } // Prepending relocation to the block // - blocks[aligned_rva].emplace_back(std::move(relocation)); + blocks[aligned_rva].emplace_back(relocation); // Accounting entry // - section_size += sizeof(win::reloc_entry_t); + section_header.size_raw_data += sizeof(win::reloc_entry_t); } // Obtaining a pointer to the directory header // - auto* dir_header = image->get_directory(win::directory_id::directory_entry_basereloc); - if (dir_header == nullptr) { + if (const auto* dir_header = image->get_directory(win::directory_id::directory_entry_basereloc); dir_header == nullptr) { throw std::runtime_error("pe: rebuilder: .reloc header not found"); } // Inserting the new section with our relocations // - auto& new_section = image->new_section(sections::e_section_t::RELOC, section_size); + auto& new_section = image->new_section(section_header); auto section_data = memory::address{new_section.raw_data.data()}; - auto section_end = section_data.offset(new_section.raw_data.size()); + const auto section_end = section_data.offset(static_cast(new_section.raw_data.size())); // Serializing reloc entries // @@ -138,7 +136,7 @@ namespace pe::detail { // const auto reloc_encoded = win::reloc_entry_t{ .offset = (relocation.rva - rva).as(), - .type = relocation.type, + .type = static_cast(relocation.type), }; // Writing it @@ -151,17 +149,29 @@ namespace pe::detail { // Mark as sec with relocs // - new_section.set_contained_dir(win::directory_id::directory_entry_basereloc, 0, section_size); + new_section.set_contained_dir(DirectoryType::Reloc, 0, new_section.size_raw_data); } - template - void update_relocations_(Img* image, std::vector& data [[maybe_unused]]) { - erase_relocations(image); + template + void update_relocations_(Image* image) { + erase_relocations(image); assemble_relocations(image); } } // namespace - void update_relocations(const ImgWrapped image, std::vector& data) { - UNWRAP_IMAGE(void, update_relocations_); + void update_relocations(Image* image, const std::vector& data) { + std::ignore = data; + switch (image->mode()) { + case ImageMode::X64: { + update_relocations_(image); + break; + } + case ImageMode::X86: { + update_relocations_(image); + break; + } + default: + throw std::out_of_range("cont::pe::detail::update_relocations: Unsupported image mode"); + } } -} // namespace pe::detail \ No newline at end of file +} // namespace cont::pe::detail \ No newline at end of file diff --git a/src/lib/cont/pe/rebuilder/rebuilder.hpp b/src/lib/cont/pe/rebuilder/rebuilder.hpp new file mode 100644 index 0000000..2c0b73e --- /dev/null +++ b/src/lib/cont/pe/rebuilder/rebuilder.hpp @@ -0,0 +1,53 @@ +#pragma once +#include "cont/pe/image.hpp" +#include + +namespace cont::pe { + namespace detail { + void update_relocations(Image* /*image*/, const std::vector& data); + void init_header(Image* image, std::vector& data); + void copy_sections(Image* image, std::vector& data); + void update_checksum(const Image* image, std::vector& data); + void erase_metadata(const Image* image, std::vector& data); + } // namespace detail + + struct RebuilderContext { + Image* image; + }; + + [[nodiscard]] inline std::vector rebuild_pe(RebuilderContext& ctx) { + // Init result data + // + std::vector result = {}; + auto progress = progress::Progress("pe: rebuilding", 5); + + // Updating .reloc section + // + detail::update_relocations(ctx.image, result); + progress.step(); + + // Reserving and copying the original header first + // + detail::init_header(ctx.image, result); + progress.step(); + + // Copying sections + // + detail::copy_sections(ctx.image, result); + progress.step(); + + // Update checksum + // + detail::update_checksum(ctx.image, result); + progress.step(); + + /// Wipe metadata + // + detail::erase_metadata(ctx.image, result); + progress.step(); + + // We are done here + // + return result; + } +} // namespace cont::pe diff --git a/src/lib/easm/assembler/assembler.cpp b/src/lib/easm/assembler/assembler.cpp index 1d6e9fe..e739337 100644 --- a/src/lib/easm/assembler/assembler.cpp +++ b/src/lib/easm/assembler/assembler.cpp @@ -25,7 +25,8 @@ namespace easm { const auto& insn_info = node_insn->getDetail(program.getMode()); if (!insn_info) { - throw std::runtime_error("Unable to estimate program size: unable to get if instr info"); + throw std::runtime_error(std::format("Unable to estimate program size: unable to get if instr info {}", + ZydisMnemonicGetString(static_cast(node_insn->getMnemonic().value())))); } result += insn_info->getLength(); diff --git a/src/lib/easm/debug/debug.hpp b/src/lib/easm/debug/debug.hpp index c8fe77d..4de0eb2 100644 --- a/src/lib/easm/debug/debug.hpp +++ b/src/lib/easm/debug/debug.hpp @@ -2,13 +2,13 @@ #include "analysis/common/debug.hpp" namespace easm { - inline void dump_program(const zasm::Program& program) { + inline void dump_program(const zasm::Program& program, const bool errors_only = true) { /// Iterating over the program nodes /// for (auto* node = program.getHead(); node != nullptr; node = node->getNext()) { // Handling `zasm::Data` // - if (const auto* node_data = node->getIf(); node_data != nullptr) { + if (const auto* node_data = node->getIf(); !errors_only && node_data != nullptr) { logger::info("Data: {:#x}", node_data->valueAsU64()); continue; } @@ -17,11 +17,18 @@ namespace easm { // Handling `zasm::Instruction` // if (auto* node_insn = node->getIf(); node_insn != nullptr) { + bool is_error = false; if (const auto& insn_info = node_insn->getDetail(program.getMode()); !insn_info) { - throw std::runtime_error("Unable to debug dump program: unable to get if instr info"); + is_error = true; + logger::error<1>("unable to get instn info -- this is broken"); + } + + if (!is_error && errors_only) { + continue; } logger::info("Instruction: {}", ZydisMnemonicGetString(static_cast(node_insn->getMnemonic().value()))); + if (const auto ops_count = node_insn->getOperandCount(); ops_count > 0) { logger::info<1>("Operands:"); @@ -42,13 +49,11 @@ namespace easm { if (const auto* p_mem = node_insn->getOperandIf(i); p_mem) { const auto reg = static_cast(p_mem->getBase().getId()); + logger::info<3>("expr: [{} + {:#x}]", reg == 0 ? "none" : ZydisRegisterGetString(reg), p_mem->getDisplacement()); + } - std::string reg_str = reg == 0 ? "none" : std::to_string(reg); - if (p_mem->getBase().isIP()) { - reg_str = "ip"; - } - - logger::info<3>("expr: [{} + {:#x}]", reg_str, p_mem->getDisplacement()); + if (const auto* p_reg = node_insn->getOperandIf(i); p_reg) { + logger::info<3>("name: {}", ZydisRegisterGetString(static_cast(p_reg->getId()))); } } } diff --git a/src/lib/easm/misc/misc.hpp b/src/lib/easm/misc/misc.hpp index cb98b62..79ca614 100644 --- a/src/lib/easm/misc/misc.hpp +++ b/src/lib/easm/misc/misc.hpp @@ -1,5 +1,5 @@ #pragma once -#include "pe/pe.hpp" +#include "cont/pe/image.hpp" #include #include @@ -8,13 +8,17 @@ namespace easm { constexpr size_t kMaxEntryInstructionSize = 5; // jump in our case - template - constexpr zasm::x86::Gp sp_for_arch() { - if constexpr (pe::is_x64_v) { + constexpr zasm::x86::Gp sp_for_arch(const zasm::MachineMode machine_mode) { + switch (machine_mode) { + case zasm::MachineMode::AMD64: { return zasm::x86::rsp; - } else { + } + case zasm::MachineMode::I386: { return zasm::x86::esp; } + default: + throw std::out_of_range("easm: sp_for_arch: Unsupported machine mode"); + } } constexpr zasm::BitSize sp_size_for_arch(const zasm::MachineMode machine_mode) { @@ -29,22 +33,18 @@ namespace easm { } } - template - constexpr zasm::BitSize sp_size_for_arch() { - if constexpr (pe::is_x64_v) { - return sp_size_for_arch(zasm::MachineMode::AMD64); - } else { - return sp_size_for_arch(zasm::MachineMode::I386); - } - } - - template - constexpr zasm::Mem ptr(TArgs... args) { - if constexpr (pe::is_x64_v) { + template + constexpr zasm::Mem ptr(const zasm::MachineMode machine_mode, TArgs... args) { + switch (machine_mode) { + case zasm::MachineMode::AMD64: { return zasm::x86::qword_ptr(std::forward(args)...); - } else { + } + case zasm::MachineMode::I386: { return zasm::x86::dword_ptr(std::forward(args)...); } + default: + throw std::out_of_range("easm: ptr: Unsupported image mode"); + } } inline bool is_jcc_or_jmp(const zasm::Instruction& insn) { @@ -78,14 +78,21 @@ namespace easm { } bool result = false; - for (std::size_t i = 0; i < insn.getOperandCount() && !result; ++i) { - if (const auto* op_reg = insn.getOperandIf(i); op_reg != nullptr) { - result = op_reg->isIP(); - continue; - } - if (const auto* op_mem = insn.getOperandIf(i); op_mem != nullptr) { - result = op_mem->getBase().isIP(); + if (auto detail = insn.getDetail(zasm::MachineMode::AMD64); detail.hasValue()) { + for (const auto& op : detail->getOperands()) { + if (result) { + break; + } + + if (const auto* op_reg = op.getIf(); op_reg != nullptr) { + result = op_reg->isIP(); + continue; + } + + if (const auto* op_mem = op.getIf(); op_mem != nullptr) { + result = op_mem->getBase().isIP(); + } } } @@ -205,7 +212,12 @@ namespace easm { inline void assert_operand_size(const zasm::MachineMode machine_mode [[maybe_unused]], const zasm::Instruction* insn [[maybe_unused]], const std::size_t index [[maybe_unused]], const zasm::Reg reg [[maybe_unused]]) { - assert(get_operand_size(machine_mode, insn, index) == reg.getBitSize(machine_mode)); + const auto operand_size = get_operand_size(machine_mode, insn, index); + const auto reg_size = reg.getBitSize(machine_mode); + if (operand_size != reg_size) { + assert(false); + throw std::runtime_error("assert_operand_size: Operand size does not match register size"); + } } inline bool is_sp(const zasm::MachineMode machine_mode, const zasm::Reg reg) { @@ -248,4 +260,27 @@ namespace easm { return result; } + + inline bool affects_flags(const zasm::Instruction& insn) { + bool result = false; + + if (auto detail = insn.getDetail(zasm::MachineMode::AMD64); detail.hasValue()) { + for (const auto& op : detail->getOperands()) { + if (result) { + break; + } + + if (const auto* op_reg = op.getIf(); op_reg != nullptr) { + result = op_reg->getClass() == static_cast(ZYDIS_REGCLASS_FLAGS); + continue; + } + + if (const auto* op_mem = op.getIf(); op_mem != nullptr) { + result = op_mem->getBase().getClass() == static_cast(ZYDIS_REGCLASS_FLAGS); + } + } + } + + return result; + } } // namespace easm diff --git a/src/lib/func_parser/common/sanitizer.hpp b/src/lib/func_parser/common/sanitizer.hpp index 324bebd..bc5ad9e 100644 --- a/src/lib/func_parser/common/sanitizer.hpp +++ b/src/lib/func_parser/common/sanitizer.hpp @@ -1,14 +1,13 @@ #pragma once +#include "cont/base.hpp" #include "func_parser/common/common.hpp" -#include "pe/pe.hpp" #include namespace func_parser::sanitizer { - template - void sanitize_function_list(function_list_t& items, const Img* image) noexcept { + inline void sanitize_function_list(function_list_t& items, const cont::ImageBase* image) noexcept { // Obtaining executable sections // - const auto exec_sections = image->find_sections_if([](const pe::section_t& sec) -> bool { // + const auto exec_sections = image->find_sections_if([](const cont::Section& sec) -> bool { // return sec.characteristics.mem_execute; }); @@ -25,7 +24,7 @@ namespace func_parser::sanitizer { // Checking whether function is in an executable section or not // bool in_exec_mem = false; - for (auto& sec : exec_sections) { + for (const auto& sec : exec_sections) { in_exec_mem = item.rva >= sec.virtual_address && item.rva <= (sec.virtual_address + sec.virtual_size); if (in_exec_mem) { diff --git a/src/lib/func_parser/map/map.cpp b/src/lib/func_parser/map/map.cpp index 3154d5d..3191186 100644 --- a/src/lib/func_parser/map/map.cpp +++ b/src/lib/func_parser/map/map.cpp @@ -17,12 +17,12 @@ namespace func_parser::map { } } // namespace - function_list_t discover_functions(const std::filesystem::path& map_path, const std::vector& sections) { + function_list_t discover_functions(const std::filesystem::path& map_path, const std::vector& sections) { // Reading map file // const auto map_content = files::read_file(map_path); if (!map_content.has_value() || map_content->empty()) { - throw std::runtime_error("Empty map file"); + return {}; } // Converting to string stream diff --git a/src/lib/func_parser/map/map.hpp b/src/lib/func_parser/map/map.hpp index ce96089..14638ed 100644 --- a/src/lib/func_parser/map/map.hpp +++ b/src/lib/func_parser/map/map.hpp @@ -1,8 +1,8 @@ #pragma once +#include "cont/base.hpp" #include "func_parser/common/common.hpp" -#include "pe/common/types.hpp" #include namespace func_parser::map { - function_list_t discover_functions(const std::filesystem::path& map_path, const std::vector& sections); + function_list_t discover_functions(const std::filesystem::path& map_path, const std::vector& sections); } // namespace func_parser::map diff --git a/src/lib/func_parser/parser.cpp b/src/lib/func_parser/parser.cpp index 4e36930..2ebd8e3 100644 --- a/src/lib/func_parser/parser.cpp +++ b/src/lib/func_parser/parser.cpp @@ -1,12 +1,13 @@ #include "func_parser/parser.hpp" + +#include "cont/pe/image.hpp" #include "func_parser/common/combiner.hpp" #include "func_parser/map/map.hpp" #include "func_parser/pdb/pdb.hpp" #include namespace func_parser { - template - void Instance::collect_functions() { + void Instance::collect_functions() { // Parsing from all sources possible // parse(); @@ -28,17 +29,20 @@ namespace func_parser { logger::debug("func_parser: discovered {} functions", function_list_.size()); } - template - void Instance::parse() { - parse_pdb(); + void Instance::parse() { + if (image_->image_type() == cont::ContImageType::PE) { + parse_pdb(); + } progress_step(); parse_map(); progress_step(); + + parse_manual(); + progress_step(); } - template - void Instance::parse_pdb() { + void Instance::parse_pdb() { // If force disabled // if (!config_.pdb_enabled) { @@ -47,7 +51,7 @@ namespace func_parser { // Obtaining base of code // - const auto base_of_code = image_->raw_image->get_nt_headers()->optional_header.base_of_code; + const auto base_of_code = reinterpret_cast(image_)->get_base_of_code(); // Trying to parse from a custom pdb path first // @@ -58,10 +62,10 @@ namespace func_parser { } // Trying to parse from a codeview path - // - if (push(pdb::discover_functions(image_->find_codeview70(), base_of_code))) { - return; - } + // \todo @es3n1n: implement + // if (push(pdb::discover_functions(image_->find_codeview70(), base_of_code))) { + // return; + // } // Trying to find .pdb near the executable // @@ -70,8 +74,7 @@ namespace func_parser { push(pdb::discover_functions(pdb_path, base_of_code)); } - template - void Instance::parse_map() { + void Instance::parse_map() { // If force disabled // if (!config_.map_enabled) { @@ -93,5 +96,20 @@ namespace func_parser { push(map::discover_functions(map_path, image_->sections)); } - PE_DECL_TEMPLATE_CLASSES(Instance); + void Instance::parse_manual() { + std::vector functions = {}; + + for (auto& func : function_configurations_) { + if (!func.rva.has_value()) { + continue; + } + + auto& new_func = functions.emplace_back(); + new_func.rva = func.rva.value(); + new_func.name = func.name(); + new_func.valid = true; + } + + push(functions); + } } // namespace func_parser \ No newline at end of file diff --git a/src/lib/func_parser/parser.hpp b/src/lib/func_parser/parser.hpp index c23fd4d..07ca1e2 100644 --- a/src/lib/func_parser/parser.hpp +++ b/src/lib/func_parser/parser.hpp @@ -6,17 +6,18 @@ #include namespace func_parser { - template class Instance { public: DEFAULT_CT_CTOR_DTOR(Instance); - DEFAULT_COPY(Instance); + NON_COPYABLE(Instance); - void setup(Img* image, const config_parser::func_parser_config_t& config, const config_parser::obfuscator_config_t& obfuscator_config) { + void setup(cont::ImageBase* image, const config_parser::func_parser_config_t& config, const config_parser::obfuscator_config_t& obfuscator_config, + const std::vector& function_configurations) { image_ = image; config_ = config; obfuscator_config_ = obfuscator_config; - progress_.emplace("func_parser: discovering functions", 4); + function_configurations_ = function_configurations; + progress_.emplace("func_parser: discovering functions", 5); } void collect_functions(); @@ -34,6 +35,7 @@ namespace func_parser { void parse(); void parse_pdb(); void parse_map(); + void parse_manual(); bool push(function_list_t items) { if (items.empty()) { @@ -55,11 +57,12 @@ namespace func_parser { progress_->step(); } - Img* image_ = nullptr; + cont::ImageBase* image_ = nullptr; std::vector function_lists_; function_list_t function_list_; // function_lists_ combined and sanitized basically config_parser::func_parser_config_t config_ = {}; config_parser::obfuscator_config_t obfuscator_config_ = {}; + std::vector function_configurations_; std::optional progress_ = std::nullopt; }; } // namespace func_parser \ No newline at end of file diff --git a/src/lib/func_parser/pdb/pdb.hpp b/src/lib/func_parser/pdb/pdb.hpp index 10a2c0d..c80a574 100644 --- a/src/lib/func_parser/pdb/pdb.hpp +++ b/src/lib/func_parser/pdb/pdb.hpp @@ -1,6 +1,6 @@ #pragma once +#include "cont/base.hpp" #include "func_parser/common/common.hpp" -#include "pe/pe.hpp" #include // \todo: @es3n1n: validate that the pdb could be used for provided file @@ -8,14 +8,4 @@ namespace func_parser::pdb { function_list_t discover_functions(const std::filesystem::path& pdb_path, std::uint64_t base_of_code = 0ULL); - - inline function_list_t discover_functions(const win::cv_pdb70_t* code_view, const std::uint64_t base_of_code = 0ULL) { - // Return an empty set if there's no code view - // - if (code_view == nullptr) { - return {}; - } - - return discover_functions(code_view->pdb_name, base_of_code); - } } // namespace func_parser::pdb diff --git a/src/lib/obfuscator/config_merger/config_merger.hpp b/src/lib/obfuscator/config_merger/config_merger.hpp index 77cdb9f..f56bbc6 100644 --- a/src/lib/obfuscator/config_merger/config_merger.hpp +++ b/src/lib/obfuscator/config_merger/config_merger.hpp @@ -10,8 +10,7 @@ namespace obfuscator::config_merger { /// \param transform Transform pointer /// \param type Var type /// \return vector of names - template - std::vector get_all_required_vars_for(Transform* transform, const TransformConfig::Var::Type type) { + inline std::vector get_all_required_vars_for(Transform* transform, const TransformConfig::Var::Type type) { std::vector required_var_names = {}; transform->iter_vars([&required_var_names, type](const TransformConfig::Var& var) -> void { @@ -33,9 +32,8 @@ namespace obfuscator::config_merger { /// \param type var type /// \param values config values /// \param shared_config shared config reference - template - void apply_vars(Transform* transform, const TransformConfig::Var::Type type, const std::unordered_map& values, - TransformSharedConfig& shared_config) { + inline void apply_vars(Transform* transform, const TransformConfig::Var::Type type, const std::unordered_map& values, + TransformSharedConfig& shared_config) { /// Collect all required vars auto required_var_names = get_all_required_vars_for(transform, type); @@ -78,15 +76,14 @@ namespace obfuscator::config_merger { /// \brief Apply transform global vars /// \tparam Img X64 or X86 image /// \param config config reference - template - void apply_global_vars(config_parser::Config& config) { + inline void apply_global_vars(config_parser::Config& config) { /// Get the scheduler - auto& scheduler = TransformScheduler::get().for_arch(); + const auto& scheduler = TransformScheduler::get().container; /// Iterate over the global defined vars for the transform for (auto& [tag, values] : config.global_transforms_config()) { /// Get the transform, its shared config - auto& transform = scheduler.transforms.at(tag); + const auto& transform = scheduler.transforms.at(tag); auto& shared_config = TransformSharedConfigStorage::get().get_for(tag); /// Apply vars @@ -97,11 +94,10 @@ namespace obfuscator::config_merger { /// \brief Apply user-defined configuration for the transform /// \tparam Img X64 or X86 image /// \param transform_config user-defined options - template - void apply_config(const config_parser::transform_configuration_t& transform_config) { + inline void apply_config(const config_parser::transform_configuration_t& transform_config) { /// Get all the needed stuff - auto& scheduler = TransformScheduler::get().for_arch(); - auto& transform = scheduler.transforms.at(transform_config.tag); + const auto& scheduler = TransformScheduler::get().container; + const auto& transform = scheduler.transforms.at(transform_config.tag); auto& shared_config = TransformSharedConfigStorage::get().get_for(transform_config.tag); /// Reset all PER_FUNCTION vars diff --git a/src/lib/obfuscator/function.hpp b/src/lib/obfuscator/function.hpp index c127a40..bbde5eb 100644 --- a/src/lib/obfuscator/function.hpp +++ b/src/lib/obfuscator/function.hpp @@ -5,25 +5,24 @@ namespace obfuscator { /// \brief Function information representation /// \tparam Img PE Image type, either x64 or x86 - template struct Function { DEFAULT_DTOR(Function); NON_COPYABLE(Function); - Function(const analysis::Function& func, const Img* image) + explicit Function(const analysis::Function& func) : parsed_func(func.parsed_func), lru_reg(func.lru_reg), bb_storage(func.bb_storage), program(func.program), assembler(func.assembler), cursor(std::make_unique(program, assembler)), observer(func.observer), bb_provider(func.bb_provider), - machine_mode(image->guess_machine_mode()) { } + machine_mode(func.lru_reg.machine_mode()) { } /// \brief Construct new var allocator /// \return varalloc instance auto var_alloc() { - return analysis::VarAlloc(&lru_reg); + return analysis::VarAlloc(&lru_reg); } - /// \brief Parsed information from PDB/MAP/etc - func_parser::function_t parsed_func; + /// \brief Parsed information from PDB/MAP/etc. nullopt for nameless functions. + std::optional parsed_func; /// \brief Least recently used register cache - analysis::LRUReg lru_reg; + analysis::LRUReg lru_reg; /// \brief A storage with basic blocks std::shared_ptr bb_storage; /// \brief Zasm routine diff --git a/src/lib/obfuscator/obfuscator.cpp b/src/lib/obfuscator/obfuscator.cpp index a674e23..1b25f4d 100644 --- a/src/lib/obfuscator/obfuscator.cpp +++ b/src/lib/obfuscator/obfuscator.cpp @@ -4,6 +4,7 @@ #include "obfuscator/config_merger/config_merger.hpp" #include "obfuscator/function.hpp" #include "obfuscator/transforms/scheduler.hpp" +#include "util/sections.hpp" #include #include @@ -12,11 +13,15 @@ namespace obfuscator { constexpr size_t kTextSectionAlignment = 0x10; - template - void Instance::setup() { + void Instance::setup() { + /// Make sure that image is set + if (!image_.has_value()) { + throw std::runtime_error("obfuscator: unable to setup with image_ being nullopt"); + } + // Initializing instances // - func_parser_.setup(image_, config_.func_parser_config(), config_.obfuscator_config()); + func_parser_.setup(*image_, config_.func_parser_config(), config_.obfuscator_config(), config_.function_configurations()); // Running setup tasks // @@ -38,173 +43,207 @@ namespace obfuscator { } } - template - void Instance::add_function(const config_parser::function_configuration_t& configuration) { - /// We don't want to obfuscate functions with 0 transforms - // if (configuration.transform_configurations.empty()) { - // logger::warn("collect: excluding function {} from obfuscation list", configuration.function_name); - // return; - //} + Instance::nameless_function_t& Instance::add_function(std::span raw_function_bytes, + const config_parser::nameless_function_configuration_t& configuration) { + assert(!raw_function_bytes.empty()); /// what are you doing man + + /// Schedule transforms + schedule_transforms(configuration.transform_configurations); + + /// Analyse function and store it + return nameless_functions_.emplace_back(nameless_function_t{ + .analysed = analysis::analyse(image_mode_, raw_function_bytes), + .configuration = configuration, + }); + } + + Instance::function_t& Instance::add_function(const config_parser::function_configuration_t& configuration) { + /// Make sure image is set + if (!image_.has_value()) { + throw std::runtime_error("obfuscator: unable to add non-nameless function with image_ being nullopt"); + } /// Try to find function info from map/pdb - const auto function_info = - func_parser_.find_if([&configuration](const func_parser::function_t& func) -> bool { return func.name == configuration.function_name; }); + const auto function_info = func_parser_.find_if([&configuration](const func_parser::function_t& func) -> bool { + if (configuration.rva.has_value() && func.rva == *configuration.rva) { + return true; + } + if (configuration.function_name.has_value() && func.name == *configuration.function_name) { + return true; + } + return false; + }); if (!function_info.has_value()) { - throw std::runtime_error(std::format("collect: function {} not found", configuration.function_name)); + throw std::runtime_error(std::format("collect: function {} not found", configuration.name())); } - /// Enable needed transforms - auto& scheduler = TransformScheduler::get(); - for (const auto& [tag, _] : configuration.transform_configurations) { - scheduler.enable_transform(tag); - } + /// Schedule transforms + schedule_transforms(configuration.transform_configurations); /// Store function info - functions_.emplace_back(function_t{ - .analysed = analysis::analyse(image_, function_info.value()), + return functions_.emplace_back(function_t{ + .analysed = analysis::analyse(*image_, function_info.value()), .configuration = configuration, }); } - template - void Instance::obfuscate() { + void Instance::obfuscate() { /// Debug log - logger::info("obfuscator: got {} function(s) to obfuscate", functions_.size()); + logger::info("obfuscator: got {} function(s) to obfuscate", functions_.size() + nameless_functions_.size()); - if (functions_.empty()) { + if (functions_.empty() && nameless_functions_.empty()) { throw std::runtime_error("obfuscator: got 0 functions to protect"); } - /// Obtain transform scheduler for the platform - auto& scheduler = TransformScheduler::get().for_arch(); - config_merger::apply_global_vars(config_); + /// Apply global vars from the config + config_merger::apply_global_vars(config_); - /// Iterate over functions that we need to obfuscate + /// Iterate over the named functions and obfuscate them for (const auto& func : functions_) { - /// Init the `obfuscator::Function` that is going to be used within - /// transforms - auto obf_func = obfuscator::Function(func.analysed, image_); - - /// Export tags that this function would need - auto tags = std::views::all(func.configuration.transform_configurations) | - std::views::transform([](const config_parser::transform_configuration_t& it) -> TransformTag { return it.tag; }) | - std::ranges::to(); - - /// Export transforms - auto transforms = scheduler.select_transforms(tags); - - /// Init the progress bar - auto progress = progress::Progress(std::format("obfuscator: obfuscating {}", obf_func.parsed_func.name), transforms.size()); - - /// An util that would check the chances and all this other crap, that would be - /// needed for like every possible function/transform - auto execute_transform = [func](const TransformTag tag, const std::function& callback, - const bool check_chances = true) -> void { - auto preset = std::ranges::find_if(func.configuration.transform_configurations, [tag](auto&& it) -> bool { - return it.tag == tag; // - }); - if (preset == std::end(func.configuration.transform_configurations)) { - throw std::runtime_error(std::format("obfuscate: unable to find configuration for transform {}", tag)); - } + auto obf_func = Function(func.analysed); + obfuscate(func.configuration.transform_configurations, obf_func, func.configuration.name()); + } - /// Apply the preset - config_merger::apply_config(*preset); + /// Iterate over the nameless functions and obfuscate them + for (const auto& func : nameless_functions_) { + auto obf_func = Function(func.analysed); + obfuscate(func.configuration.transform_configurations, obf_func, "nameless"); + } + } - /// Get the shared config and check the chance - auto& cfg = TransformSharedConfigStorage::get().get_for(tag); + void Instance::obfuscate(const config_parser::transform_configurations_t& configurations, Function& function, const std::string& function_name) const { + auto& scheduler = TransformScheduler::get().container; + + /// Export tags that this function would need + auto tags = std::views::all(configurations) | + std::views::transform([](const config_parser::transform_configuration_t& it) -> TransformTag { return it.tag; }) | + std::ranges::to(); + + /// Export transforms + auto transforms = scheduler.select_transforms(tags); + + /// Init the progress bar + auto progress = progress::Progress(std::format("obfuscator: obfuscating {}", function_name), transforms.size()); + + /// An util that would check the chances and all this other crap, that would be + /// needed for like every possible function/transform + auto execute_transform = [configurations](const TransformTag tag, const std::function& callback, + const bool check_chances = true) -> void { + const auto preset = std::ranges::find_if(configurations, [tag](auto&& it) -> bool { + return it.tag == tag; // + }); + if (preset == std::end(configurations)) { + throw std::runtime_error(std::format("obfuscate: unable to find configuration for transform {}", tag)); + } - /// Check the chance - /// \todo @es3n1n: Check for chance feature - if (check_chances && !rnd::chance(cfg.chance())) { - return; - } + /// Apply the preset + config_merger::apply_config(*preset); - /// Otherwise run this method - for (std::size_t i = 0; i < cfg.repeat_times(); ++i) { - /// Init context, run the task - auto context = TransformContext(cfg); + /// Get the shared config and check the chance + auto& cfg = TransformSharedConfigStorage::get().get_for(tag); - do { - context.rerun_me = false; - callback(context); - } while (context.rerun_me); - } - }; - auto execute_transform_no_chances = [&](const TransformTag tag, const std::function& callback) -> void { - return execute_transform(tag, callback, false); - }; - - /// \note @es3n1n: We can't iterate through the insns/bbs and execute transforms - /// from there as it would break the scheduling order - for (auto& [tag, transform] : transforms) { - /// Apply function transform - if (transform->feature(TransformFeaturesSet::HAS_FUNCTION_TRANSFORM)) { - execute_transform_no_chances(tag, [&obf_func, &transform](auto& ctx) -> void { - transform->run_on_function(ctx, &obf_func); // - }); - } + /// Check the chance + /// \todo @es3n1n: Check for chance feature + if (check_chances && !rnd::chance(cfg.chance())) { + return; + } - /// Apply basic block transforms - if (transform->feature(TransformFeaturesSet::HAS_BB_TRANSFORM)) { - for (auto& basic_block : obf_func.bb_storage->temp_copy()) { - execute_transform(tag, [&obf_func, &transform, &basic_block](auto& ctx) -> void { - transform->run_on_bb(ctx, &obf_func, basic_block.get()); // - }); - } - } + /// Otherwise run this method + for (std::size_t i = 0; i < cfg.repeat_times(); ++i) { + /// Init context, run the task + auto context = TransformContext(cfg); - /// Apply analysis insn transforms - if (transform->feature(TransformFeaturesSet::HAS_INSN_TRANSFORM)) { - for (auto& basic_block : obf_func.bb_storage->temp_copy()) { - for (auto& insn : basic_block->temp_insns_copy()) { - execute_transform(tag, [&obf_func, &transform, &insn](auto& ctx) -> void { - transform->run_on_insn(ctx, &obf_func, insn.get()); // - }); - } - } + do { + context.rerun_me = false; + callback(context); + } while (context.rerun_me); + } + }; + auto execute_transform_no_chances = [&](const TransformTag tag, const std::function& callback) -> void { + execute_transform(tag, callback, false); + }; + + /// \note @es3n1n: We can't iterate through the insns/bbs and execute transforms + /// from there as it would break the scheduling order + for (auto& [tag, transform] : transforms) { + /// Apply function transform + if (transform->feature(TransformFeaturesSet::HAS_FUNCTION_TRANSFORM)) { + execute_transform_no_chances(tag, [&function, &transform](auto& ctx) -> void { + transform->run_on_function(ctx, &function); // + }); + } + + /// Apply basic block transforms + if (transform->feature(TransformFeaturesSet::HAS_BB_TRANSFORM)) { + for (auto& basic_block : function.bb_storage->temp_copy()) { + execute_transform(tag, [&function, &transform, &basic_block](auto& ctx) -> void { + transform->run_on_bb(ctx, &function, basic_block.get()); // + }); } + } - /// Apply program nodes transform - if (transform->feature(TransformFeaturesSet::HAS_NODE_TRANSFORM)) { - for (auto* node = obf_func.program->getHead(); node != nullptr; node = node->getNext()) { - /// Transform nodes - execute_transform(tag, [&obf_func, &transform, &node](auto& ctx) -> void { - transform->run_on_node(ctx, &obf_func, node); // + /// Apply analysis insn transforms + if (transform->feature(TransformFeaturesSet::HAS_INSN_TRANSFORM)) { + for (const auto& basic_block : function.bb_storage->temp_copy()) { + for (auto& insn : basic_block->temp_insns_copy()) { + execute_transform(tag, [&function, &transform, &insn](auto& ctx) -> void { + transform->run_on_insn(ctx, &function, insn.get()); // }); } } + } - /// Increment progress bar - progress.step(); + /// Apply program nodes transform + if (transform->feature(TransformFeaturesSet::HAS_NODE_TRANSFORM)) { + for (auto* node = function.program->getHead(); node != nullptr; node = node->getNext()) { + /// Transform nodes + execute_transform(tag, [&function, &transform, &node](auto& ctx) -> void { + transform->run_on_node(ctx, &function, node); // + }); + } } - /// We are done here + /// Increment progress bar + progress.step(); } + + /// We are done here } - template - void Instance::assemble() { + void Instance::assemble() { + /// Make sure that we have an image to deal with + if (!image_.has_value()) { + throw std::runtime_error("obfuscate: no image available for assembling"); + } + + /// Notify the user that they would not see nameless function in their binary + if (!nameless_functions_.empty()) { + logger::warn("please note that {} 'nameless' functions would not be assembled in the output PE", nameless_functions_.size()); + } + /// Estimating section size auto size_estimation_progress = progress::Progress("obfuscator: estimating section size", functions_.size()); - std::size_t section_size = 0; + auto sec_header = sections::get(sections::e_section_t::CODE); for (auto& func : functions_) { + easm::dump_program(*func.analysed.program); const auto program_size = easm::estimate_program_size(*func.analysed.program); - section_size += memory::address{program_size}.align_up(kTextSectionAlignment).as(); + sec_header.size_raw_data += memory::address{program_size}.align_up(kTextSectionAlignment).as(); size_estimation_progress.step(); } - logger::debug("assemble: estimated new section size: {:#x}", section_size); + logger::debug("assemble: estimated new section size: {:#x}", sec_header.size_raw_data); /// Allocate new section - auto img_base = image_->raw_image->get_nt_headers()->optional_header.image_base; - auto& new_sec = image_->new_section(sections::e_section_t::CODE, section_size); + auto img_base = (*image_)->get_image_base(); + auto& new_sec = (*image_)->new_section(sec_header); memory::address virt_address = new_sec.virtual_address; /// Iterate over the obfuscated functions auto linking_progress = progress::Progress("obfuscator: linking functions", functions_.size()); - for (auto& [func, _] : functions_) { - /// \todo @es3n1n: perhaps i should split this monstrosity into a separate functions - /// Erase the original function code + /// Iterating over the enabled functions + /// \todo @es3n1n: perhaps i should split this monstrosity into a separate functions + for (auto& [func, _] : functions_) { + /// Erase the original functions for (auto& basic_block : *func.bb_storage) { for (auto& insn : basic_block) { /// Clang-tidy is working a bit weird with smart pointers and `bugprone-unchecked-optional-access` @@ -219,26 +258,26 @@ namespace obfuscator { const auto randomized = rnd::bytes(*raw_ptr->length); /// Replace instruction with junk - auto* insn_ptr = image_->rva_to_ptr(*raw_ptr->rva); + auto* insn_ptr = (*image_)->rva_to_ptr(*raw_ptr->rva); std::memcpy(insn_ptr, randomized.data(), randomized.size()); /// Remove pe relocation, if there's any if (raw_ptr->reloc.type == analysis::insn_reloc_t::e_type::HEADER) { - image_->relocations.erase(*raw_ptr->rva + raw_ptr->reloc.offset.value_or(0)); + (*image_)->relocations.erase(*raw_ptr->rva + raw_ptr->reloc.offset.value_or(0)); } } } /// Insert the jmp to obfuscated routine at the very beginning of the function - auto* func_start_ptr = image_->rva_to_ptr(func.range.start); - auto jmp_data = easm::encode_jmp(image_->guess_machine_mode(), func.range.start + img_base, virt_address + img_base); + auto* func_start_ptr = (*image_)->rva_to_ptr(func.range.start); + auto jmp_data = easm::encode_jmp(machine_mode(), func.range.start + img_base, virt_address + img_base); if (!jmp_data.has_value()) { throw std::runtime_error("assemble: unable to encode jmp"); } std::memcpy(func_start_ptr, jmp_data->data(), jmp_data->size()); /// Assemble the obfuscated function - auto assemble_progress = progress::Progress(std::format("obfuscator: assembling {}", func.parsed_func.name), 1); + auto assemble_progress = progress::Progress(std::format("obfuscator: assembling {}", func.parsed_func.value().name), 1); const auto assembled = easm::assemble_program(virt_address + img_base, *func.program); assemble_progress.step(); @@ -266,14 +305,14 @@ namespace obfuscator { } /// Store the new relocation data - image_->relocations[relocation.address - img_base] = - pe::relocation_t{.rva = memory::address{static_cast(relocation.address - img_base)}, + (*image_)->relocations[relocation.address - img_base] = + cont::Relocation{.rva = memory::address{static_cast(relocation.address - img_base)}, .size = static_cast(getBitSize(relocation.size) / CHAR_BIT), - .type = win_reloc_type}; + .type = to_cont(win_reloc_type)}; } /// Align size and increment offset - const auto aligned_size = memory::address{assembled.data.size()}.align_up(kTextSectionAlignment).as(); + const auto aligned_size = memory::address{assembled.data.size()}.align_up(kTextSectionAlignment).as(); virt_address = virt_address.offset(aligned_size); /// Increment progress bar @@ -283,10 +322,14 @@ namespace obfuscator { logger::info("assemble: assembled {} functions", functions_.size()); } - template - std::filesystem::path Instance::save() { + std::filesystem::path Instance::save() { + /// We can't rebuild PE without any source PE data :shrug: + if (!image_.has_value()) { + throw std::runtime_error("assemble: no image for saving"); + } + logger::info("obfuscator: saving.."); - auto new_img = image_->rebuild_pe_image(); + auto new_img = (*image_)->rebuild_image(); auto out_path = config_.obfuscator_config().binary_path; @@ -303,13 +346,18 @@ namespace obfuscator { return out_path; } - template - std::filesystem::path Instance::run() { + std::filesystem::path Instance::run() { setup(); obfuscate(); assemble(); return save(); } - PE_DECL_TEMPLATE_CLASSES(Instance); + void Instance::schedule_transforms(const config_parser::transform_configurations_t& configurations) { + /// Enable needed transforms + auto& scheduler = TransformScheduler::get(); + for (const auto& [tag, _] : configurations) { + scheduler.enable_transform(tag); + } + } } // namespace obfuscator diff --git a/src/lib/obfuscator/obfuscator.hpp b/src/lib/obfuscator/obfuscator.hpp index 075ae71..55353b0 100644 --- a/src/lib/obfuscator/obfuscator.hpp +++ b/src/lib/obfuscator/obfuscator.hpp @@ -2,35 +2,67 @@ #include "analysis/analysis.hpp" #include "config_parser/config_parser.hpp" #include "func_parser/parser.hpp" -#include "pe/pe.hpp" +#include "obfuscator/function.hpp" #include "util/structs.hpp" namespace obfuscator { - template class Instance { public: - Instance(Img* image, config_parser::Config& config): image_(image), config_(std::move(config)) { } + Instance(cont::ImageBase* image, config_parser::Config& config): image_(image), image_mode_(image->mode()), config_(std::move(config)) { } + explicit Instance(const cont::ImageMode image_mode): image_(std::nullopt), image_mode_(image_mode) { } + DEFAULT_DTOR(Instance); NON_COPYABLE(Instance); + struct function_t { + analysis::Function analysed; + config_parser::function_configuration_t configuration; + }; + + struct nameless_function_t { + analysis::Function analysed; + config_parser::nameless_function_configuration_t configuration; + }; + + /// Parses functions using `func_parser_`, enables transforms in the scheduler void setup(); - void add_function(const config_parser::function_configuration_t& configuration); + + /// Adds a function to obfuscate from its raw binary representation (mostly used by tests) + nameless_function_t& add_function(std::span raw_function_bytes, + const config_parser::nameless_function_configuration_t& configuration); + /// Adds a function to obfuscate by name (will be looked up in `func_parser_`) + function_t& add_function(const config_parser::function_configuration_t& configuration); + + /// Run over the added functions and apply enabled transforms void obfuscate(); + + /// 1. Allocates new section for our code + /// 2. Erases original function's code/relocations + /// 3. Links obfuscated functions + /// 4. Writes the obfuscated functions to the new section + /// 4. Saves new relocations void assemble(); + + /// Rebuilds PE with new data we got and saves it to disk std::filesystem::path save(); - // setup() -> obfuscate() -> assemble() -> save() + /// setup() -> obfuscate() -> assemble() -> save() std::filesystem::path run(); - struct function_t { - analysis::Function analysed; - config_parser::function_configuration_t configuration; - }; - private: - Img* image_ = nullptr; + static void schedule_transforms(const config_parser::transform_configurations_t& configurations); + void obfuscate(const config_parser::transform_configurations_t& configurations, Function& function, const std::string& function_name) const; + + /// \fixme @es3n1n: this is wrong + [[nodiscard]] zasm::MachineMode machine_mode() const { + return image_mode_ == cont::ImageMode::X64 ? zasm::MachineMode::AMD64 : zasm::MachineMode::I386; + } + + std::optional image_ = nullptr; + cont::ImageMode image_mode_; config_parser::Config config_; - func_parser::Instance func_parser_; + func_parser::Instance func_parser_; std::vector functions_; + std::vector nameless_functions_; }; } // namespace obfuscator \ No newline at end of file diff --git a/src/lib/obfuscator/transforms/configs.hpp b/src/lib/obfuscator/transforms/configs.hpp index 07c1563..5b4f77e 100644 --- a/src/lib/obfuscator/transforms/configs.hpp +++ b/src/lib/obfuscator/transforms/configs.hpp @@ -140,7 +140,7 @@ namespace obfuscator { } /// \brief Get transform config using the transform type - template