Fixing tests and Architecture lookup (#27)

* minor fixes

* added pwn::shellcode::patternfind

* Fixed minor bugs when looking up architectures

Modules/Assembly/Source/Disassembler.cpp

@@ -301,22 +301,22 @@ Disassembler::Print(std::vector<u8> const& bytes, std::optional<Architecture> ar
 void
 Disassembler::X64(std::vector<u8> const& bytes)
 {
- Disassembler::Print(bytes, Architectures["x64"]);
+ Disassembler::Print(bytes, Architectures[ArchitectureType::x64]);
 }
 
 
 void
 Disassembler::X86(std::vector<u8> const& bytes)
 {
- Disassembler::Print(bytes, Architectures["x86"]);
+ Disassembler::Print(bytes, Architectures[ArchitectureType::x86]);
 }
 #endif // PWN_DISASSEMBLE_X86
 
 #ifdef PWN_DISASSEMBLE_ARM64
 void
 Disassembler::ARM64(std::vector<u8> const& bytes)
 {
- Disassembler::Print(bytes, Architectures["arm64"]);
+ Disassembler::Print(bytes, Architectures[ArchitectureType::arm64]);
 }
 #endif
 

Modules/Assembly/Tests/CMakeLists.txt

@@ -1,12 +1,12 @@
 enable_testing()
 set(TEST_EXECUTABLE_NAME tests_pwn_${PROJECT_NAME})
-list(APPEND SOURCE_FILES
+set(FILES
 
  ${TEST_DIR}/main.cpp
  ${TEST_DIR}/pwn_disasm.cpp
 )
 
-add_executable(${TEST_EXECUTABLE_NAME} ${SOURCE_FILES})
+add_executable(${TEST_EXECUTABLE_NAME} ${FILES})
 add_executable(PWN::Tests::${PROJECT_NAME} ALIAS ${TEST_EXECUTABLE_NAME})
 add_dependencies(${TEST_EXECUTABLE_NAME} PWN::Deps::Catch2 PWN::${PROJECT_NAME})
 target_link_libraries(${TEST_EXECUTABLE_NAME} PUBLIC Catch2::Catch2WithMain PWN::${PROJECT_NAME})

Modules/Binary/Include/Win32/PE.hpp

@@ -213,6 +213,14 @@ class PE
  return m_PeHeader;
  }
 
+ uptr const
+ EntryPointAddress() const
+ {
+ auto const hdrs = Header();
+ return Is64b() ? std::get<PeHeader64>(hdrs).OptionalHeader.AddressOfEntryPoint :
+ std::get<PeHeader32>(hdrs).OptionalHeader.AddressOfEntryPoint;
+ }
+
 
  std::vector<PeDataDirectory> const&
  DataDirectories() const

Modules/Binary/Tests/CMakeLists.txt

@@ -1,15 +1,10 @@
 enable_testing()
 set(TEST_BINARY_PE tests_pwn_${PROJECT_NAME})
-list(APPEND SOURCE_FILES
-
- ${TEST_DIR}/main.cpp
- ${TEST_DIR}/pwn_binary_pe.cpp
-)
 
 set(DEPS PWN::Common PWN::FileSystem)
 file(GLOB DLL_TEST_FILES "C:/Windows/System32/*.dll")
 
-add_executable(${TEST_BINARY_PE} ${SOURCE_FILES})
+add_executable(${TEST_BINARY_PE} ${TEST_DIR}/main.cpp ${TEST_DIR}/pwn_binary_pe.cpp)
 add_executable(PWN::Tests::${PROJECT_NAME}::PE ALIAS ${TEST_BINARY_PE})
 add_dependencies(${TEST_BINARY_PE} PWN::Deps::Catch2 PWN::${PROJECT_NAME} ${DEPS})
 target_link_libraries(${TEST_BINARY_PE} PUBLIC Catch2::Catch2WithMain PWN::${PROJECT_NAME} ${DEPS})

Modules/Common/Include/Architecture.hpp

@@ -87,25 +87,27 @@ struct Architecture
 
  ///
  ///@brief Find an architecture by name. The function will throw `std::range_error`
- // if not found
+ // if not found.
+ ///
+ ///@param architecture_name
  ///
- ///@param sv
  ///@return Architecture
+ ///@throw range_error if not found
  ///
- static Architecture
- Find(std::string_view const& sv);
+ static Architecture const&
+ Find(std::string_view const& architecture_name);
 };
 
 
 ///
 ///@brief Supported architecture declarations
 ///
-static constexpr CMap<std::string_view, Architecture, 4> Architectures {
+static constexpr CMap<ArchitectureType, Architecture, 4> Architectures {
  {{
- {"x64"sv, {"X86_64"sv, ArchitectureType::x64, 8, Endianess::little}},
- {"x86"sv, {"X86_32"sv, ArchitectureType::x86, 4, Endianess::little}},
- {"arm64"sv, {"ARM_AARCH64"sv, ArchitectureType::arm64, 8, Endianess::little}},
- {"arm"sv, {"ARM_AARCH64"sv, ArchitectureType::arm, 4, Endianess::little}},
+ {ArchitectureType::x64, {"X86_64"sv, ArchitectureType::x64, 8, Endianess::little}},
+ {ArchitectureType::x86, {"X86_32"sv, ArchitectureType::x86, 4, Endianess::little}},
+ {ArchitectureType::arm64, {"ARM_AARCH64"sv, ArchitectureType::arm64, 8, Endianess::little}},
+ {ArchitectureType::arm, {"ARM_AARCH64"sv, ArchitectureType::arm, 4, Endianess::little}},
  }},
 };
 

Modules/Common/Include/Common.hpp

@@ -33,10 +33,6 @@
 // clang-format on
 #pragma warning(pop)
 
-#ifndef UNREFERENCED_PARAMETER
-#define UNREFERENCED_PARAMETER(x) ((void)x)
-#endif // UNREFERENCED_PARAMETER
-
 #elif defined(PWN_BUILD_FOR_LINUX)
 
 //
@@ -62,11 +58,12 @@
 #define MAX_PATH 260
 #endif // MAX_PATH
 
+#endif // defined(PWN_BUILD_FOR_LINUX)
+
 #ifndef UNREFERENCED_PARAMETER
 #define UNREFERENCED_PARAMETER(x) ((void)x)
 #endif // UNREFERENCED_PARAMETER
 
-#endif // defined(PWN_BUILD_FOR_LINUX)
 
 #ifndef UnusedParameter
 #define UnusedParameter UNREFERENCED_PARAMETER
@@ -114,9 +111,6 @@ using usize = std::size_t;
 using ssize = std::intptr_t;
 using uptr = std::uintptr_t;
 
-#ifndef UnreferencedParameter
-#define UnreferencedParameter(x) ((void)(x))
-#endif // UnreferencedParameter
 
 using namespace std::literals::string_view_literals;
 using namespace std::literals::chrono_literals;
@@ -150,7 +144,8 @@ constexpr_concat(std::string const& arg, Args... args)
 template<typename Key, typename Value, usize Size>
 struct CMap
 {
- std::array<std::pair<Key, Value>, Size> data;
+ using CMapEntry = std::pair<Key, Value>;
+ std::array<CMapEntry, Size> data;
 
  [[nodiscard]] constexpr Value
  at(const Key& key) const
@@ -166,10 +161,7 @@ struct CMap
  {
  return itr->second;
  }
- else
- {
- throw std::range_error("Not Found");
- }
+ throw std::range_error("Not Found");
  }
 
  [[nodiscard]] constexpr Value

Modules/Common/Include/Context.hpp

@@ -41,6 +41,12 @@ class GlobalContext
  return;
  }
 
+ if constexpr ( std::is_same_v<T, ArchitectureType> )
+ {
+ SetArchitecture(arg);
+ return;
+ }
+
  if constexpr ( std::is_same_v<T, Log::LogLevel> )
  {
  SetLogLevel(arg);
@@ -65,7 +71,16 @@ class GlobalContext
 
 private:
  ///
- ///@brief Set the Architecture object
+ /// @brief Set the Architecture object
+ ///
+ /// @param arch
+ ///
+ void
+ SetArchitecture(ArchitectureType const& arch);
+
+
+ ///
+ ///@brief Set the Architecture object from a string
  ///
  ///@param type
  ///

Modules/Common/Source/Architecture.cpp

@@ -36,10 +36,21 @@ operator<<(std::wostream& wos, Endianess e)
 }
 
 
-Architecture
+const Architecture&
 Architecture::Find(std::string_view const& sv)
 {
- return Architectures.at(sv);
+ auto entry = std::find_if(
+ Architectures.data.cbegin(),
+ Architectures.data.cend(),
+ [&](auto const& e)
+ {
+ return e.second.name == sv;
+ });
+ if ( entry == Architectures.data.cend() )
+ {
+ throw std::range_error("Architecture not found");
+ }
+ return entry->second;
 }
 
 

Modules/Common/Source/Context.cpp

@@ -13,14 +13,20 @@ GlobalContext::GlobalContext()
  Set("x64");
 };
 
+void
+GlobalContext::SetArchitecture(ArchitectureType const& archtype)
+{
+ architecture = Architectures.at(archtype);
+}
 
 void
 GlobalContext::SetArchitecture(std::string_view const& type)
 {
- architecture = Architecture::Find(type);
- endianess = architecture.endian;
- ptrsize = architecture.ptrsize;
- dbg("Selecting '{}'", architecture);
+ auto arch = Architecture::Find(type);
+ architecture = arch;
+ endianess = arch.endian;
+ ptrsize = arch.ptrsize;
+ dbg("Selecting '{}'", arch);
 }
 
 

Modules/Common/Source/Utils.cpp

@@ -209,7 +209,7 @@ Random::Next() -> u64
 
 
 auto
-Random::Next(u64 const max, u64 const min) noexcept -> u64
+Random::Next(u64 const min, u64 const max) noexcept -> u64
 {
  return (XorShift64() + min) % max;
 }
@@ -521,7 +521,7 @@ Pack::p8(u8 v, Endianess e)
 void
 Pause()
 {
- dbg("Pausing, press enter to resume...");
+ info("Pausing, press enter to resume...");
  std::cin.get();
 }
 

Modules/Process/Source/Win32/Process.cpp

@@ -94,9 +94,6 @@ Process::Process(u32 Pid) : m_ProcessId {Pid}
  throw std::runtime_error("Process initialization error");
  }
 
- xdbg("Process handle with {}", ProcessAccessToString(m_ProcessHandleAccessMask).c_str());
-
-
  // Process PPID
  {
  auto BasicInfo = Value(Query<PROCESS_BASIC_INFORMATION>(PROCESSINFOCLASS::ProcessBasicInformation));
@@ -107,15 +104,16 @@ Process::Process(u32 Pid) : m_ProcessId {Pid}
  // Full path
  {
  auto NativeFilePath = Value(Query<UNICODE_STRING>(PROCESSINFOCLASS::ProcessImageFileName));
- m_NativePath = std::wstring {NativeFilePath->Buffer};
+ m_NativePath = (NativeFilePath->Length) ? std::wstring {NativeFilePath->Buffer} : std::wstring {L""};
  }
 }
 
 
 Process::Process(HANDLE&& hProcess) : Process(::GetProcessId(hProcess))
 {
  m_ProcessHandle.reset(std::move(hProcess));
- m_ProcessHandleAccessMask = PROCESS_ALL_ACCESS;
+ // TODO: fix by querying existing access
+ m_ProcessHandleAccessMask = PROCESS_QUERY_LIMITED_INFORMATION;
 }
 
 
@@ -525,6 +523,8 @@ Process::ReOpenProcessWith(const DWORD DesiredAccess)
  //
  m_ProcessHandle = UniqueHandle {hProcess};
  m_ProcessHandleAccessMask = NewAccessMask;
+
+ xdbg("Process handle with {}", ProcessAccessToString(m_ProcessHandleAccessMask).c_str());
  return Ok(true);
 }
 

Modules/Shellcode/CMakeLists.txt

@@ -7,8 +7,12 @@ set(SOURCE_DIR ${PROJECT_DIR}/Source)
 set(HEADER_DIR ${SOURCE_DIR}/Include)
 set(TEST_DIR ${PROJECT_DIR}/Tests)
 
+set(SOURCE_FILES
+ ${SOURCE_DIR}/Pattern.cpp
+)
+
 if(WIN32)
- set(SOURCE_FILES
+ list(APPEND SOURCE_FILES
  ${SOURCE_DIR}/Win32/Kernel.cpp
 
  $<$<STREQUAL:${CMAKE_GENERATOR_PLATFORM},x64>:${SOURCE_DIR}/Win32/asm/x64/copy_system_token.asm>

Modules/Shellcode/Include/Pattern.hpp

@@ -0,0 +1,32 @@
+#pragma once
+
+#include "Common.hpp"
+
+
+namespace pwn::Shellcode
+{
+
+///
+///@brief
+///
+///@param Bytes
+///@param Needle
+///@param NeedleMask
+///
+///@return ssize
+///
+ssize
+PatternFind(std::vector<u8> const& Bytes, std::vector<u8> const& Needle, std::vector<u8> const& NeedleMask);
+
+///
+///@brief
+///
+///@param Bytes
+///@param Needle
+///
+///@return ssize
+///
+ssize
+PatternFind(std::vector<u8> const& Bytes, std::vector<u8> const& Needle);
+
+} // namespace pwn::Shellcode