From ac8c5d9a1033fe7c18fe66fe8f94c570a139c4f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=B2=81=E6=A0=91=E4=BA=BA?= Date: Fri, 18 Oct 2024 08:11:34 +0900 Subject: [PATCH] refactor: use threads! --- src/jobs.hpp | 170 ++++++++++++++++++++++++++++++++++++++++++ src/main.cpp | 141 +++++++++++++---------------------- src/qmc2/qmc2.h | 12 +-- src/qmc2/qmc2_map.cpp | 2 +- src/qmc2/qmc2_rc4.cpp | 14 ++-- 5 files changed, 236 insertions(+), 103 deletions(-) create mode 100644 src/jobs.hpp diff --git a/src/jobs.hpp b/src/jobs.hpp new file mode 100644 index 0000000..8ec92bd --- /dev/null +++ b/src/jobs.hpp @@ -0,0 +1,170 @@ +#pragma once +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "qmc2/qmc2.h" + +class KggTask { + public: + explicit KggTask(std::filesystem::path kgg_path, std::filesystem::path out_dir) + : kgg_path_(std::move(kgg_path)), out_dir_(std::move(out_dir)) {} + + void warning(const wchar_t* msg) const { fwprintf(stderr, L"[WARN] %s (%s)\n", msg, kgg_path_.filename().c_str()); } + void warning(const std::wstring& msg) const { warning(msg.c_str()); } + + void error(const wchar_t* msg) const { fwprintf(stderr, L"[ERR ] %s (%s)\n", msg, kgg_path_.filename().c_str()); } + void error(const std::wstring& msg) const { error(msg.c_str()); } + + void info(const wchar_t* msg) const { fwprintf(stderr, L"[INFO] %s (%s)\n", msg, kgg_path_.filename().c_str()); } + void info(const std::wstring& msg) const { info(msg.c_str()); } + + void Execute(const Infra::kgm_ekey_db_t& ekey_db) { + constexpr static std::array kMagicHeader{0x7C, 0xD5, 0x32, 0xEB, 0x86, 0x02, 0x7F, 0x4B, + 0xA8, 0xAF, 0xA6, 0x8E, 0x0F, 0xFF, 0x99, 0x14}; + + std::ifstream kgg_stream_in(kgg_path_, std::ios::binary); + char header[0x100]{}; + kgg_stream_in.read(header, sizeof(kgg_stream_in)); + if (std::equal(kMagicHeader.cbegin(), kMagicHeader.cend(), header)) { + warning(L"invalid kgg header"); + return; + } + uint32_t offset_to_audio = *reinterpret_cast(&header[0x10]); + uint32_t encrypt_mode = *reinterpret_cast(&header[0x14]); + if (encrypt_mode != 5) { + warning(std::format(L"unsupported enc_version (expect=0x05, got 0x{:02x})", encrypt_mode)); + return; + } + uint32_t audio_hash_len = *reinterpret_cast(&header[0x44]); + if (audio_hash_len != 0x20) { + warning(std::format(L"audio hash length invalid (expect=0x20, got 0x{:02x})", audio_hash_len)); + return; + } + std::string audio_hash(&header[0x48], &header[0x48 + audio_hash_len]); + std::string ekey{}; + if (auto it = ekey_db.find(audio_hash); it != ekey_db.end()) { + ekey = it->second; + } else { + warning(L"ekey not found"); + return; + } + + auto qmc2 = QMC2::Create(ekey); + + std::string magic(4, 0); + kgg_stream_in.seekg(offset_to_audio, std::ios::beg); + kgg_stream_in.read(magic.data(), 4); + qmc2->Decrypt(std::span(reinterpret_cast(magic.data()), 4), 0); + auto real_ext = DetectRealExt(magic); + auto out_path = out_dir_ / std::format(L"{}_kgg-dec.{}", kgg_path_.stem().wstring(), real_ext); + + if (exists(out_path)) { + warning(std::format(L"output file already exists: {}", out_path.filename().wstring())); + return; + } + + kgg_stream_in.seekg(offset_to_audio, std::ios::beg); + std::ofstream ofs_decrypted(out_path, std::ios::binary); + if (!ofs_decrypted.is_open()) { + error(L"failed to open output file"); + return; + } + + size_t offset{0}; + thread_local std::array temp_buffer{}; + auto buf_signed = std::span(reinterpret_cast(temp_buffer.data()), temp_buffer.size()); + auto buf_unsigned = std::span(temp_buffer); + + while (!kgg_stream_in.eof()) { + kgg_stream_in.read(buf_signed.data(), buf_signed.size()); + const auto n = static_cast(kgg_stream_in.gcount()); + qmc2->Decrypt(buf_unsigned.subspan(0, n), offset); + ofs_decrypted.write(buf_signed.data(), n); + offset += n; + } + + info(std::format(L"** OK ** -> {}", out_path.filename().wstring())); + } + + private: + std::filesystem::path kgg_path_; + std::filesystem::path out_dir_; + + static const wchar_t* DetectRealExt(std::string_view magic) { + if (magic == "fLaC") { + return L"flac"; + } + if (magic == "OggS") { + return L"ogg"; + } + return L"mp3"; + } +}; + +class KggTaskQueue { + public: + explicit KggTaskQueue(Infra::kgm_ekey_db_t ekey_db) : ekey_db_(std::move(ekey_db)) {} + + void Push(std::unique_ptr task) { + std::lock_guard lock(mutex_); + tasks_.push(std::move(task)); + signal_.notify_one(); + } + + std::unique_ptr Pop() { + std::unique_lock lock(mutex_); + signal_.wait(lock, [this] { return !tasks_.empty() || thread_end_; }); + if (tasks_.empty()) { + return {}; + } + + auto task = std::move(tasks_.front()); + tasks_.pop(); + return task; + } + + [[nodiscard]] bool Finished() { + std::lock_guard lock(mutex_); + return tasks_.empty(); + } + + void AddWorkerThread() { threads_.emplace(&KggTaskQueue::WorkerThreadBody, this); } + + void Join() { + thread_end_ = true; + signal_.notify_all(); + + for (int i = 1; !threads_.empty(); i++) { + threads_.front().join(); + threads_.pop(); +#ifndef NDEBUG + fprintf(stderr, "[INFO] thread %d joined\n", i); +#endif + } + } + + private: + bool thread_end_{false}; + Infra::kgm_ekey_db_t ekey_db_; + void WorkerThreadBody() { + SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_BELOW_NORMAL); + + std::unique_ptr task{nullptr}; + while ((task = Pop())) { + task->Execute(ekey_db_); + } + } + + std::mutex mutex_{}; + std::condition_variable signal_; + std::queue> tasks_{}; + std::queue threads_{}; +}; diff --git a/src/main.cpp b/src/main.cpp index 9faca2c..b3b57ed 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -1,3 +1,8 @@ +#include "infra/infra.h" +#include "infra/sqlite_error.h" +#include "jobs.hpp" +#include "qmc2/qmc2.h" + // clang-format off #include #include @@ -5,102 +10,41 @@ #include // clang-format on +#include #include #include -#include #include #include #include - -#include "infra/infra.h" -#include "infra/sqlite_error.h" -#include "qmc2/qmc2.h" +#include using Infra::kgm_ekey_db_t; -constexpr std::array kMagicHeader{0x7C, 0xD5, 0x32, 0xEB, 0x86, 0x02, 0x7F, 0x4B, - 0xA8, 0xAF, 0xA6, 0x8E, 0x0F, 0xFF, 0x99, 0x14}; +void WalkFileOrDir(KggTaskQueue& queue, const std::filesystem::path& input_path, bool scan_all_exts) { + std::queue file_queue; + file_queue.push(absolute(input_path)); -void DecryptKGG(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& kgg_file) { - 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"); - return; - } - uint32_t offset_to_audio = *reinterpret_cast(&header[0x10]); - uint32_t encrypt_mode = *reinterpret_cast(&header[0x14]); - if (encrypt_mode != 5) { - fprintf(stderr, "[WARN] ! invalid enc_version (expect=0x05, got 0x%02x), skip.\n", encrypt_mode); - return; - } - uint32_t audio_hash_len = *reinterpret_cast(&header[0x44]); - if (audio_hash_len != 0x20) { - 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]); - std::string ekey{}; - if (auto it = ekey_db.find(audio_hash); it != ekey_db.end()) { - ekey = it->second; - } else { - fprintf(stderr, "[WARN] ! ekey not found, skip.\n"); - return; - } + while (!file_queue.empty()) { + auto target_path = std::move(file_queue.front()); + file_queue.pop(); - auto qmc2 = QMC2::Create(ekey); - if (!qmc2) { - fprintf(stderr, "[WARN] ! create qmc2 failed, skip.\n"); - return; - } + if (is_regular_file(target_path)) { + if (!scan_all_exts && target_path.extension() != L".kgg") { + continue; + } - ifs_kgg.seekg(offset_to_audio, std::ios::beg); - std::vector buffer(1024 * 1024); - ifs_kgg.read(reinterpret_cast(buffer.data()), 4); - auto n = static_cast(ifs_kgg.gcount()); - qmc2->Decrypt(std::span(buffer.begin(), n), 0); - - auto decrypted_path = kgg_file; - auto magic = std::span(buffer.cbegin(), 4); - if (std::equal(magic.begin(), magic.end(), "fLaC")) { - decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.flac"); - } else if (std::equal(magic.begin(), magic.end(), "OggS")) { - decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.ogg"); - } else { - 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()); - return; - } - - std::ofstream ofs_decrypted(decrypted_path, std::ios::binary); - ofs_decrypted.write(reinterpret_cast(buffer.data()), n); - size_t offset{n}; - while (!ifs_kgg.eof()) { - ifs_kgg.read(reinterpret_cast(buffer.data()), buffer.size()); - n = static_cast(ifs_kgg.gcount()); - qmc2->Decrypt(std::span(buffer.begin(), n), offset); - ofs_decrypted.write(reinterpret_cast(buffer.data()), n); - offset += n; - } - qmc2.reset(); - - auto kgg_fname = kgg_file.filename(); - 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) { - if (is_directory(input_path)) { - for (auto const& dir_entry : std::filesystem::directory_iterator{input_path}) { - DecryptKGG(ekey_db, dir_entry.path()); + queue.Push(std::make_unique(target_path, target_path.parent_path())); + continue; } - } 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()); + + if (is_directory(target_path)) { + for (auto const& dir_entry : std::filesystem::directory_iterator{target_path}) { + file_queue.push(dir_entry.path()); + } + continue; + } + + fwprintf(stderr, L"[WARN] invalid path: %s\n", target_path.c_str()); } } @@ -144,6 +88,7 @@ int main() { fputs( "Usage: kgg-dec " "[--infra-dll infra.dll] " + "[--scan-all-file-ext 0] " "[--db /path/to/KGMusicV3.db] " "[--] [kgg-dir... = '.']\n\n", stderr); @@ -164,6 +109,11 @@ int main() { return 1; } + bool scan_all_exts{false}; + if (auto it = named_args.find(L"scan-all-file-ext"); it != named_args.end()) { + scan_all_exts = it->second == L"1"; + } + if (auto it = named_args.find(L"db"); it != named_args.end()) { kgm_db_path = std::filesystem::path{it->second}; } else { @@ -190,16 +140,29 @@ int main() { } db.Close(); -#ifdef _DEBUG - printf("ekey_db:\n"); +#ifndef NDEBUG + fprintf(stderr, "ekey_db:\n"); for (auto& [a, b] : ekey_db) { - printf("%s --> %s\n", a.c_str(), b.c_str()); + fprintf(stderr, "%s --> %s\n", a.c_str(), b.c_str()); } #endif - for (auto& positional_arg : positional_args) { - WalkFileOrDir(ekey_db, positional_arg); + KggTaskQueue queue(ekey_db); + auto thread_count = +#ifndef NDEBUG + 1; +#else + std::max(static_cast(std::thread::hardware_concurrency()) - 2, 2); +#endif + + for (int i = 0; i < thread_count; i++) { + queue.AddWorkerThread(); } + for (auto& positional_arg : positional_args) { + WalkFileOrDir(queue, positional_arg, scan_all_exts); + } + queue.Join(); + return 0; } diff --git a/src/qmc2/qmc2.h b/src/qmc2/qmc2.h index eb2fab3..928160b 100644 --- a/src/qmc2/qmc2.h +++ b/src/qmc2/qmc2.h @@ -11,10 +11,10 @@ namespace QMC2 { class QMC2_Base { public: - QMC2_Base(){}; + QMC2_Base() = default; virtual ~QMC2_Base() = default; - virtual void Decrypt(std::span data, size_t offset) = 0; + virtual void Decrypt(std::span data, size_t offset) const = 0; }; constexpr size_t kMapIndexOffset = 71214; @@ -24,7 +24,7 @@ class QMC2_MAP : public QMC2_Base { public: explicit QMC2_MAP(std::span key); - void Decrypt(std::span data, size_t offset) override; + void Decrypt(std::span data, size_t offset) const override; private: std::array key_{}; @@ -38,15 +38,15 @@ class QMC2_RC4 : public QMC2_Base { public: explicit QMC2_RC4(std::span key); - void Decrypt(std::span data, size_t offset) override; + void Decrypt(std::span data, size_t offset) const override; private: std::vector key_{}; double hash_{0}; std::array key_stream_{}; - size_t DecryptFirstSegment(std::span data, size_t offset); - size_t DecryptOtherSegment(std::span data, size_t offset); + [[nodiscard]] size_t DecryptFirstSegment(std::span data, size_t offset) const; + [[nodiscard]] size_t DecryptOtherSegment(std::span data, size_t offset) const; }; std::unique_ptr Create(std::string_view ekey); diff --git a/src/qmc2/qmc2_map.cpp b/src/qmc2/qmc2_map.cpp index d724543..0fe9fba 100644 --- a/src/qmc2/qmc2_map.cpp +++ b/src/qmc2/qmc2_map.cpp @@ -13,7 +13,7 @@ QMC2_MAP::QMC2_MAP(std::span key) { } } -void QMC2_MAP::Decrypt(std::span data, size_t offset) { +void QMC2_MAP::Decrypt(std::span data, size_t offset) const { for (auto& it : data) { size_t idx = (offset <= kMapOffsetBoundary) ? offset : (offset % kMapOffsetBoundary); it ^= key_[idx % key_.size()]; diff --git a/src/qmc2/qmc2_rc4.cpp b/src/qmc2/qmc2_rc4.cpp index 2138003..07d076c 100644 --- a/src/qmc2/qmc2_rc4.cpp +++ b/src/qmc2/qmc2_rc4.cpp @@ -14,7 +14,7 @@ inline double hash(const uint8_t* key, size_t len) { } // Overflow check. - uint32_t next_hash = hash * (uint32_t)(*key); + uint32_t next_hash = hash * static_cast(*key); if (next_hash <= hash) { break; } @@ -47,7 +47,7 @@ class RC4 { for (auto& it : buffer) { i = (i + 1) % n; - j = (j + (size_t)s[i]) % n; + j = (j + s[i]) % n; std::swap(s[i], s[j]); const size_t final_idx = (s[i] + s[j]) % n; @@ -70,8 +70,8 @@ inline size_t get_segment_key(double key_hash, size_t segment_id, uint8_t seed) return 0; } - double result = key_hash / (double)(seed * (segment_id + 1)) * 100.0; - return (size_t)result; + const double result = key_hash / static_cast(seed * (segment_id + 1)) * 100.0; + return static_cast(result); } QMC2_RC4::QMC2_RC4(std::span key) { @@ -82,7 +82,7 @@ QMC2_RC4::QMC2_RC4(std::span key) { rc4.Derive(std::span(key_stream_)); } -void QMC2_RC4::Decrypt(std::span data, size_t offset) { +void QMC2_RC4::Decrypt(std::span data, size_t offset) const { if (offset < kFirstSegmentSize) { const auto n = DecryptFirstSegment(data, offset); offset += n; @@ -96,7 +96,7 @@ void QMC2_RC4::Decrypt(std::span data, size_t offset) { } } -size_t QMC2_RC4::DecryptFirstSegment(std::span data, size_t offset) { +size_t QMC2_RC4::DecryptFirstSegment(std::span data, size_t offset) const { const uint8_t* key = key_.data(); const size_t n = this->key_.size(); @@ -109,7 +109,7 @@ size_t QMC2_RC4::DecryptFirstSegment(std::span data, size_t offset) { return process_len; } -size_t QMC2_RC4::DecryptOtherSegment(std::span data, size_t offset) { +size_t QMC2_RC4::DecryptOtherSegment(std::span data, size_t offset) const { const size_t n = this->key_.size(); size_t segment_idx = offset / kOtherSegmentSize;