aboutsummaryrefslogtreecommitdiff
path: root/memory.hpp
blob: ff40d547057bcad9138e8789443223fee4d1f91d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
#ifndef MEMORY_H
#define MEMORY_H 1

#include <EASTL/string.h>
#include <EASTL/vector.h>
#include <cstdint>
#include <cstdlib>
#include <ntifs.h>
#include <wdm.h>

#include "stringify.hpp"

struct Process {
  uint32_t NumberOfThreads;
  eastl::wstring ProcessName;
  uint64_t UniqueProcessId;
  uint32_t HandleCount;
};

struct Module {
  uint64_t DllBase;
  uint64_t EntryPoint;
  uint32_t SizeOfImage;
  eastl::wstring FullDllName;
  eastl::wstring BaseDllName;
  uint32_t Flags;
  uint16_t LoadCount;
  uint16_t TlsIndex;
};

struct Page {
  eastl::string toString() const {
    return ::toString(BaseAddress, RegionSize, Type, State, Protect);
  }

  uint64_t BaseAddress;
  uint64_t AllocationBase;
  uint32_t AllocationProtect;
  size_t RegionSize;
  uint32_t State;
  uint32_t Protect;
  uint32_t Type;
};

void SetLdrInitWaitPrefs(int waitCount, LONGLONG waitTimeout);
eastl::vector<Process> GetProcesses();
NTSTATUS OpenProcess(_In_ HANDLE pid, _Out_ PEPROCESS *pep, _Out_ HANDLE *process);
NTSTATUS CloseProcess(_In_ _Out_ PEPROCESS *pep, _In_ _Out_ HANDLE *process);
eastl::vector<Page> GetPages(_In_ HANDLE obj, SIZE_T maxPages = 1024,
                             ULONG_PTR startAddress = 0);
eastl::vector<Module> GetModules(_In_ PEPROCESS Process, _In_ BOOLEAN isWow64);
NTSTATUS ProtectVirtualMemory(_In_ PEPROCESS pep,
                              _In_ uint64_t addr,
                              _In_ SIZE_T size, _In_ ULONG newProt,
                              _Out_ ULONG *oldProt);
NTSTATUS RestoreProtectVirtualMemory(_In_ PEPROCESS pep, _In_ uint64_t addr,
                                     _In_ SIZE_T siz, _In_ ULONG old_prot);
NTSTATUS ReadVirtualMemory(_In_ PEPROCESS pep, _In_ uint64_t sourceAddress,
                           _Out_ UCHAR *targetAddress, _In_ _Out_ SIZE_T *size);
NTSTATUS WriteVirtualMemory(_In_ PEPROCESS pep, _In_ const UCHAR *sourceAddress,
                            _Out_ uint64_t targetAddress,
                            _In_ _Out_ SIZE_T *size);

auto get_process_cr3(PEPROCESS pe_process) -> uint64_t;
auto swap_process(PEPROCESS new_process) -> PEPROCESS;

class Memory {
public:
    Memory(_In_ PEPROCESS& pep) : m_pep(pep) {
        ClearLastErrorAndSize();
    }
    Memory(const Memory&) = delete;

    void ClearLastErrorAndSize() {
        m_last_error = STATUS_SUCCESS;
        m_last_size = 0;
    }

    static bool IsValidAddress(uint64_t address) {
        return address >= 0x10000 && address < 0x000F000000000000;
    }

    template<typename T>
    T Read(uint64_t sourceAddress) {
        if (!IsValidAddress(sourceAddress))
            return T();
        T value;
        SIZE_T size = sizeof(value);
        m_last_error = ReadVirtualMemory(m_pep, sourceAddress, reinterpret_cast<UCHAR*>(&value), &size);
        m_last_size = size;
        if (m_last_error == STATUS_SUCCESS && m_last_size == sizeof(T))
            return value;
        return T();
    }

    template<typename T>
    T ReadChain(uint64_t sourceAddress, const eastl::vector<uint64_t>& chainedOffsets) {
        for (const auto& offset : chainedOffsets) {
            if (offset == chainedOffsets.back())
                break;
            sourceAddress = Read<uint64_t>(sourceAddress + offset);
            if (!sourceAddress)
                break;
        }
        if (chainedOffsets.size() == 0)
            return Read<T>(sourceAddress);
        else
            return Read<T>(sourceAddress + chainedOffsets.back());
    }

    template<typename T, size_t N>
    bool ReadBuffer(uint64_t sourceAddress, T out[N]) {
        if (!IsValidAddress(sourceAddress))
            return false;
        SIZE_T size = sizeof(T) * N;
        m_last_error = ReadVirtualMemory(m_pep, sourceAddress, reinterpret_cast<UCHAR*>(out), &size);
        m_last_size = size;
        return Succeeded<T, N>();
    }

    template<size_t N>
    bool ReadString(uint64_t sourceAddress, char out[N]) {
        const auto retval = ReadBuffer<char, N-1>(sourceAddress, out);
        out[m_last_size] = '\0';
        return retval;
    }

    template<typename T>
    bool Write(uint64_t targetAddress, const T& writeData) {
        if (!IsValidAddress(targetAddress))
            return false;
        SIZE_T size = sizeof(T);
        m_last_error = WriteVirtualMemory(m_pep, reinterpret_cast<const UCHAR*>(&writeData), targetAddress, &size);
        m_last_size = size;
        return Succeeded<T, 1>();
    }

    template<typename T, size_t N>
    bool Succeeded() { return m_last_error == STATUS_SUCCESS && m_last_size == sizeof(T) * N; }

    NTSTATUS LastError() { return m_last_error; }
    SIZE_T LastSize() { return m_last_size; }

private:
    PEPROCESS& m_pep;
    NTSTATUS m_last_error;
    SIZE_T m_last_size;
};
#endif