refactor: allow custom path to db and infra.dll

src/main.cpp

@@ -22,25 +22,22 @@ constexpr std::array<uint8_t, 16> kMagicHeader{0x7C, 0xD5, 0x32, 0xEB, 0x86, 0x0
                                                0xA8, 0xAF, 0xA6, 0x8E, 0x0F, 0xFF, 0x99, 0x14};
 
 void DecryptKGG(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& kgg_file) {
-    auto kgg_fname = kgg_file.filename().wstring();
-    fwprintf(stderr, L"[INFO] reading %s...\n", kgg_fname.c_str());
-
     std::ifstream ifs_kgg(kgg_file, std::ios::binary);
     char header[0x100]{};
     ifs_kgg.read(header, sizeof(ifs_kgg));
     if (std::equal(kMagicHeader.cbegin(), kMagicHeader.cend(), header)) {
-        fprintf(stderr, "[WARN] invalid kgg header, skip.\n");
+        fprintf(stderr, "[WARN] ! invalid kgg header, skip.\n");
         return;
     }
     uint32_t offset_to_audio = *reinterpret_cast<uint32_t*>(&header[0x10]);
     uint32_t encrypt_mode = *reinterpret_cast<uint32_t*>(&header[0x14]);
     if (encrypt_mode != 5) {
-        fprintf(stderr, "[WARN] invalid enc_version (expect=0x05, got 0x%02x), skip.\n", encrypt_mode);
+        fprintf(stderr, "[WARN] ! invalid enc_version (expect=0x05, got 0x%02x), skip.\n", encrypt_mode);
         return;
     }
     uint32_t audio_hash_len = *reinterpret_cast<uint32_t*>(&header[0x44]);
     if (audio_hash_len != 0x20) {
-        fprintf(stderr, "[WARN] audio hash length invalid (expect=0x20, got 0x%02x), skip.\n", audio_hash_len);
+        fprintf(stderr, "[WARN] ! audio hash length invalid (expect=0x20, got 0x%02x), skip.\n", audio_hash_len);
         return;
     }
     std::string audio_hash(&header[0x48], &header[0x48 + audio_hash_len]);
@@ -48,19 +45,19 @@ void DecryptKGG(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& kgg_f
     if (auto it = ekey_db.find(audio_hash); it != ekey_db.end()) {
         ekey = it->second;
     } else {
-        fprintf(stderr, "[WARN] ekey not found, skip.\n");
+        fprintf(stderr, "[WARN] ! ekey not found, skip.\n");
         return;
     }
 
     auto qmc2 = QMC2::Create(ekey);
     if (!qmc2) {
-        fprintf(stderr, "[WARN] create qmc2 failed, skip.\n");
+        fprintf(stderr, "[WARN] ! create qmc2 failed, skip.\n");
         return;
     }
 
     ifs_kgg.seekg(offset_to_audio, std::ios::beg);
     std::vector<uint8_t> buffer(1024 * 1024);
-    ifs_kgg.read(reinterpret_cast<char*>(buffer.data()), buffer.size());
+    ifs_kgg.read(reinterpret_cast<char*>(buffer.data()), 4);
     auto n = static_cast<size_t>(ifs_kgg.gcount());
     qmc2->Decrypt(std::span(buffer.begin(), n), 0);
 
@@ -74,7 +71,7 @@ void DecryptKGG(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& kgg_f
         decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.mp3");
     }
     if (exists(decrypted_path)) {
-        fprintf(stderr, "[WARN] output file '%s' exists, skip.\n", decrypted_path.filename().string().c_str());
+        fprintf(stderr, "[WARN] ! output file '%s' exists, skip.\n", decrypted_path.filename().string().c_str());
         return;
     }
 
@@ -88,9 +85,11 @@ void DecryptKGG(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& kgg_f
         ofs_decrypted.write(reinterpret_cast<char*>(buffer.data()), n);
         offset += n;
     }
+    qmc2.reset();
 
-    auto decrypted_fname = decrypted_path.filename().wstring();
+    auto kgg_fname = kgg_file.filename();
-    fwprintf(stderr, L"[INFO] OK! kgg --> %s\n", decrypted_fname.c_str());
+    auto decrypted_fname = decrypted_path.filename();
+    fwprintf(stderr, L"[INFO] * OK! %s --> %s\n", kgg_fname.c_str(), decrypted_fname.c_str());
 }
 
 void WalkFileOrDir(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& input_path) {
@@ -100,23 +99,86 @@ void WalkFileOrDir(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& in
         }
     } else if (is_regular_file(input_path)) {
         DecryptKGG(ekey_db, input_path);
+    } else {
+        fwprintf(stderr, L"[WARN] invalid path: %s\n", input_path.c_str());
     }
 }
 
+std::pair<std::vector<std::wstring>, std::unordered_map<std::wstring, std::wstring>> ParseCommandLine() {
+    int argc;
+    wchar_t** argv = CommandLineToArgvW(GetCommandLineW(), &argc);
+
+    std::vector<std::wstring> positional_args{};
+    std::unordered_map<std::wstring, std::wstring> named_args{};
+
+    bool positional_only{false};
+    for (int i = 1; i < argc; i++) {
+        std::wstring arg{argv[i]};
+        if (arg == L"--") {
+            positional_only = true;
+            continue;
+        }
+
+        if (!positional_only && arg.starts_with(L"--")) {
+            auto pos = arg.find(L'=');
+            if (pos != std::wstring::npos) {
+                named_args[arg.substr(2, pos - 2)] = arg.substr(pos + 1);
+            } else if (++i < argc) {
+                named_args[arg.substr(2)] = argv[i];
+            } else {
+                named_args[arg.substr(2)] = L"";
+            }
+        } else {
+            positional_args.push_back(arg);
+        }
+    }
+
+    return std::make_pair(positional_args, named_args);
+}
+
 int main() {
     SetConsoleOutputCP(CP_UTF8);
     setlocale(LC_ALL, ".UTF8");
 
     fputs("kgg-dec v0.1 by LSR\n", stderr);
-    fputs("Usage: kgg-dec [kgg-dir=.]\n\n", stderr);
+    fputs(
+        "Usage: kgg-dec "
+        "[--infra-dll infra.dll] "
+        "[--db /path/to/KGMusicV3.db] "
+        "[--] [kgg-dir... = '.']\n\n",
+        stderr);
+
+    auto [positional_args, named_args] = ParseCommandLine();
+    if (positional_args.empty()) {
+        positional_args.emplace_back(L".");
+    }
+
+    std::filesystem::path kgm_db_path{};
+    std::filesystem::path infra_dll_path{L"infra.dll"};
+    if (auto it = named_args.find(L"infra-dll"); it != named_args.end()) {
+        infra_dll_path = std::filesystem::path{it->second};
+    }
+    infra_dll_path = absolute(infra_dll_path);
+    if (!exists(infra_dll_path)) {
+        fputs("[ERR ] infra.dll not found\n", stderr);
+        return 1;
+    }
+
+    if (auto it = named_args.find(L"db"); it != named_args.end()) {
+        kgm_db_path = std::filesystem::path{it->second};
+    } else {
         PWSTR pAppDirPath{};
         SHGetKnownFolderPath(FOLDERID_RoamingAppData, 0, nullptr, &pAppDirPath);
-    std::filesystem::path kgmPath{pAppDirPath};
+        kgm_db_path = std::filesystem::path{pAppDirPath} / L"Kugou8" / L"KGMusicV3.db";
         CoTaskMemFree(pAppDirPath);
-    kgmPath = kgmPath / L"Kugou8" / L"KGMusicV3.db";
+    }
+    if (!exists(kgm_db_path)) {
+        fputs("[ERR ] KGMusicV3.db not found\n", stderr);
+        return 1;
+    }
 
     int error{-1};
-    Infra::KugouDb db{L"infra.dll", kgmPath};
+    Infra::KugouDb db{infra_dll_path, kgm_db_path};
     if (!db.IsOpen()) {
         fprintf(stderr, "[ERR ] db init error: is infra.dll ok?\n");
         return 1;
@@ -135,12 +197,9 @@ int main() {
     }
 #endif
 
-    int argc;
+    for (auto& positional_arg : positional_args) {
-    wchar_t** argv = CommandLineToArgvW(GetCommandLineW(), &argc);
+        WalkFileOrDir(ekey_db, positional_arg);
-
+    }
-    std::filesystem::path input_dir{(argc < 2) ? L"." : argv[1]};
-    input_dir = absolute(input_dir);
-    WalkFileOrDir(ekey_db, input_dir);
 
     return 0;
 }