From c831a0785e7137bb8ced7715bd1e80f18cc5da25 Mon Sep 17 00:00:00 2001 From: zhupengfei Date: Thu, 30 Apr 2020 18:05:58 +0800 Subject: [PATCH] core: Generalize QuickDecryptor This is now a class template. This would be useful when we implement NCCH decryption. --- src/core/decryptor.cpp | 19 ++++++-- src/core/decryptor.h | 5 +- src/core/quick_decryptor.cpp | 94 ++++++++++++++---------------------- src/core/quick_decryptor.h | 28 +++++++---- 4 files changed, 73 insertions(+), 73 deletions(-) diff --git a/src/core/decryptor.cpp b/src/core/decryptor.cpp index 68c9fd3..103fc1e 100644 --- a/src/core/decryptor.cpp +++ b/src/core/decryptor.cpp @@ -17,8 +17,7 @@ namespace Core { -SDMCDecryptor::SDMCDecryptor(const std::string& root_folder_) - : root_folder(root_folder_), quick_decryptor(root_folder) { +SDMCDecryptor::SDMCDecryptor(const std::string& root_folder_) : root_folder(root_folder_) { ASSERT_MSG(Key::IsNormalKeyAvailable(Key::SDKey), "SD Key must be available in order to decrypt"); @@ -54,8 +53,20 @@ void SDMCDecryptor::Reset(std::size_t total_size) { } bool SDMCDecryptor::DecryptAndWriteFile(const std::string& source, const std::string& destination, - const QuickDecryptor::ProgressCallback& callback) { - return quick_decryptor.DecryptAndWriteFile(source, destination, callback); + const ProgressCallback& callback) { + if (!FileUtil::CreateFullPath(destination)) { + LOG_ERROR(Core, "Could not create path {}", destination); + return false; + } + + auto source_file = std::make_unique(root_folder + source, "rb"); + auto size = source_file->GetSize(); + auto destination_file = std::make_unique(destination, "wb"); + auto key = Key::GetNormalKey(Key::SDKey); + auto ctr = GetFileCTR(source); + return quick_decryptor.DecryptAndWriteFile(std::move(source_file), size, + std::move(destination_file), std::move(key), + std::move(ctr), callback); } void SDMCDecryptor::Abort() { diff --git a/src/core/decryptor.h b/src/core/decryptor.h index e1d05c1..2a88b19 100644 --- a/src/core/decryptor.h +++ b/src/core/decryptor.h @@ -32,8 +32,7 @@ public: * @return true on success, false otherwise */ bool DecryptAndWriteFile(const std::string& source, const std::string& destination, - const QuickDecryptor::ProgressCallback& callback = [](std::size_t, - std::size_t) {}); + const ProgressCallback& callback = [](std::size_t, std::size_t) {}); void Abort(); @@ -53,7 +52,7 @@ public: private: std::string root_folder; - QuickDecryptor quick_decryptor; + QuickDecryptor<> quick_decryptor; }; /// Interface for reading an SDMC file like a normal IOFile. This is read-only. diff --git a/src/core/quick_decryptor.cpp b/src/core/quick_decryptor.cpp index b975f17..bd993c7 100644 --- a/src/core/quick_decryptor.cpp +++ b/src/core/quick_decryptor.cpp @@ -11,54 +11,27 @@ #include "common/assert.h" #include "common/file_util.h" #include "common/string_util.h" -#include "core/key/key.h" +#include "core/decryptor.h" #include "core/quick_decryptor.h" namespace Core { -QuickDecryptor::QuickDecryptor(const std::string& root_folder_) : root_folder(root_folder_) { - ASSERT_MSG(Key::IsNormalKeyAvailable(Key::SDKey), - "SD Key must be available in order to decrypt"); +template +QuickDecryptor::QuickDecryptor() = default; - if (root_folder.back() == '/' || root_folder.back() == '\\') { - // Remove '/' or '\' character at the end as we will add them back when combining path - root_folder.erase(root_folder.size() - 1); - } -} +template +QuickDecryptor::~QuickDecryptor() = default; -QuickDecryptor::~QuickDecryptor() = default; - -namespace { -std::array GetFileCTR(const std::string& path) { - auto path_utf16 = Common::UTF8ToUTF16(path); - std::vector path_data(path_utf16.size() * 2 + 2, 0); // Add the '\0' character - std::memcpy(path_data.data(), path_utf16.data(), path_utf16.size() * 2); - - CryptoPP::SHA256 sha; - std::array hash; - sha.CalculateDigest(hash.data(), path_data.data(), path_data.size()); - - std::array ctr; - for (int i = 0; i < 16; i++) { - ctr[i] = hash[i] ^ hash[16 + i]; - } - return ctr; -} -} // namespace - -bool QuickDecryptor::DecryptAndWriteFile(const std::string& source_, - const std::string& destination_, - const ProgressCallback& callback_) { +template +bool QuickDecryptor::DecryptAndWriteFile(std::unique_ptr source_, std::size_t size, + std::unique_ptr destination_, + Core::Key::AESKey key_, Core::Key::AESKey ctr_, + const ProgressCallback& callback_) { if (is_running) { LOG_ERROR(Core, "Decryptor is running"); return false; } - if (!FileUtil::CreateFullPath(destination_)) { - LOG_ERROR(Core, "Could not create path {}", destination_); - return false; - } - for (auto& event : data_read_event) { event.Reset(); } @@ -70,15 +43,13 @@ bool QuickDecryptor::DecryptAndWriteFile(const std::string& source_, } completion_event.Reset(); - source = source_; - destination = destination_; + source = std::move(source_); + destination = std::move(destination_); + key = std::move(key_); + ctr = std::move(ctr_); callback = callback_; - current_total_size = FileUtil::GetSize(root_folder + source); - if (current_total_size == 0) { - LOG_ERROR(Core, "Could not open file {}", root_folder + source); - return false; - } + current_total_size = size; is_good = is_running = true; @@ -93,17 +64,21 @@ bool QuickDecryptor::DecryptAndWriteFile(const std::string& source_, decrypt_thread->join(); write_thread->join(); + // Release the files + source.reset(); + destination.reset(); + bool ret = is_good; is_good = true; return ret; } -void QuickDecryptor::DataReadLoop() { +template +void QuickDecryptor::DataReadLoop() { std::size_t current_buffer = 0; bool is_first_run = true; - FileUtil::IOFile file(root_folder + source, "rb"); - if (!file) { + if (!*source) { is_good = false; completion_event.Set(); return; @@ -121,7 +96,7 @@ void QuickDecryptor::DataReadLoop() { } const auto bytes_to_read = std::min(BufferSize, file_size); - if (file.ReadBytes(buffers[current_buffer].data(), bytes_to_read) != bytes_to_read) { + if (source->ReadBytes(buffers[current_buffer].data(), bytes_to_read) != bytes_to_read) { is_good = false; completion_event.Set(); return; @@ -133,9 +108,8 @@ void QuickDecryptor::DataReadLoop() { } } -void QuickDecryptor::DataDecryptLoop() { - auto ctr = GetFileCTR(source); - auto key = Key::GetNormalKey(Key::SDKey); +template +void QuickDecryptor::DataDecryptLoop() { CryptoPP::CTR_Mode::Decryption aes; aes.SetKeyWithIV(key.data(), key.size(), ctr.data()); @@ -156,11 +130,11 @@ void QuickDecryptor::DataDecryptLoop() { } } -void QuickDecryptor::DataWriteLoop() { +template +void QuickDecryptor::DataWriteLoop() { std::size_t current_buffer = 0; - FileUtil::IOFile file(destination, "wb"); - if (!file) { + if (!*destination) { is_good = false; completion_event.Set(); return; @@ -181,7 +155,8 @@ void QuickDecryptor::DataWriteLoop() { data_decrypted_event[current_buffer].Wait(); const auto bytes_to_write = std::min(BufferSize, file_size); - if (file.WriteBytes(buffers[current_buffer].data(), bytes_to_write) != bytes_to_write) { + if (destination->WriteBytes(buffers[current_buffer].data(), bytes_to_write) != + bytes_to_write) { is_good = false; completion_event.Set(); return; @@ -196,16 +171,21 @@ void QuickDecryptor::DataWriteLoop() { completion_event.Set(); } -void QuickDecryptor::Abort() { +template +void QuickDecryptor::Abort() { if (is_running.exchange(false)) { is_good = false; completion_event.Set(); } } -void QuickDecryptor::Reset(std::size_t total_size_) { +template +void QuickDecryptor::Reset(std::size_t total_size_) { total_size = total_size_; imported_size = 0; } +template class QuickDecryptor; +template class QuickDecryptor; + } // namespace Core diff --git a/src/core/quick_decryptor.h b/src/core/quick_decryptor.h index d82977a..5f7735d 100644 --- a/src/core/quick_decryptor.h +++ b/src/core/quick_decryptor.h @@ -10,23 +10,32 @@ #include #include "common/common_types.h" #include "common/thread.h" +#include "core/key/key.h" namespace Core { +/// (current_size, total_size) +using ProgressCallback = std::function; + +/** + * Helper that reads, decrypts and writes data. This uses three threads to process the data + * and call progress callbacks occasionally. + * + * While this is a template, it really should only be used with IOFile and SDMCFile. + */ +template class QuickDecryptor { public: - /// (current_size, total_size) - using ProgressCallback = std::function; - /** * Initializes the decryptor. * @param root_folder Path to the "Nintendo 3DS//" folder. */ - explicit QuickDecryptor(const std::string& root_folder); - + explicit QuickDecryptor(); ~QuickDecryptor(); - bool DecryptAndWriteFile(const std::string& source, const std::string& destination, + bool DecryptAndWriteFile(std::unique_ptr source, std::size_t size, + std::unique_ptr destination, Core::Key::AESKey key, + Core::Key::AESKey ctr, const ProgressCallback& callback = [](std::size_t, std::size_t) {}); void DataReadLoop(); @@ -41,9 +50,10 @@ public: private: static constexpr std::size_t BufferSize = 16 * 1024; // 16 KB - std::string root_folder; - std::string source; - std::string destination; + std::unique_ptr source; + std::unique_ptr destination; + Core::Key::AESKey key; + Core::Key::AESKey ctr; // Total size of this content, may consist of multiple files std::size_t total_size{};