core: Generalize QuickDecryptor

This is now a class template. This would be useful when we implement NCCH decryption.
This commit is contained in:
zhupengfei
2020-04-30 18:05:58 +08:00
parent 576053d995
commit c831a0785e
4 changed files with 73 additions and 73 deletions
+15 -4
View File
@@ -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<FileUtil::IOFile>(root_folder + source, "rb");
auto size = source_file->GetSize();
auto destination_file = std::make_unique<FileUtil::IOFile>(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() {
+2 -3
View File
@@ -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.
+37 -57
View File
@@ -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 <typename In, typename Out>
QuickDecryptor<In, Out>::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 <typename In, typename Out>
QuickDecryptor<In, Out>::~QuickDecryptor() = default;
QuickDecryptor::~QuickDecryptor() = default;
namespace {
std::array<u8, 16> GetFileCTR(const std::string& path) {
auto path_utf16 = Common::UTF8ToUTF16(path);
std::vector<u8> 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<u8, CryptoPP::SHA256::DIGESTSIZE> hash;
sha.CalculateDigest(hash.data(), path_data.data(), path_data.size());
std::array<u8, 16> 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 <typename In, typename Out>
bool QuickDecryptor<In, Out>::DecryptAndWriteFile(std::unique_ptr<In> source_, std::size_t size,
std::unique_ptr<Out> 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 <typename In, typename Out>
void QuickDecryptor<In, Out>::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 <typename In, typename Out>
void QuickDecryptor<In, Out>::DataDecryptLoop() {
CryptoPP::CTR_Mode<CryptoPP::AES>::Decryption aes;
aes.SetKeyWithIV(key.data(), key.size(), ctr.data());
@@ -156,11 +130,11 @@ void QuickDecryptor::DataDecryptLoop() {
}
}
void QuickDecryptor::DataWriteLoop() {
template <typename In, typename Out>
void QuickDecryptor<In, Out>::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 <typename In, typename Out>
void QuickDecryptor<In, Out>::Abort() {
if (is_running.exchange(false)) {
is_good = false;
completion_event.Set();
}
}
void QuickDecryptor::Reset(std::size_t total_size_) {
template <typename In, typename Out>
void QuickDecryptor<In, Out>::Reset(std::size_t total_size_) {
total_size = total_size_;
imported_size = 0;
}
template class QuickDecryptor<FileUtil::IOFile, FileUtil::IOFile>;
template class QuickDecryptor<SDMCFile, FileUtil::IOFile>;
} // namespace Core
+19 -9
View File
@@ -10,23 +10,32 @@
#include <string>
#include "common/common_types.h"
#include "common/thread.h"
#include "core/key/key.h"
namespace Core {
/// (current_size, total_size)
using ProgressCallback = std::function<void(std::size_t, std::size_t)>;
/**
* 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 <typename In = FileUtil::IOFile, typename Out = FileUtil::IOFile>
class QuickDecryptor {
public:
/// (current_size, total_size)
using ProgressCallback = std::function<void(std::size_t, std::size_t)>;
/**
* Initializes the decryptor.
* @param root_folder Path to the "Nintendo 3DS/<ID0>/<ID1>" 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<In> source, std::size_t size,
std::unique_ptr<Out> 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<In> source;
std::unique_ptr<Out> destination;
Core::Key::AESKey key;
Core::Key::AESKey ctr;
// Total size of this content, may consist of multiple files
std::size_t total_size{};