Locale update - improve remove, extract, read

src/main.cpp

@@ -359,8 +359,9 @@ int main(int argc, char **argv) {
         }
 
         LCID locale = LangToLocale(baseLocale);
-        RemoveFile(hArchive, baseFile, locale);
+        int result = RemoveFile(hArchive, baseFile, locale);
         CloseMpqArchive(hArchive);
+        return result;
     }
 
     // Handle subcommand: List
@@ -397,11 +398,18 @@ int main(int argc, char **argv) {
             std::cout << "[!] Warning: The locale '" << baseLocale << "' is unknown. Will use default locale instead." << std::endl;
         }
 
+        int result;
         if (baseFile != "default") {
-            ExtractFile(hArchive, baseOutput, baseFile, extractKeepFolderStructure, locale);
+            result = ExtractFile(hArchive, baseOutput, baseFile, extractKeepFolderStructure, locale);
         } else {
-            ExtractFiles(hArchive, baseOutput, baseListfileName, locale);
+            result = ExtractFiles(hArchive, baseOutput, baseListfileName, locale);
         }
+        CloseMpqArchive(hArchive);
+
+        if (result != 0) {
+            std::cerr << std::endl << "[!] Failed to extract all files." << std::endl;
+        }
+        return result;
     }
 
     // Handle subcommand: Read

src/mpq.cpp

@@ -31,6 +31,20 @@ int CloseMpqArchive(HANDLE hArchive) {
     return 1;
 }
 
+bool FileExistsInArchiveForLocale(const HANDLE hArchive, const std::string& filePath, const LCID locale) {
+    bool fileExists = false;
+    SFileSetLocale(locale);
+    HANDLE hFile;
+    if (SFileOpenFileEx(hArchive, filePath.c_str(), SFILE_OPEN_FROM_MPQ, &hFile)) {
+        const auto fileLocale = GetFileInfo<int32_t>(hFile, SFileInfoLocale);
+        if (fileLocale == locale) {
+            fileExists = true;
+        }
+    }
+    SFileCloseFile(hFile);
+    return fileExists;
+}
+
 int SignMpqArchive(HANDLE hArchive) {
     if (!SFileSignArchive(hArchive, SIGNATURE_TYPE_WEAK)) {
         std::cerr << "[!] Failed to sign MPQ archive." << std::endl;
@@ -49,39 +63,42 @@ int ExtractFiles(HANDLE hArchive, const std::string& output, const std::string&
     if (findHandle == NULL) {
         std::cerr << "[!] Failed to find first file in MPQ archive." << std::endl;
         SFileCloseArchive(hArchive);
-        return -1;
+        return 1;
     }
 
+    int32_t result = 0;
     do {
-        int32_t result = ExtractFile(
+        result |= ExtractFile(
             hArchive,
             output,
             findData.cFileName,
             true,  // Keep folder structure
             preferredLocale
         );
-        if (result != 0) {
-            return result;
-        }
     } while (SFileFindNextFile(
         findHandle,
         &findData));
 
-    return 0;
+    return result;
 }
 
 int ExtractFile(HANDLE hArchive, const std::string& output, const std::string& fileName, bool keepFolderStructure, LCID preferredLocale) {
     SFileSetLocale(preferredLocale);
     const char *szFileName = fileName.c_str();
-    if (!SFileHasFile(hArchive, szFileName)) {
-        std::cerr << "[!] Failed: File doesn't exist: " << szFileName << std::endl;
-        return -1;
+    if (
+        !FileExistsInArchiveForLocale(hArchive, szFileName, preferredLocale) &&
+        !FileExistsInArchiveForLocale(hArchive, szFileName, defaultLocale)
+    ) {
+        std::cerr << "[!] Failed: File doesn't exist"
+            << PrettyPrintLocale(preferredLocale, " for locale ", true)
+            << ": " << szFileName << std::endl;
+        return 1;
     }
 
     HANDLE hFile;
     if (!SFileOpenFileEx(hArchive, szFileName, SFILE_OPEN_FROM_MPQ, &hFile)) {
         std::cerr << "[!] Failed: File cannot be opened: " << szFileName << std::endl;
-        return -1;
+        return 1;
     }
 
     // Change forward slashes on non-Windows systems
@@ -109,7 +126,7 @@ int ExtractFile(HANDLE hArchive, const std::string& output, const std::string& f
     } else {
         int32_t error = SErrGetLastError();
         std::cerr << "[!] Failed: " << "(" << error << ") " << szFileName << std::endl;
-        return error;
+        return 1;
     }
 
     return 0;
@@ -194,18 +211,11 @@ int AddFile(
         return -1;
     }
 
-    // Check if file exists in MPQ archive
-    SFileSetLocale(locale);
-    HANDLE hFile;
-    if (SFileOpenFileEx(hArchive, archiveFilePath.c_str(), SFILE_OPEN_FROM_MPQ, &hFile)) {
-        int32_t fileLocale = GetFileInfo<int32_t>(hFile, SFileInfoLocale);
-        if (fileLocale == locale) {
-            std::cerr << "[!] File" << PrettyPrintLocale(locale, " for locale ") << " already exists in MPQ archive: " << archiveFilePath
-                      << " - Skipping..." << std::endl;
-            return -1;
-        }
+    if (FileExistsInArchiveForLocale(hArchive, archiveFilePath, locale)) {
+        std::cerr << "[!] File" << PrettyPrintLocale(locale, " for locale ") << " already exists in MPQ archive: "
+            << archiveFilePath << " - Skipping..." << std::endl;
+        return -1;
     }
-    SFileCloseFile(hFile);
     std::cout << "[+] Adding file" << PrettyPrintLocale(locale, " for locale ") << ": " << archiveFilePath << std::endl;
 
     // Verify that we are not exceeding maxFile size of the archive, and if we do, increase it
@@ -256,14 +266,18 @@ int RemoveFile(HANDLE hArchive, const std::string& archiveFilePath, LCID locale)
     SFileSetLocale(locale);
     std::cout << "[-] Removing file" << PrettyPrintLocale(locale, " for locale ") <<": " << archiveFilePath << std::endl;
 
-    if (!SFileHasFile(hArchive, archiveFilePath.c_str())) {
-        std::cerr << "[!] Failed: File doesn't exist" << PrettyPrintLocale(locale, " for locale ") << ": " << archiveFilePath << std::endl;
-        return -1;
+    if (!FileExistsInArchiveForLocale(hArchive, archiveFilePath, locale)) {
+        std::cerr << "[!] Failed: File doesn't exist"
+            << PrettyPrintLocale(locale, " for locale ", true)
+            << ": " << archiveFilePath << std::endl;
+        return 1;
     }
 
     if (!SFileRemoveFile(hArchive, archiveFilePath.c_str(), 0)) {
-        std::cerr << "[!] Failed: File cannot be removed" << PrettyPrintLocale(locale, " for locale ") << ": " << archiveFilePath << std::endl;
-        return -1;
+        std::cerr << "[!] Failed: File cannot be removed"
+            << PrettyPrintLocale(locale, " for locale ", true)
+            << ": " << archiveFilePath << std::endl;
+        return 1;
     }
 
     return 0;
@@ -460,8 +474,13 @@ int ListFiles(HANDLE hArchive, const std::string& listfileName, bool listAll, bo
 
 char* ReadFile(HANDLE hArchive, const char *szFileName, unsigned int *fileSize, LCID preferredLocale) {
     SFileSetLocale(preferredLocale);
-    if (!SFileHasFile(hArchive, szFileName)) {
-        std::cerr << "[!] Failed: File doesn't exist: " << szFileName << std::endl;
+    if (
+        !FileExistsInArchiveForLocale(hArchive, szFileName, preferredLocale) &&
+        !FileExistsInArchiveForLocale(hArchive, szFileName, defaultLocale)
+    ) {
+        std::cerr << "[!] Failed: File doesn't exist"
+            << PrettyPrintLocale(preferredLocale, " for locale ", true)
+            << ": " << szFileName << std::endl;
         return NULL;
     }
 

test/conftest.py

@@ -134,7 +134,7 @@ def generate_mpq_without_internal_listfile(binary_path):
     shutil.rmtree(locales_files_dir, ignore_errors=True)
     locales_files_dir.mkdir(parents=True, exist_ok=True)
 
-    mpq_file_name = data_dir / "mpq_without_internal_listfile2.mpq"
+    mpq_file_name = data_dir / "mpq_without_internal_listfile.mpq"
     mpq_file_name.unlink(missing_ok=True)
 
     content = [