feat: add RCDATA replacing by --set-rcdata key (#77)

AniSkyWorker · zcbenz · commit 3a8e8239daad · 2019-07-01T15:53:29.000+09:00
diff --git a/src/main.cc b/src/main.cc
@@ -23,7 +23,8 @@ void print_help() {
 "  --set-requested-execution-level <level>    Pass nothing to see usage\n"
 "  --application-manifest <path-to-file>      Set manifest file\n"
 "  --set-resource-string <key> <value>        Set resource string\n"
-"  --get-resource-string <key>                Get resource string\n");
+"  --get-resource-string <key>                Get resource string\n"
+"  --set-rcdata <key> <path-to-file>          Replace RCDATA by integer id\n");
 }
 
 bool print_error(const char* message) {
@@ -154,6 +155,18 @@ int wmain(int argc, const wchar_t* argv[]) {
       if (!updater.ChangeString(key_id, value))
         return print_error("Unable to change string");
 
+    } else if (wcscmp(argv[i], L"--set-rcdata") == 0) {
+      if (argc - i < 3)
+        return print_error("--set-rcdata requires int 'Key' and path to resource 'Value'");
+
+      const wchar_t* key = argv[++i];
+      unsigned int key_id = 0;
+      if (swscanf_s(key, L"%d", &key_id) != 1)
+        return print_error("Unable to parse id");
+
+      const wchar_t* pathToResource = argv[++i];
+      if (!updater.ChangeRcData(key_id, pathToResource))
+        return print_error("Unable to change RCDATA");
     } else if (wcscmp(argv[i], L"--get-resource-string") == 0 ||
       wcscmp(argv[i], L"-grs") == 0) {
       if (argc - i < 2)
diff --git a/src/rescle.cc b/src/rescle.cc
@@ -13,6 +13,7 @@
 #include <iomanip> // setw, setfill
 #include <fstream>
 #include <codecvt>
+#include <algorithm>
 
 namespace rescle {
 
@@ -418,6 +419,7 @@ bool ResourceUpdater::Load(const WCHAR* filename) {
   EnumResourceNamesW(module_, RT_GROUP_ICON, OnEnumResourceName, reinterpret_cast<LONG_PTR>(this));
   EnumResourceNamesW(module_, RT_ICON, OnEnumResourceName, reinterpret_cast<LONG_PTR>(this));
   EnumResourceNamesW(module_, RT_MANIFEST, OnEnumResourceManifest, reinterpret_cast<LONG_PTR>(this));
+  EnumResourceNamesW(module_, RT_RCDATA, OnEnumResourceName, reinterpret_cast<LONG_PTR>(this));
 
   return true;
 }
@@ -554,6 +556,43 @@ bool ResourceUpdater::ChangeString(UINT id, const WCHAR* value) {
   return ChangeString(langId, id, value);
 }
 
+bool ResourceUpdater::ChangeRcData(UINT id, const WCHAR* pathToResource) {
+  auto rcDataLngPairIt = std::find_if(rcDataLngMap_.begin(), rcDataLngMap_.end(), [=](const auto& rcDataLngPair) {
+    return rcDataLngPair.second.find(id) != rcDataLngPair.second.end();
+  });
+
+  if (rcDataLngPairIt == rcDataLngMap_.end()) {
+    fprintf(stderr, "Cannot find RCDATA with id '%u'\n", id);
+    return false;
+  }
+
+  wchar_t abspath[MAX_PATH] = { 0 };
+  const auto filePath = _wfullpath(abspath, pathToResource, MAX_PATH) ? abspath : pathToResource;
+  ScopedFile newRcDataFile(filePath);
+  if (newRcDataFile == INVALID_HANDLE_VALUE) {
+    fprintf(stderr, "Cannot open new data file '%ws'\n", filePath);
+    return false;
+  }
+
+  const auto dwFileSize = GetFileSize(newRcDataFile, NULL);
+  if (dwFileSize == INVALID_FILE_SIZE) {
+    fprintf(stderr, "Cannot get file size for '%ws'\n", filePath);
+    return false;
+  }
+
+  auto& rcData = rcDataLngPairIt->second[id];
+  rcData.clear();
+  rcData.resize(dwFileSize);
+
+  DWORD dwBytesRead{ 0 };
+  if (!ReadFile(newRcDataFile, rcData.data(), dwFileSize, &dwBytesRead, NULL)) {
+    fprintf(stderr, "Cannot read file '%ws'\n", filePath);
+    return false;
+  }
+
+  return true;
+}
+
 const WCHAR* ResourceUpdater::GetString(WORD languageId, UINT id) {
   StringTable& table = stringTableMap_[languageId];
 
@@ -762,6 +801,15 @@ bool ResourceUpdater::Commit() {
     }
   }
 
+  for (const auto& rcDataLangPair : rcDataLngMap_) {
+    for (const auto&rcDataMap : rcDataLangPair.second) {
+      if (!UpdateResourceW(ru.Get(), RT_RCDATA, reinterpret_cast<LPWSTR>(rcDataMap.first),
+                           rcDataLangPair.first, (LPVOID)rcDataMap.second.data(), rcDataMap.second.size())) {
+        return false;
+      }
+    }
+  }
+
   for (const auto& iLangIconInfoPair : iconBundleMap_) {
     auto langId = iLangIconInfoPair.first;
     auto maxIconId = iLangIconInfoPair.second.maxIconId;
@@ -864,6 +912,19 @@ BOOL CALLBACK ResourceUpdater::OnEnumResourceLanguage(HANDLE hModule, LPCWSTR lp
         instance->iconBundleMap_[wIDLanguage].iconBundles[iconId] = nullptr;
         break;
       }
+      case reinterpret_cast<ptrdiff_t>(RT_RCDATA): {
+        const auto moduleHandle = HMODULE(hModule);
+        HRSRC hResInfo = FindResource(moduleHandle, lpszName, lpszType);
+        DWORD cbResource = SizeofResource(moduleHandle, hResInfo);
+        HGLOBAL hResData = LoadResource(moduleHandle, hResInfo);
+
+        const auto *pResource = (const BYTE *)LockResource(hResData);
+        const auto resId = reinterpret_cast<ptrdiff_t>(lpszName);
+        instance->rcDataLngMap_[wIDLanguage][resId] = std::vector<BYTE>(pResource, pResource + cbResource);
+
+        UnlockResource(hResData);
+        FreeResource(hResData);
+      }
       default:
         break;
     }
diff --git a/src/rescle.h b/src/rescle.h
@@ -109,6 +109,9 @@ class ResourceUpdater {
   typedef std::map<WORD, StringTable> StringTableMap;
   typedef std::map<LANGID, VersionInfo> VersionStampMap;
   typedef std::map<UINT, std::unique_ptr<IconsValue>> IconTable;
+  typedef std::vector<BYTE> RcDataValue;
+  typedef std::map<ptrdiff_t, RcDataValue> RcDataMap;
+  typedef std::map<LANGID, RcDataMap> RcDataLangMap;
 
   struct IconResInfo {
     UINT maxIconId = 0;
@@ -131,6 +134,7 @@ class ResourceUpdater {
   bool SetFileVersion(unsigned short v1, unsigned short v2, unsigned short v3, unsigned short v4);
   bool ChangeString(WORD languageId, UINT id, const WCHAR* value);
   bool ChangeString(UINT id, const WCHAR* value);
+  bool ChangeRcData(UINT id, const WCHAR* pathToResource);
   const WCHAR* GetString(WORD languageId, UINT id);
   const WCHAR* GetString(UINT id);
   bool SetIcon(const WCHAR* path, const LANGID& langId, UINT iconBundle);
@@ -158,6 +162,7 @@ class ResourceUpdater {
   VersionStampMap versionStampMap_;
   StringTableMap stringTableMap_;
   IconTableMap iconBundleMap_;
+  RcDataLangMap rcDataLngMap_;
 };
 
 class ScopedResourceUpdater {
