um/kgg-dec
6
6
Fork 0
Code Issues Pull Requests Releases 4 Wiki Activity

refactor: use threads!

Browse Source
This commit is contained in:
鲁树人 2024-10-18 08:11:34 +09:00
parent b23e7d0302
commit ac8c5d9a10
5 changed files with 236 additions and 103 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

170
src/jobs.hpp Normal file
View File

@ -0,0 +1,170 @@
#pragma once
#include <windows.h>
#include <filesystem>
#include <fstream>
#include <mutex>
#include <queue>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
#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<uint8_t, 16> 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<uint32_t*>(&header[0x10]);
uint32_t encrypt_mode = *reinterpret_cast<uint32_t*>(&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<uint32_t*>(&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<uint8_t*>(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<uint8_t, 1024 * 1024> temp_buffer{};
auto buf_signed = std::span(reinterpret_cast<char*>(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<size_t>(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<KggTask> task) {
std::lock_guard lock(mutex_);
tasks_.push(std::move(task));
signal_.notify_one();
}
std::unique_ptr<KggTask> 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<KggTask> task{nullptr};
while ((task = Pop())) {
task->Execute(ekey_db_);
}
}
std::mutex mutex_{};
std::condition_variable signal_;
std::queue<std::unique_ptr<KggTask>> tasks_{};
std::queue<std::thread> threads_{};
};

133
src/main.cpp
View File

@ -1,3 +1,8 @@
#include "infra/infra.h"
#include "infra/sqlite_error.h"
#include "jobs.hpp"
#include "qmc2/qmc2.h"
// clang-format off
#include <windows.h>
#include <shlobj.h>
@ -5,102 +10,41 @@
#include <knownfolders.h>
// clang-format on
#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <span>
#include "infra/infra.h"
#include "infra/sqlite_error.h"
#include "qmc2/qmc2.h"
#include <utility>
using Infra::kgm_ekey_db_t;
constexpr std::array<uint8_t, 16> 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<std::filesystem::path> 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<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);
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);
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();
if (is_regular_file(target_path)) {
if (!scan_all_exts && target_path.extension() != L".kgg") {
continue;
}
auto qmc2 = QMC2::Create(ekey);
if (!qmc2) {
fprintf(stderr, "[WARN] ! create qmc2 failed, skip.\n");
return;
queue.Push(std::make_unique<KggTask>(target_path, target_path.parent_path()));
continue;
}
ifs_kgg.seekg(offset_to_audio, std::ios::beg);
std::vector<uint8_t> buffer(1024 * 1024);
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);
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 (is_directory(target_path)) {
for (auto const& dir_entry : std::filesystem::directory_iterator{target_path}) {
file_queue.push(dir_entry.path());
}
if (exists(decrypted_path)) {
fprintf(stderr, "[WARN] ! output file '%s' exists, skip.\n", decrypted_path.filename().string().c_str());
return;
continue;
}
std::ofstream ofs_decrypted(decrypted_path, std::ios::binary);
ofs_decrypted.write(reinterpret_cast<char*>(buffer.data()), n);
size_t offset{n};
while (!ifs_kgg.eof()) {
ifs_kgg.read(reinterpret_cast<char*>(buffer.data()), buffer.size());
n = static_cast<size_t>(ifs_kgg.gcount());
qmc2->Decrypt(std::span(buffer.begin(), n), offset);
ofs_decrypted.write(reinterpret_cast<char*>(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());
}
} 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());
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<int>(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;
}

12
src/qmc2/qmc2.h
View File

@ -11,10 +11,10 @@ namespace QMC2 {
class QMC2_Base {
public:
QMC2_Base(){};
QMC2_Base() = default;
virtual ~QMC2_Base() = default;
virtual void Decrypt(std::span<uint8_t> data, size_t offset) = 0;
virtual void Decrypt(std::span<uint8_t> 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<uint8_t> key);
void Decrypt(std::span<uint8_t> data, size_t offset) override;
void Decrypt(std::span<uint8_t> data, size_t offset) const override;
private:
std::array<uint8_t, kMapKeySize> key_{};
@ -38,15 +38,15 @@ class QMC2_RC4 : public QMC2_Base {
public:
explicit QMC2_RC4(std::span<uint8_t> key);
void Decrypt(std::span<uint8_t> data, size_t offset) override;
void Decrypt(std::span<uint8_t> data, size_t offset) const override;
private:
std::vector<uint8_t> key_{};
double hash_{0};
std::array<uint8_t, kRC4StreamSize> key_stream_{};
size_t DecryptFirstSegment(std::span<uint8_t> data, size_t offset);
size_t DecryptOtherSegment(std::span<uint8_t> data, size_t offset);
[[nodiscard]] size_t DecryptFirstSegment(std::span<uint8_t> data, size_t offset) const;
[[nodiscard]] size_t DecryptOtherSegment(std::span<uint8_t> data, size_t offset) const;
};
std::unique_ptr<QMC2_Base> Create(std::string_view ekey);

2
src/qmc2/qmc2_map.cpp
View File

@ -13,7 +13,7 @@ QMC2_MAP::QMC2_MAP(std::span<uint8_t> key) {
}
}
void QMC2_MAP::Decrypt(std::span<uint8_t> data, size_t offset) {
void QMC2_MAP::Decrypt(std::span<uint8_t> data, size_t offset) const {
for (auto& it : data) {
size_t idx = (offset <= kMapOffsetBoundary) ? offset : (offset % kMapOffsetBoundary);
it ^= key_[idx % key_.size()];

14
src/qmc2/qmc2_rc4.cpp
View File

@ -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<uint32_t>(*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<double>(seed * (segment_id + 1)) * 100.0;
return static_cast<size_t>(result);
}
QMC2_RC4::QMC2_RC4(std::span<uint8_t> key) {
@ -82,7 +82,7 @@ QMC2_RC4::QMC2_RC4(std::span<uint8_t> key) {
rc4.Derive(std::span<uint8_t>(key_stream_));
}
void QMC2_RC4::Decrypt(std::span<uint8_t> data, size_t offset) {
void QMC2_RC4::Decrypt(std::span<uint8_t> 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<uint8_t> data, size_t offset) {
}
}
size_t QMC2_RC4::DecryptFirstSegment(std::span<uint8_t> data, size_t offset) {
size_t QMC2_RC4::DecryptFirstSegment(std::span<uint8_t> 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<uint8_t> data, size_t offset) {
return process_len;
}
size_t QMC2_RC4::DecryptOtherSegment(std::span<uint8_t> data, size_t offset) {
size_t QMC2_RC4::DecryptOtherSegment(std::span<uint8_t> data, size_t offset) const {
const size_t n = this->key_.size();
size_t segment_idx = offset / kOtherSegmentSize;