4 files changed, 152 insertions, 70 deletions

@@ -8,41 +8,41 @@
 #include <thread>
 
 #include "Memory.h"
+#include <Random.h>
 class Randomizer {
 public:
     Randomizer(const std::shared_ptr<Memory>&) {}
-    void Randomize() {
+    void Randomize(int seed) {
         std::this_thread::sleep_for(std::chrono::milliseconds(1000));
     }
+    
+    void RandomizeChallenge(int seed) {
+        Randomize(seed);
+    }
 };
 
-#define EXE_NAME L"witness64_d3d11.exe"
-#define HEARTBEAT 0x401
+// Heartbeat is defined to 0x401 by Memory.h
 #define RANDOMIZE_READY 0x402
 #define RANDOMIZING 0403
 #define RANDOMIZE_DONE 0x404
+#define RANDOMIZE_CHALLENGE_DONE 0x405
+#define CHALLENGE_ONLY 0x406
 
 // Globals
 HWND g_hwnd;
 HWND g_seed;
 HWND g_randomizerStatus;
 HINSTANCE g_hInstance;
-auto g_witnessProc = std::make_shared<Memory>();
+auto g_witnessProc = std::make_shared<Memory>(L"witness64_d3d11.exe");
 std::shared_ptr<Randomizer> g_randomizer;
 
-// Notifications I need:
-// Game shutdown
-// Load game?
-
 LRESULT CALLBACK WndProc(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam) {
         if (message == WM_DESTROY) {
                 PostQuitMessage(0);
     } else if (message == WM_COMMAND || message == WM_TIMER) {
         switch (LOWORD(wParam)) {
             case HEARTBEAT:
-                SetTimer(g_hwnd, HEARTBEAT, 1000, NULL);
-                // Potential improvement: Change this call to be part of the HEARTBEAT message.
-                switch (g_witnessProc->Heartbeat(EXE_NAME)) {
+                switch ((ProcStatus)lParam) {
                     case ProcStatus::NotRunning:
                         // Shut down randomizer, wait for startup
                         if (g_randomizer) g_randomizer = nullptr;
@@ -60,21 +60,58 @@ LRESULT CALLBACK WndProc(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam)
                     }
                 break;
             case RANDOMIZE_READY:
-                SetWindowText(g_randomizerStatus, L"Randomize");
                 EnableWindow(g_randomizerStatus, TRUE);
+                if (IsDlgButtonChecked(hwnd, CHALLENGE_ONLY)) {
+                    SetWindowText(g_randomizerStatus, L"Randomize Challenge");
+                } else {
+                    SetWindowText(g_randomizerStatus, L"Randomize");
+                }
                 break;
             case RANDOMIZING:
-                if (!g_randomizer) break;
+                if (!g_randomizer) {
+                    assert(false);
+                    break;
+                }
                 EnableWindow(g_randomizerStatus, FALSE);
-                SetWindowText(g_randomizerStatus, L"Randomizing...");
-                std::thread([]{
-                    g_randomizer->Randomize();
-                    PostMessage(g_hwnd, WM_COMMAND, RANDOMIZE_DONE, NULL);
-                }).detach();
+
+                {
+                    int seed = 0;
+                    std::wstring text(128, L'\0');
+                    int size = GetWindowText(g_seed, text.data(), static_cast<int>(text.size()));
+                    if (size > 0) { // Set seed
+                        seed = _wtoi(text.c_str());
+                    } else { // Random seed
+                        seed = Random::RandInt(0, 999999);
+                        SetWindowText(g_seed, std::to_wstring(seed).c_str());
+                                        RedrawWindow(g_seed, NULL, NULL, RDW_UPDATENOW);
+                    }
+                    std::thread([hwnd, seed]{
+                        if (IsDlgButtonChecked(hwnd, CHALLENGE_ONLY)) {
+                            SetWindowText(g_randomizerStatus, L"Randomizing Challenge...");
+                            g_randomizer->RandomizeChallenge(seed);
+                            PostMessage(g_hwnd, WM_COMMAND, RANDOMIZE_CHALLENGE_DONE, NULL);
+                        } else {
+                            SetWindowText(g_randomizerStatus, L"Randomizing...");
+                            g_randomizer->Randomize(seed);
+                            PostMessage(g_hwnd, WM_COMMAND, RANDOMIZE_DONE, NULL);
+                        }
+                    }).detach();
+                }
                 break;
             case RANDOMIZE_DONE:
+                EnableWindow(g_randomizerStatus, FALSE);
                 SetWindowText(g_randomizerStatus, L"Randomized!");
                 break;
+            case RANDOMIZE_CHALLENGE_DONE:
+                EnableWindow(g_randomizerStatus, FALSE);
+                SetWindowText(g_randomizerStatus, L"Randomized Challenge!");
+                break;
+            case CHALLENGE_ONLY:
+                CheckDlgButton(hwnd, CHALLENGE_ONLY, !IsDlgButtonChecked(hwnd, CHALLENGE_ONLY));
+                if (IsWindowEnabled(g_randomizerStatus)) {
+                    PostMessage(g_hwnd, WM_COMMAND, RANDOMIZE_READY, NULL);
+                }
+                break;
         }
     }
     return DefWindowProc(hwnd, message, wParam, lParam);
@@ -86,20 +123,26 @@ HWND CreateLabel(int x, int y, int width, LPCWSTR text) {
                 x, y, width, 16, g_hwnd, NULL, g_hInstance, NULL);
 }
 
-HWND CreateButton(int x, int y, int width, LPCWSTR text, int message) {
+HWND CreateText(int x, int y, int width, LPCWSTR defaultText = L"") {
+        return CreateWindow(MSFTEDIT_CLASS, defaultText,
+        WS_TABSTOP | WS_VISIBLE | WS_CHILD | WS_BORDER,
+        x, y, width, 26, g_hwnd, NULL, g_hInstance, NULL);
+}
+
 #pragma warning(push)
 #pragma warning(disable: 4312)
+HWND CreateButton(int x, int y, int width, LPCWSTR text, int message) {
         return CreateWindow(L"BUTTON", text,
                 WS_TABSTOP | WS_VISIBLE | WS_CHILD | BS_DEFPUSHBUTTON,
                 x, y, width, 26, g_hwnd, (HMENU)message, g_hInstance, NULL);
-#pragma warning(pop)
 }
 
-HWND CreateText(int x, int y, int width, LPCWSTR defaultText = L"") {
-        return CreateWindow(MSFTEDIT_CLASS, L"",
-        WS_TABSTOP | WS_VISIBLE | WS_CHILD | WS_BORDER,
-        x, y, width, 26, g_hwnd, NULL, g_hInstance, NULL);
+HWND CreateCheckbox(int x, int y, int message) {
+        return CreateWindow(L"BUTTON", L"",
+        WS_VISIBLE | WS_CHILD | BS_CHECKBOX,
+        x, y, 12, 12, g_hwnd, (HMENU)message, g_hInstance, NULL);
 }
+#pragma warning(pop)
 
 int APIENTRY wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance, _In_ LPWSTR lpCmdLine, _In_ int nCmdShow) {
         LoadLibrary(L"Msftedit.dll");
@@ -126,9 +169,13 @@ int APIENTRY wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance
 
     CreateLabel(390, 15, 90, L"Version: " VERSION_STR);
     g_seed = CreateText(10, 10, 100);
-    g_randomizerStatus = CreateButton(120, 10, 110, L"Randomize", RANDOMIZING);
+    g_randomizerStatus = CreateButton(120, 10, 180, L"Randomize", RANDOMIZING);
     EnableWindow(g_randomizerStatus, FALSE);
-    PostMessage(g_hwnd, WM_COMMAND, HEARTBEAT, NULL);
+    CreateCheckbox(10, 300, CHALLENGE_ONLY);
+    CreateLabel(30, 300, 200, L"Randomize the challenge only");
+    EnableWindow(g_randomizerStatus, FALSE);
+
+    g_witnessProc->StartHeartbeat(g_hwnd);
 
         ShowWindow(g_hwnd, nCmdShow);
         UpdateWindow(g_hwnd);
@@ -2,25 +2,87 @@
 #include <psapi.h>
 #include <tlhelp32.h>
 #include <iostream>
+#include <string>
 #include <cassert>
 
 #undef PROCESSENTRY32
 #undef Process32Next
 
+Memory::Memory(const std::wstring& processName) : _processName(processName) {
+}
+
+Memory::~Memory() {
+    if (_threadActive) {
+        _threadActive = false;
+        _thread.join();
+    }
+    if (_handle != nullptr) {
+            CloseHandle(_handle);
+    }
+}
+
+void Memory::StartHeartbeat(HWND window, std::chrono::milliseconds beat) {
+    if (_threadActive) return;
+    _threadActive = true;
+    _thread = std::thread([sharedThis = shared_from_this(), window, beat]{
+        while (sharedThis->_threadActive) {
+            sharedThis->Heartbeat(window);
+            std::this_thread::sleep_for(beat);
+        }
+    });
+    _thread.detach();
+}
+
+void Memory::Heartbeat(HWND window) {
+    if (!_handle && !Initialize()) {
+        // Couldn't initialize, definitely not running
+        PostMessage(window, WM_COMMAND, HEARTBEAT, (LPARAM)ProcStatus::NotRunning);
+        return;
+    }
+
+    DWORD exitCode = 0;
+    assert(_handle);
+    GetExitCodeProcess(_handle, &exitCode);
+    if (exitCode != STILL_ACTIVE) {
+        // Process has exited, clean up.
+        _computedAddresses.clear();
+        _handle = NULL;
+        PostMessage(window, WM_COMMAND, HEARTBEAT, (LPARAM)ProcStatus::NotRunning);
+        return;
+    }
+
+#if GLOBALS == 0x5B28C0
+    int currentFrame = ReadData<int>({0x5BE3B0}, 1)[0];
+#elif GLOBALS == 0x62D0A0
+    int currentFrame = ReadData<int>({0x63954C}, 1)[0];
+#endif
+    int frameDelta = currentFrame - _previousFrame;
+    _previousFrame = currentFrame;
+    if (frameDelta < 0 && currentFrame < 250) {
+        PostMessage(window, WM_COMMAND, HEARTBEAT, (LPARAM)ProcStatus::NewGame);
+        return;
+    }
+
+    // TODO: Some way to return ProcStatus::Randomized vs ProcStatus::NotRandomized vs ProcStatus::DeRandomized;
+
+    PostMessage(window, WM_COMMAND, HEARTBEAT, (LPARAM)ProcStatus::Running);
+}
+
+
 [[nodiscard]]
-bool Memory::Initialize(const std::wstring& processName) {
+bool Memory::Initialize() {
         // First, get the handle of the process
         PROCESSENTRY32W entry;
         entry.dwSize = sizeof(entry);
         HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
         while (Process32NextW(snapshot, &entry)) {
-                if (processName == entry.szExeFile) {
+                if (_processName == entry.szExeFile) {
                         _handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, entry.th32ProcessID);
                         break;
                 }
         }
         if (!_handle) {
-                std::cerr << "Couldn't find " << processName.c_str() << ", is it open?" << std::endl;
+                std::cerr << "Couldn't find " << _processName.c_str() << ", is it open?" << std::endl;
         return false;
         }
 
@@ -33,7 +95,7 @@ bool Memory::Initialize(const std::wstring& processName) {
         for (DWORD i = 0; i < numModules / sizeof(HMODULE); i++) {
                 int length = GetModuleBaseNameW(_handle, moduleList[i], &name[0], static_cast<DWORD>(name.size()));
                 name.resize(length);
-                if (processName == name) {
+                if (_processName == name) {
                         _baseAddress = (uintptr_t)moduleList[i];
                         break;
                 }
@@ -45,42 +107,6 @@ bool Memory::Initialize(const std::wstring& processName) {
     return true;
 }
 
-ProcStatus Memory::Heartbeat(const std::wstring& processName) {
-    if (!_handle && !Initialize(processName)) {
-        // Couldn't initialize, definitely not running
-        return ProcStatus::NotRunning;
-    }
-
-    DWORD exitCode = 0;
-    GetExitCodeProcess(_handle, &exitCode);
-    if (exitCode != STILL_ACTIVE) {
-        // Process has exited, clean up.
-        _computedAddresses.clear();
-        _handle = NULL;
-        return ProcStatus::NotRunning;
-    }
-
-    int currentFrame = 0x7FFFFFFF;
-        if (GLOBALS == 0x5B28C0) {
-        currentFrame = ReadData<int>({0x5BE3B0}, 1)[0];
-        } else if (GLOBALS == 0x62D0A0) {
-        currentFrame = ReadData<int>({0x63954C}, 1)[0];
-        } else {
-        assert(false);
-        }
-    if (currentFrame < 80) return ProcStatus::NewGame;
-
-    // TODO: Some way to return ProcStatus::Randomized vs ProcStatus::NotRandomized vs ProcStatus::DeRandomized;
-
-    return ProcStatus::Running;
-}
-
-Memory::~Memory() {
-    if (_handle != nullptr) {
-            CloseHandle(_handle);
-    }
-}
-
 void Memory::AddSigScan(const std::vector<byte>& scanBytes, const std::function<void(int index)>& scanFunc)
 {
         _sigScans[scanBytes] = {scanFunc, false};
@@ -1,12 +1,14 @@
 #pragma once
 #include <functional>
 #include <map>
+#include <thread>
 #include <vector>
 #include <windows.h>
 
 // #define GLOBALS 0x5B28C0
 #define GLOBALS 0x62D0A0
 
+#define HEARTBEAT 0x401
 enum class ProcStatus {
     NotRunning,
     Running,
@@ -17,11 +19,11 @@ enum class ProcStatus {
 // http://stackoverflow.com/q/32798185
 // http://stackoverflow.com/q/36018838
 // http://stackoverflow.com/q/1387064
-class Memory {
+class Memory final : public std::enable_shared_from_this<Memory> {
 public:
-        Memory() = default;
-    ProcStatus Heartbeat(const std::wstring& processName);
+        Memory(const std::wstring& processName);
         ~Memory();
+    void StartHeartbeat(HWND window, std::chrono::milliseconds beat = std::chrono::milliseconds(1000));
 
         Memory(const Memory& memory) = delete;
         Memory& operator=(const Memory& other) = delete;
@@ -74,10 +76,15 @@ private:
                 ThrowError();
         }
 
-        bool Initialize(const std::wstring& processName);
+    void Heartbeat(HWND window);
+        bool Initialize();
         void ThrowError();
         void* ComputeOffset(std::vector<int> offsets);
 
+    int _previousFrame = 0;
+    bool _threadActive = false;
+    std::thread _thread;
+    std::wstring _processName;
         std::map<uintptr_t, uintptr_t> _computedAddresses;
         uintptr_t _baseAddress = 0;
         HANDLE _handle = nullptr;
@@ -158,9 +158,11 @@
   </ItemDefinitionGroup>
   <ItemGroup>
     <ClInclude Include="Memory.h" />
+    <ClInclude Include="Random.h" />
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="Memory.cpp" />
+    <ClCompile Include="Random.cpp" />
   </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">