From 5e721aab6c1519bfc74cfc55f70da5b6275f33a2 Mon Sep 17 00:00:00 2001
From: qtkite <kite@uwu.icu>
Date: Tue, 20 Jul 2021 03:24:54 +1000
Subject: [PATCH] refactoring dcontrol.cpp code

---
 src/defender-control/dcontrol.cpp             | 180 +-----------------
 src/defender-control/dcontrol.hpp             |  12 +-
 src/defender-control/defender-control.vcxproj |   3 +
 .../defender-control.vcxproj.filters          |  26 ++-
 src/defender-control/main.cpp                 |  13 +-
 src/defender-control/reg.cpp                  | 106 +++++++++++
 src/defender-control/reg.hpp                  |  13 ++
 src/defender-control/settings.hpp             |   3 +
 src/defender-control/util.cpp                 |  70 +++++++
 src/defender-control/util.hpp                 |  12 ++
 10 files changed, 234 insertions(+), 204 deletions(-)
 create mode 100644 src/defender-control/reg.cpp
 create mode 100644 src/defender-control/reg.hpp
 create mode 100644 src/defender-control/settings.hpp

diff --git a/src/defender-control/dcontrol.cpp b/src/defender-control/dcontrol.cpp
index c9897fb..7528c3e 100644
--- a/src/defender-control/dcontrol.cpp
+++ b/src/defender-control/dcontrol.cpp
@@ -1,194 +1,18 @@
 #include "dcontrol.hpp"
 
-namespace REG
-{
-  // reads a key from HKEY_LOCAL_MACHINE
-  //
-  DWORD read_key(const wchar_t* root_name, const wchar_t* value_name, uint32_t flags)
-  {
-    LSTATUS status;
-    HKEY hkey;
-    DWORD result{};
-    DWORD buff_sz = sizeof(DWORD);
-
-    // https://docs.microsoft.com/en-us/windows/win32/winprog64/accessing-an-alternate-registry-view
-
-    status = RegOpenKeyExW(
-      HKEY_LOCAL_MACHINE,
-      root_name,
-      0,
-      KEY_READ | KEY_WOW64_64KEY,
-      &hkey
-    );
-
-    if (status)
-    {
-      if (flags & DBG_MSG)
-        std::cout << "Error opening " << root_name << " key" << std::endl;
-
-      return -1;
-    }
-
-    status = RegQueryValueExW(
-      hkey,
-      value_name,
-      0, NULL,
-      reinterpret_cast<LPBYTE>(&result),
-      &buff_sz
-    );
-
-    if (status)
-    {
-      if (flags & DBG_MSG)
-        std::cout << "Failed to read " << result << std::endl;
-
-      return -1;
-    }
-
-    RegCloseKey(hkey);
-
-    return result;
-  }
-
-  // creates a registry in HKEY_LOCAL_MACHINE with KEY_ALL_ACCESS permissions
-  //
-  bool create_registry(const wchar_t* root_name, HKEY& hkey)
-  {
-    LSTATUS status;
-
-    DWORD dwDisposition;
-
-    status = RegCreateKeyExW(
-      HKEY_LOCAL_MACHINE,
-      root_name,
-      0,
-      0,
-      0,
-      131334,
-      0,
-      &hkey,
-      &dwDisposition
-    );
-
-    if (status)
-    {
-      std::wcout << "could not find or create " << root_name << " error: " << status << std::endl;
-      return false;
-    }
-
-    return true;
-  }
-
-  bool set_keyval(HKEY& hkey, const wchar_t* value_name, DWORD value)
-  {
-    auto ret = RegSetValueExW(hkey, value_name, 0, REG_DWORD,
-      reinterpret_cast<LPBYTE>(&value), 4);
-
-    if (ret)
-    {
-      std::cout << "Set error: " << ret << std::endl;
-      return false;
-    }
-
-    return true;
-  }
-
-  bool set_keyval_bin(HKEY& hkey, const wchar_t* value_name, DWORD value)
-  {
-    auto ret = RegSetValueExW(hkey, value_name, 0, REG_BINARY,
-      reinterpret_cast<LPBYTE>(&value), 12);
-
-    if (ret)
-    {
-      std::cout << "Set error: " << ret << std::endl;
-      return false;
-    }
-    return true;
-  }
-}
-
-namespace WMIC
-{
-}
-
 namespace DCONTROL
 {
-  // Sets the programs debug priviliges
-  bool set_privilege(LPCSTR privilege, BOOL enable)
-  {
-    TOKEN_PRIVILEGES priv = { 0,0,0,0 };
-    HANDLE token = nullptr;
-    LUID luid = { 0,0 };
-
-    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &token)) 
-    {
-      if (token)
-        CloseHandle(token);
-
-      return false;
-    }
-
-    if (!LookupPrivilegeValueA(nullptr, SE_DEBUG_NAME, &luid)) 
-    {
-      if (token)
-        CloseHandle(token);
-
-      return false;
-    }
-    priv.PrivilegeCount = 1;
-    priv.Privileges[0].Luid = luid;
-    priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
-
-    if (!AdjustTokenPrivileges(token, false, &priv, 0, nullptr, nullptr))
-    {
-      if (token)
-        CloseHandle(token);
-
-      return false;
-    }
-    if (token)
-      CloseHandle(token);
-
-    return true;
-  }
-
-  char sub_43604B()
-  {
-    char v0; // bl
-    SC_HANDLE v1; // eax
-    SC_HANDLE v2; // esi
-    void* v3; // eax
-
-    v0 = 0;
-    v1 = OpenSCManagerW(0, 0, 8u);
-    v2 = v1;
-    if (v1)
-    {
-      v3 = LockServiceDatabase(v1);
-      if (v3)
-      {
-        UnlockServiceDatabase(v3);
-        CloseServiceHandle(v2);
-        return 1;
-      }
-      if (GetLastError() == 1055)
-        v0 = 1;
-      CloseServiceHandle(v2);
-    }
-    return v0;
-  }
-
   // disables window defender
   //
   bool disable_defender()
   {
-    if (!sub_43604B())
+    if (!util::sub_43604B())
     {
       std::cout << "permission error" << std::endl;
       return false;
     }
 
-    set_privilege(SE_DEBUG_NAME, TRUE);
+    util::set_privilege(SE_DEBUG_NAME, TRUE);
 
     HKEY hkey;
 
diff --git a/src/defender-control/dcontrol.hpp b/src/defender-control/dcontrol.hpp
index 14af5f1..b64f361 100644
--- a/src/defender-control/dcontrol.hpp
+++ b/src/defender-control/dcontrol.hpp
@@ -3,15 +3,9 @@
 #include <Windows.h>
 #include <iostream>
 
-#define DBG_MSG (1 << 0)
-
-namespace REG
-{
-  DWORD read_key(const wchar_t* root_name, const wchar_t* value_name, uint32_t flags = 0);
-  bool create_registry(const wchar_t* root_name, HKEY& hkey);
-  bool set_keyval(HKEY& hkey, const wchar_t* value_name, DWORD value);
-  bool set_keyval_bin(HKEY& hkey, const wchar_t* value_name, DWORD value);
-}
+#include "settings.hpp"
+#include "reg.hpp"
+#include "util.hpp"
 
 namespace DCONTROL
 {
diff --git a/src/defender-control/defender-control.vcxproj b/src/defender-control/defender-control.vcxproj
index 60e39f9..34f1e4a 100644
--- a/src/defender-control/defender-control.vcxproj
+++ b/src/defender-control/defender-control.vcxproj
@@ -143,11 +143,14 @@
   <ItemGroup>
     <ClCompile Include="dcontrol.cpp" />
     <ClCompile Include="main.cpp" />
+    <ClCompile Include="reg.cpp" />
     <ClCompile Include="util.cpp" />
     <ClCompile Include="wmic.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="dcontrol.hpp" />
+    <ClInclude Include="settings.hpp" />
+    <ClInclude Include="reg.hpp" />
     <ClInclude Include="util.hpp" />
     <ClInclude Include="wmic.hpp" />
   </ItemGroup>
diff --git a/src/defender-control/defender-control.vcxproj.filters b/src/defender-control/defender-control.vcxproj.filters
index 560a9bd..90ef69b 100644
--- a/src/defender-control/defender-control.vcxproj.filters
+++ b/src/defender-control/defender-control.vcxproj.filters
@@ -16,12 +16,15 @@
     <Filter Include="Source Files\defender-control">
       <UniqueIdentifier>{8a88e18b-d3f3-447e-a3b0-9867c153c3c1}</UniqueIdentifier>
     </Filter>
-    <Filter Include="Source Files\defender-control\wmic">
-      <UniqueIdentifier>{db30358b-f563-460e-92fb-eacafe3a68cb}</UniqueIdentifier>
+    <Filter Include="Source Files\reg">
+      <UniqueIdentifier>{1d1c7c6d-6f90-49ce-a823-f7deee4d955d}</UniqueIdentifier>
     </Filter>
-    <Filter Include="Source Files\defender-control\util">
+    <Filter Include="Source Files\util">
       <UniqueIdentifier>{0cabcba0-6285-4ad6-9295-add0563d0d48}</UniqueIdentifier>
     </Filter>
+    <Filter Include="Source Files\wmic">
+      <UniqueIdentifier>{db30358b-f563-460e-92fb-eacafe3a68cb}</UniqueIdentifier>
+    </Filter>
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="main.cpp">
@@ -31,10 +34,13 @@
       <Filter>Source Files\defender-control</Filter>
     </ClCompile>
     <ClCompile Include="wmic.cpp">
-      <Filter>Source Files\defender-control\wmic</Filter>
+      <Filter>Source Files\wmic</Filter>
     </ClCompile>
     <ClCompile Include="util.cpp">
-      <Filter>Source Files\defender-control\util</Filter>
+      <Filter>Source Files\util</Filter>
+    </ClCompile>
+    <ClCompile Include="reg.cpp">
+      <Filter>Source Files\reg</Filter>
     </ClCompile>
   </ItemGroup>
   <ItemGroup>
@@ -42,10 +48,16 @@
       <Filter>Source Files\defender-control</Filter>
     </ClInclude>
     <ClInclude Include="wmic.hpp">
-      <Filter>Source Files\defender-control\wmic</Filter>
+      <Filter>Source Files\wmic</Filter>
     </ClInclude>
     <ClInclude Include="util.hpp">
-      <Filter>Source Files\defender-control\util</Filter>
+      <Filter>Source Files\util</Filter>
+    </ClInclude>
+    <ClInclude Include="reg.hpp">
+      <Filter>Source Files\reg</Filter>
+    </ClInclude>
+    <ClInclude Include="settings.hpp">
+      <Filter>Source Files</Filter>
     </ClInclude>
   </ItemGroup>
 </Project>
\ No newline at end of file
diff --git a/src/defender-control/main.cpp b/src/defender-control/main.cpp
index 005aafd..9ddb22d 100644
--- a/src/defender-control/main.cpp
+++ b/src/defender-control/main.cpp
@@ -1,24 +1,17 @@
-#include "dcontrol.hpp"
-#include "wmic.hpp"
-
 // to-do:
 // write argument parser
 // create cli program
 // maybe make a ui for this
-
-// entrypoint
 //
+#include "dcontrol.hpp"
+#include "wmic.hpp"
+
 int main()
 {
   printf(DCONTROL::check_defender() ?
     "Windows defender is ACTIVE turning off..\n" :
     "Windows defender is OFF turning on...\n");
 
-  //if (DCONTROL::check_defender())
-  //  wmic::test_exec(true);
-  //else
-  //  wmic::test_exec(false);
-
   auto helper = new wmic::helper(
     "Root\\Microsoft\\Windows\\Defender",
     "MSFT_MpPreference",
diff --git a/src/defender-control/reg.cpp b/src/defender-control/reg.cpp
new file mode 100644
index 0000000..8653e3c
--- /dev/null
+++ b/src/defender-control/reg.cpp
@@ -0,0 +1,106 @@
+#include "reg.hpp"
+
+namespace REG
+{
+  // reads a key from HKEY_LOCAL_MACHINE
+  //
+  DWORD read_key(const wchar_t* root_name, const wchar_t* value_name, uint32_t flags)
+  {
+    LSTATUS status;
+    HKEY hkey;
+    DWORD result{};
+    DWORD buff_sz = sizeof(DWORD);
+
+    // https://docs.microsoft.com/en-us/windows/win32/winprog64/accessing-an-alternate-registry-view
+
+    status = RegOpenKeyExW(
+      HKEY_LOCAL_MACHINE,
+      root_name,
+      0,
+      KEY_READ | KEY_WOW64_64KEY,
+      &hkey
+    );
+
+    if (status)
+    {
+      if (flags & DBG_MSG)
+        std::cout << "Error opening " << root_name << " key" << std::endl;
+
+      return -1;
+    }
+
+    status = RegQueryValueExW(
+      hkey,
+      value_name,
+      0, NULL,
+      reinterpret_cast<LPBYTE>(&result),
+      &buff_sz
+    );
+
+    if (status)
+    {
+      if (flags & DBG_MSG)
+        std::cout << "Failed to read " << result << std::endl;
+
+      return -1;
+    }
+
+    RegCloseKey(hkey);
+
+    return result;
+  }
+
+  // creates a registry in HKEY_LOCAL_MACHINE with KEY_ALL_ACCESS permissions
+  //
+  bool create_registry(const wchar_t* root_name, HKEY& hkey)
+  {
+    LSTATUS status;
+
+    DWORD dwDisposition;
+
+    status = RegCreateKeyExW(
+      HKEY_LOCAL_MACHINE,
+      root_name,
+      0, 0, 0,
+      131334,
+      0,
+      &hkey,
+      &dwDisposition
+    );
+
+    if (status)
+    {
+      std::wcout << "could not find or create " << root_name << " error: " << status << std::endl;
+      return false;
+    }
+
+    return true;
+  }
+
+  bool set_keyval(HKEY& hkey, const wchar_t* value_name, DWORD value)
+  {
+    auto ret = RegSetValueExW(hkey, value_name, 0, REG_DWORD,
+      reinterpret_cast<LPBYTE>(&value), 4);
+
+    if (ret)
+    {
+      std::cout << "Set error: " << ret << std::endl;
+      return false;
+    }
+
+    return true;
+  }
+
+  bool set_keyval_bin(HKEY& hkey, const wchar_t* value_name, DWORD value)
+  {
+    auto ret = RegSetValueExW(hkey, value_name, 0, REG_BINARY,
+      reinterpret_cast<LPBYTE>(&value), 12);
+
+    if (ret)
+    {
+      std::cout << "Set error: " << ret << std::endl;
+      return false;
+    }
+    return true;
+  }
+}
diff --git a/src/defender-control/reg.hpp b/src/defender-control/reg.hpp
new file mode 100644
index 0000000..3aaa55a
--- /dev/null
+++ b/src/defender-control/reg.hpp
@@ -0,0 +1,13 @@
+#pragma once
+#include <Windows.h>
+#include <cstdint>
+#include <iostream>
+#include "settings.hpp"
+
+namespace REG
+{
+  DWORD read_key(const wchar_t* root_name, const wchar_t* value_name, uint32_t flags = 0);
+  bool create_registry(const wchar_t* root_name, HKEY& hkey);
+  bool set_keyval(HKEY& hkey, const wchar_t* value_name, DWORD value);
+  bool set_keyval_bin(HKEY& hkey, const wchar_t* value_name, DWORD value);
+}
diff --git a/src/defender-control/settings.hpp b/src/defender-control/settings.hpp
new file mode 100644
index 0000000..31288f0
--- /dev/null
+++ b/src/defender-control/settings.hpp
@@ -0,0 +1,3 @@
+#pragma once
+
+#define DBG_MSG (1 << 0)
diff --git a/src/defender-control/util.cpp b/src/defender-control/util.cpp
index 0557d88..f04b66b 100644
--- a/src/defender-control/util.cpp
+++ b/src/defender-control/util.cpp
@@ -2,6 +2,8 @@
 
 namespace util
 {
+  // Converts a string to wide
+  //
   std::wstring string_to_wide(const std::string& s)
   {
     std::wstring temp(s.length(), L' ');
@@ -9,9 +11,77 @@ namespace util
     return temp;
   }
 
+  // Converts a wide to string
+  //
   std::string wide_to_string(const std::wstring& s) {
     std::string temp(s.length(), ' ');
     std::copy(s.begin(), s.end(), temp.begin());
     return temp;
   }
+
+  // Sets the programs debug priviliges
+  //
+  bool set_privilege(LPCSTR privilege, BOOL enable)
+  {
+    TOKEN_PRIVILEGES priv = { 0,0,0,0 };
+    HANDLE token = nullptr;
+    LUID luid = { 0,0 };
+
+    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &token)) 
+    {
+      if (token)
+        CloseHandle(token);
+
+      return false;
+    }
+
+    if (!LookupPrivilegeValueA(nullptr, SE_DEBUG_NAME, &luid)) 
+    {
+      if (token)
+        CloseHandle(token);
+
+      return false;
+    }
+    priv.PrivilegeCount = 1;
+    priv.Privileges[0].Luid = luid;
+    priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
+
+    if (!AdjustTokenPrivileges(token, false, &priv, 0, nullptr, nullptr))
+    {
+      if (token)
+        CloseHandle(token);
+
+      return false;
+    }
+    if (token)
+      CloseHandle(token);
+
+    return true;
+  }
+
+  char sub_43604B()
+  {
+    char v0; // bl
+    SC_HANDLE v1; // eax
+    SC_HANDLE v2; // esi
+    void* v3; // eax
+
+    v0 = 0;
+    v1 = OpenSCManagerW(0, 0, 8u);
+    v2 = v1;
+    if (v1)
+    {
+      v3 = LockServiceDatabase(v1);
+      if (v3)
+      {
+        UnlockServiceDatabase(v3);
+        CloseServiceHandle(v2);
+        return 1;
+      }
+      if (GetLastError() == 1055)
+        v0 = 1;
+      CloseServiceHandle(v2);
+    }
+    return v0;
+  }
 }
\ No newline at end of file
diff --git a/src/defender-control/util.hpp b/src/defender-control/util.hpp
index 3b96b7d..379c663 100644
--- a/src/defender-control/util.hpp
+++ b/src/defender-control/util.hpp
@@ -1,8 +1,20 @@
 #pragma once
+#include <Windows.h>
 #include <string>
 
 namespace util
 {
+  // Converts a string to wide
+  //
   std::wstring string_to_wide(const std::string& s);
+
+  // Converts a wide to string
+  //
   std::string wide_to_string(const std::wstring& s);
+
+  // Sets the programs debug priviliges
+  //
+  bool set_privilege(LPCSTR privilege, BOOL enable);
+
+  char sub_43604B();
 }