diff --git a/src/crypto.cpp b/src/crypto.cpp index d125e1e..58fd701 100644 --- a/src/crypto.cpp +++ b/src/crypto.cpp @@ -12,29 +12,11 @@ #include "utils.h" namespace horcrux { -/** Cipher context class - * This is a wrapper for OpenSSL context object, just for RAII pattern - */ -class EvpCipherCtx { - EVP_CIPHER_CTX *ptr; -public: - EvpCipherCtx() { - ptr = EVP_CIPHER_CTX_new(); - } - - ~EvpCipherCtx() { - EVP_CIPHER_CTX_free(ptr); - } - - EVP_CIPHER_CTX* get() { - return ptr; - } -}; - - -EvpCipherCtx AES256_CBC::init(Mode mode, const raw_data& key, +void AES256_CBC::init(Mode new_mode, const raw_data& key, const raw_data& iv) { + // force mode change + mode = new_mode; if (key.size() != kKeySize) { throw std::runtime_error("Wrong key size"); } @@ -42,24 +24,54 @@ EvpCipherCtx AES256_CBC::init(Mode mode, const raw_data& key, throw std::runtime_error("Wrong IV size"); } - EvpCipherCtx ctx; + ctx = std::make_unique(); if (mode == Mode::kEncrypt) { - if (EVP_EncryptInit(ctx.get(), EVP_aes_256_cbc(), + if (EVP_EncryptInit(ctx->get(), EVP_aes_256_cbc(), key.data(), iv.data()) == 0) { throw std::runtime_error("EVP_EncryptInit"); } } else if (mode == Mode::kDecrypt) { - if (EVP_DecryptInit(ctx.get(), EVP_aes_256_cbc(), + if (EVP_DecryptInit(ctx->get(), EVP_aes_256_cbc(), key.data(), iv.data()) == 0) { throw std::runtime_error("EVP_DecryptInit"); } } else { throw std::invalid_argument("Invalid Cipher mode"); } - return ctx; } -size_t AES256_CBC::process_chunk(Mode mode, EvpCipherCtx& ctx, +size_t AES256_CBC::process_start(Mode mode, + raw_data::const_iterator begin, + raw_data::const_iterator end, + raw_data& output, + bool resize_in, bool resize_out) { + auto current = begin; + size_t output_offset = 0; + raw_data iv(kIvSize); + if (mode == Mode::kEncrypt) { + // Make sure output is large enough to add IV + output.resize(kIvSize); + // generate IV + iv = generate_random(kIvSize); + // write IV to output + std::copy(iv.begin(), iv.end(), output.begin()); + output_offset += kIvSize; + } else { + // kDecrypt + if (end - begin < kIvSize){ + throw std::invalid_argument( + "First encrypted chunk must contain the whole IV"); + } + // read IV from input + std::copy(current, current + kIvSize, iv.begin()); + current += kIvSize; + } + init(mode, encryption_key, iv); + return output_offset + process_chunk(mode, current, end, output, + output_offset, resize_in, resize_out); +} + +size_t AES256_CBC::process_chunk(Mode mode, raw_data::const_iterator begin, raw_data::const_iterator end, raw_data& output, size_t output_offset, @@ -71,13 +83,13 @@ size_t AES256_CBC::process_chunk(Mode mode, EvpCipherCtx& ctx, output.resize(output_offset + chunk_size + kIvSize); } if (mode == Mode::kEncrypt) { - if (1 != EVP_EncryptUpdate(ctx.get(), + if (1 != EVP_EncryptUpdate(ctx->get(), output.data() + output_offset, &len, &*begin, chunk_size)) { throw std::runtime_error("EVP_EncryptUpdate"); } } else { - if (1 != EVP_DecryptUpdate(ctx.get(), + if (1 != EVP_DecryptUpdate(ctx->get(), output.data() + output_offset, &len, &*begin, chunk_size)) { throw std::runtime_error("EVP_DecryptUpdate"); @@ -88,7 +100,7 @@ size_t AES256_CBC::process_chunk(Mode mode, EvpCipherCtx& ctx, return len; } -size_t AES256_CBC::process_final(Mode mode, EvpCipherCtx& ctx, +size_t AES256_CBC::process_final(Mode mode, raw_data& output, size_t output_offset, bool resize_in, bool resize_out) { int len; @@ -97,12 +109,12 @@ size_t AES256_CBC::process_final(Mode mode, EvpCipherCtx& ctx, output.resize(output_offset + kIvSize); } if (mode == Mode::kEncrypt) { - if (1 != EVP_EncryptFinal_ex(ctx.get(), output.data() + output_offset, + if (1 != EVP_EncryptFinal_ex(ctx->get(), output.data() + output_offset, &len)) { throw std::runtime_error("EVP_EncryptFinal"); } } else { - if (1 != EVP_DecryptFinal_ex(ctx.get(), (output.data()) + output_offset, + if (1 != EVP_DecryptFinal_ex(ctx->get(), (output.data()) + output_offset, &len)) { throw std::runtime_error("EVP_DecryptFinal"); } @@ -112,62 +124,57 @@ size_t AES256_CBC::process_final(Mode mode, EvpCipherCtx& ctx, return len; } -size_t AES256_CBC::process_all(Mode mode, EvpCipherCtx& ctx, +/* this function does not use process_start, hence it does not handle the IV */ +size_t AES256_CBC::process_all(Mode mode, raw_data::const_iterator begin, raw_data::const_iterator end, raw_data& output, size_t output_offset, bool resize_in, bool resize_out) { - int len = process_chunk(mode, ctx, begin, end, output, output_offset, + int len = process_chunk(mode, begin, end, output, output_offset, resize_in, false); - len += process_final(mode, ctx, output, output_offset + len, false, + len += process_final(mode, output, output_offset + len, false, resize_out); return len; } -raw_data AES256_CBC::process(const raw_data& inputdata) { - if (mode == Mode::kEncrypt) - return encrypt(inputdata); +raw_data AES256_CBC::process(const raw_data& inputdata, const Flags& flags) { + raw_data output; + size_t len; + if (flags.to_ulong() & flag::kBegin) + len = process_start(mode, inputdata.begin(), inputdata.end(), output); else - return decrypt(inputdata); + len = process_chunk(mode, inputdata.begin(), inputdata.end(), output, 0); + if (flags.to_ulong() & flag::kEnd) + process_final(mode, output, len); + return output; } -raw_data AES256_CBC::encrypt(const raw_data& plaintext) { - return encrypt(encryption_key, horcrux::generate_random(kIvSize), - plaintext); +raw_data AES256_CBC::encrypt(const raw_data& plaintext, + const std::bitset& flags) { + mode = Mode::kEncrypt; + return process(plaintext, flags); } -raw_data AES256_CBC::decrypt(const raw_data& ciphertext) { - return decrypt(encryption_key, ciphertext); + +raw_data AES256_CBC::decrypt(const raw_data& ciphertext, + const std::bitset& flags) { + mode = Mode::kDecrypt; + return process(ciphertext, flags); } +/* OVERLOADED FUNCTIONS FOR TEST PURPOSE */ + raw_data AES256_CBC::encrypt(const raw_data& key, const raw_data& iv, const raw_data& input) { - auto ctx = init(Mode::kEncrypt, key, iv); + init(Mode::kEncrypt, key, iv); // Make sure ouput is large enough to contain IV + encrypted data + padding raw_data output(input.size() + (2 * kIvSize)); std::copy(iv.begin(), iv.end(), output.begin()); - process_all(Mode::kEncrypt, ctx, input.begin(), input.end(), output, + process_all(Mode::kEncrypt, input.begin(), input.end(), output, kIvSize); return output; } -/* - -int AES256_CBC::encrypt(const raw_data& key, const raw_data& iv, - std::istream input, const size_t input_len, - std::ostream output, size_t& output_len) { - auto inbuf = raw_data(input_len); - auto outbuf = raw_data(input_len + 16); - input.read(reinterpret_cast(inbuf.data()), input_len); - - encrypt(key, iv, inbuf, outbuf); - - output.write(reinterpret_cast(outbuf.data()), outbuf.size()); - - return 0; -} - -*/ raw_data AES256_CBC::decrypt(const raw_data& key, const raw_data& ciphertext) { raw_data iv(ciphertext.begin(), ciphertext.begin() + kIvSize); return decrypt(key, iv, ciphertext.begin() + kIvSize, ciphertext.end()); @@ -181,9 +188,9 @@ raw_data AES256_CBC::decrypt(const raw_data& key, const raw_data& iv, raw_data AES256_CBC::decrypt(const raw_data& key, const raw_data& iv, raw_data::const_iterator begin, raw_data::const_iterator end) { - auto ctx = init(Mode::kDecrypt, key, iv); raw_data output; - process_all(Mode::kDecrypt, ctx, begin, end, output, 0); + init(Mode::kDecrypt, key, iv); + process_all(Mode::kDecrypt, begin, end, output, 0); return output; } }; // namespace horcrux diff --git a/src/crypto.h b/src/crypto.h index 97a5336..b623785 100644 --- a/src/crypto.h +++ b/src/crypto.h @@ -1,8 +1,10 @@ #ifndef HORCRUX_SRC_CRYPTO_H #define HORCRUX_SRC_CRYPTO_H +#include #include #include #include +#include #include "utils.h" @@ -15,6 +17,10 @@ enum class Mode { kEncrypt, kDecrypt }; */ class Cipher { public: + enum flag { kBegin = 0x01, kEnd = 0x02}; + constexpr static int Flags_num = 2; + typedef std::bitset Flags; + /** @brief default constructor */ Cipher() = delete; @@ -34,10 +40,14 @@ public: protected: /** @brief key used for encryption / decryption */ raw_data encryption_key; + /** @brief current cipher mode */ + Mode mode; public: - /** @brief Cipher mode (encrypt or decrypt) */ - const Mode mode; + /** @brief get current Cipher mode (encrypt or decrypt) */ + const Mode& get_mode() { + return mode; + } /** @brief get Cipher encryption key * @return the encryption key @@ -47,29 +57,61 @@ public: } /** @brief process inputdata according to Cipher mode + * @details + * set kBegin flag to specify this is the first chunk to process. + * set kEnd flag to specify this is the last chunk to process. * @param inputdata buffer containing data to process + * @param flags flag bitset. Default (kBegin | kEnd) * @return processed data */ - virtual raw_data process( - const raw_data& inputdata) = 0; + virtual raw_data process(const raw_data& inputdata, + const std::bitset& flags = + Flags(kBegin | kEnd)) = 0; /** @brief encrypt the content of a buffer using Cipher key + * @details + * set kBegin flag to specify this is the first chunk to process. + * set kEnd flag to specify this is the last chunk to process. * @param plaintext the input buffer + * @param flags flag bitset. Default (kBegin | kEnd) * @return encrypted data */ - virtual raw_data encrypt( - const raw_data& plaintext) = 0; + virtual raw_data encrypt(const raw_data& plaintext, + const std::bitset& flags = + Flags(kBegin | kEnd)) = 0; /** @brief decrypt the content of a buffer using Cipher key + * @details + * set kBegin flag to specify this is the first chunk to process. + * set kEnd flag to specify this is the last chunk to process. * @param decrypt the input buffer + * @param flags flag bitset. Default (kBegin | kEnd) * @return decrypted data */ - virtual raw_data decrypt( - const raw_data& ciphertext) = 0; + virtual raw_data decrypt(const raw_data& ciphertext, + const std::bitset& flags = + Flags(kBegin | kEnd)) = 0; }; -// forward declaration -class EvpCipherCtx; +/** Cipher context class + * This is a wrapper for OpenSSL context object, just for RAII pattern + */ +class EvpCipherCtx { + EVP_CIPHER_CTX *ptr; + +public: + EvpCipherCtx() { + ptr = EVP_CIPHER_CTX_new(); + } + + ~EvpCipherCtx() { + EVP_CIPHER_CTX_free(ptr); + } + + EVP_CIPHER_CTX* get() { + return ptr; + } +}; /** @brief AES256 Cipher * @details @@ -89,6 +131,8 @@ class AES256_CBC : public Cipher { const size_t kKeySize = 32; /** AES256 CBC initialization vector size, from AES256 specs */ const size_t kIvSize = 16; + /** pointer to Cipher context */ + std::unique_ptr ctx = nullptr; public: /** @brief Create cipher in encrypt mode with random key */ @@ -107,93 +151,68 @@ public: virtual ~AES256_CBC() = default; /** @brief process inputdata according to Cipher mode + * @details + * if kBegin flag is set, this function will handle the initial IV. + * if kEnd flag is set, this function will handle the final padding. * @param inputdata buffer containing data to process + * @param flags flag bitset. Default (kBegin | kEnd) * @return processed data */ - raw_data process( - const raw_data& inputdata) override; + raw_data process(const raw_data& inputdata, + const std::bitset& flags = + Flags(kBegin | kEnd)) override; /** @brief encrypt the content of a buffer using Cipher key + * @details + * if kBegin flag is set, this function will handle the initial IV. + * if kEnd flag is set, this function will handle the final padding. * @param plaintext the input buffer + * @param flags flag bitset. Default (kBegin | kEnd) * @return encrypted data */ - raw_data encrypt( - const raw_data& plaintext) override; + raw_data encrypt(const raw_data& plaintext, + const std::bitset& flags = + Flags(kBegin | kEnd)) override; /** @brief decrypt the content of a buffer using Cipher key + * @details + * if kBegin flag is set, this function will handle the initial IV. + * if kEnd flag is set, this function will handle the final padding. * @param decrypt the input buffer + * @param flags flag bitset. Default (kBegin | kEnd) * @return decrypted data */ - raw_data decrypt( - const raw_data& ciphertext) override; - - /** @brief Read the whole input and return a buffer with the encrypted data - * @param key the encryption key - * @param iv the initialization vector that will be copied to the beginning - * of the output - * @param input the input buffer - * @return the encrypted data buffer - */ - raw_data encrypt(const raw_data& key, - const raw_data& iv, - const raw_data& input); - -/* - int encrypt(const raw_data& key, - const raw_data& iv, - std::istream input, const size_t input_len, - std::ostream output, size_t& output_len); - -*/ - /** @brief Read the whole input and return a buffer with the decrypted data. - * Read the iv from the beginning of ciphertext - * @param key the encryption key - * @param ciphertext the input buffer - * @return the encrypted data buffer - */ - raw_data decrypt( - const raw_data& key, - const raw_data& ciphertext); - - /** @brief Read the whole input and return a buffer with the decrypted data. - * ciphertext does not contain the iv. - * @param key the encryption key - * @param iv the initialization vector - * @param ciphertext the input buffer - * @return the encrypted data buffer - */ - raw_data decrypt( - const raw_data& key, - const raw_data& iv, - const raw_data& ciphertext); - - /** @brief Read some input from begin to end and return a buffer with the - * decrypted data. Output data is finalized. - * @param key the encryption key - * @param iv the initialization vector - * @param begin beginning of the input buffer - * @param end end of the input buffer - * @return the encrypted data buffer - */ - raw_data decrypt( - const raw_data& key, - const raw_data& iv, - raw_data::const_iterator begin, - raw_data::const_iterator end); + raw_data decrypt(const raw_data& ciphertext, + const std::bitset& flags = + Flags(kBegin | kEnd)) override; private: - /** @brief Create a Cipher context object needed for enc/dec operations + /** @brief Initialize the Cipher context object needed for enc/dec operations + * @details this operation will reset the cipher, any pending operation will + * be lost. * @param mode encrypt or decrypt * @param key the enc/dec key * @param the initialization vector - * @return the Cipher context object */ - EvpCipherCtx init(Mode mode, const raw_data& key, - const raw_data& iv); + void init(Mode mode, const raw_data& key, const raw_data& iv); + + /** @brief process the first chunk of input, it handles the IV. + * @param mode encrypt or decrypt + * @param begin beginning of the input chunk + * @param end end of the input chunk + * @param output output buffer, it will be overwritten from the beginning + * @param resize_in resize the output buffer before appending data to it in + * order to make sure there is enough room. + * @param resize_out shrink the output buffer at the end + */ + size_t process_start(Mode mode, + raw_data::const_iterator begin, + raw_data::const_iterator end, + raw_data& output, + bool resize_in = true, bool resize_out = true); /** @brief process a chunk of input * @param mode encrypt or decrypt - * @param ctx the Cipher context object * @param begin beginning of the input chunk * @param end end of the input chunk * @param output output buffer @@ -203,7 +222,7 @@ private: * order to make sure there is enough room. * @param resize_out shrink the output buffer at the end */ - size_t process_chunk(Mode mode, EvpCipherCtx& ctx, + size_t process_chunk(Mode mode, raw_data::const_iterator begin, raw_data::const_iterator end, raw_data& output, @@ -213,7 +232,6 @@ private: /** @brief Finalize a encryption/decryption process. Write remaining bytes * to the output * @param mode encrypt or decrypt - * @param ctx the Cipher context object * @param output output buffer * @param output_offset where to start to append the last data within * the output buffer @@ -221,14 +239,14 @@ private: * order to make sure there is enough room. * @param resize_out shrink the output buffer at the end */ - size_t process_final(Mode mode, EvpCipherCtx& ctx, + size_t process_final(Mode mode, raw_data& output, size_t output_offset, bool resize_in = true, bool resize_out = true); /** @brief process a input and finalize it. + * @node this function does not handle the IV * @param mode encrypt or decrypt - * @param ctx the Cipher context object * @param begin beginning of the input chunk * @param end end of the input chunk * @param output output buffer @@ -238,11 +256,55 @@ private: * order to make sure there is enough room. * @param resize_out shrink the output buffer at the end */ - size_t process_all(Mode mode, EvpCipherCtx& ctx, + size_t process_all(Mode mode, raw_data::const_iterator begin, raw_data::const_iterator end, raw_data& output, size_t output_offset, bool resize_in = true, bool resize_out = true); + +public: + /* OVERLOADED FUNCTIONS FOR TEST PURPOSE */ + + /** @brief Read the whole input and return a buffer with the encrypted data + * @param key the encryption key + * @param iv the initialization vector that will be copied to the beginning + * of the output + * @param input the input buffer + * @return the encrypted data buffer + */ + raw_data encrypt(const raw_data& key, const raw_data& iv, + const raw_data& input); + + /** @brief Read the whole input and return a buffer with the decrypted data. + * Read the iv from the beginning of ciphertext + * @param key the encryption key + * @param ciphertext the input buffer + * @return the encrypted data buffer + */ + raw_data decrypt(const raw_data& key, const raw_data& ciphertext); + + /** @brief Read the whole input and return a buffer with the decrypted data. + * ciphertext does not contain the iv. + * @param key the encryption key + * @param iv the initialization vector + * @param ciphertext the input buffer + * @return the encrypted data buffer + */ + raw_data decrypt(const raw_data& key, const raw_data& iv, + const raw_data& ciphertext); + + /** @brief Read some input from begin to end and return a buffer with the + * decrypted data. Output data is finalized. + * @param key the encryption key + * @param iv the initialization vector + * @param begin beginning of the input buffer + * @param end end of the input buffer + * @return the encrypted data buffer + */ + raw_data decrypt(const raw_data& key, const raw_data& iv, + raw_data::const_iterator begin, + raw_data::const_iterator end); + }; }; // namespace horcrux diff --git a/src/horcrux.cpp b/src/horcrux.cpp index 0a64cf2..fefcdc4 100644 --- a/src/horcrux.cpp +++ b/src/horcrux.cpp @@ -5,16 +5,30 @@ namespace horcrux { +void Horcrux::process_chunks(size_t size) { + raw_data buf = input->read(size); + //std::cout << "read " << buf.size() << " " << size << std::endl; + output->write_chunk(cipher->process(buf, + Cipher::Flags(Cipher::flag::kBegin))); + while ((buf = input->read(size)).size() == size) { + //std::cout << "read " << buf.size() << std::endl; + output->write_chunk(cipher->process(buf, Cipher::Flags())); + } + output->write_chunk(cipher->process(buf, Cipher::Flags(Cipher::flag::kEnd))); +} + void Horcrux::init() { if (options.count <= 0) { - throw std::invalid_argument("Invalid horcrux count"); + throw std::invalid_argument("init: Invalid horcrux count"); } switch (options.mode) { case Mode::kEncrypt: cipher = std::make_unique(); input = std::make_unique(options.input[0]); output = std::make_unique(options.output, options.count, - std::filesystem::path(options.input[0]).filename().string()); + std::filesystem::path(options.input[0]).filename().string(), + static_cast(input.get())->get_file_size() / + options.count); options.base64_key = to_base64(cipher->get_encryption_key()); break; case Mode::kDecrypt: @@ -30,11 +44,23 @@ void Horcrux::init() { } } -void Horcrux::run() { - if (options.mode == Mode::kEncrypt) - output->write(cipher->encrypt(input->read())); - else - output->write(cipher->decrypt(input->read())); +void Horcrux::run(size_t buffer_size) { + if (buffer_size == 0){ + // read the whole input and process it. It may use a lot of memory. + output->write(cipher->process(input->read())); + } else { + // read and process the input buffer_size bytes at a time + process_chunks(buffer_size); + } + /* + if (options.mode == Mode::kEncrypt){ + //output->write(cipher->encrypt(input->read())); + read_chunks(buffer_size); + } else { + //output->write(cipher->decrypt(input->read())); + read_chunks(buffer_size); + } + */ } static void print_usage(const std::string& program) { diff --git a/src/horcrux.h b/src/horcrux.h index dfe9261..1cbd551 100644 --- a/src/horcrux.h +++ b/src/horcrux.h @@ -39,6 +39,12 @@ class Horcrux { * can be selected in this function */ void init(); + /** @brief read from Input, process with Cipher and write to Output + * one chunk at a time + * @ size chunk buffer size. + */ + void process_chunks(size_t size); + public: /** remove default constructor */ Horcrux() = delete; @@ -52,7 +58,7 @@ public: /** @brief execute the horcrux process * @details read from Input, process with Cipher and write to output */ - void run(); + void run(size_t size = 0); }; /** @brief utility function to parse the command line arguments * @param arguments arguments as read from command line diff --git a/src/io.cpp b/src/io.cpp index 7386735..f874cd0 100644 --- a/src/io.cpp +++ b/src/io.cpp @@ -14,12 +14,17 @@ FsPlainInput::FsPlainInput(const std::string& path) if (file_size == 0) { throw std::invalid_argument("Input is empty, nothing to encrypt"); } - file_stream = std::ifstream(file_path, std::ios::binary); } raw_data FsPlainInput::read() { - raw_data result(file_size); - file_stream.read(reinterpret_cast(result.data()), file_size); + return read(fs::file_size(file_path)); +} +raw_data FsPlainInput::read(size_t size) { + raw_data result(size); + if (!file_stream.is_open()) + file_stream.open(file_path, std::ios::binary); + file_stream.read(reinterpret_cast(result.data()), size); + result.resize(file_stream.gcount()); return result; } @@ -41,13 +46,28 @@ FsCryptoInput::FsCryptoInput(const std::vector& filenames) { } raw_data FsCryptoInput::read() { - raw_data result(total_size); + // start from the beginning + current_file = 0; + file_stream.close(); + return read(total_size); +} + +raw_data FsCryptoInput::read(size_t size) { + raw_data result(size); size_t data_read{0}; - for (auto& f : file_paths) { - std::ifstream ifstream(f.first); - ifstream.read(reinterpret_cast(result.data()) + data_read, f.second); - data_read += f.second; + while (data_read < size && current_file < file_paths.size()) { + if (!file_stream.is_open()) + file_stream.open(file_paths[current_file].first, std::ios::binary); + file_stream.read( + reinterpret_cast(result.data()) + data_read, + size - data_read); + data_read += file_stream.gcount(); + if(file_stream.eof()) { + file_stream.close(); + ++current_file; + } } + result.resize(data_read); return result; } @@ -58,9 +78,16 @@ FsPlainOutput::FsPlainOutput(const std::string& filename) throw std::invalid_argument("Output file is not a regular file"); } } + size_t FsPlainOutput::write(const raw_data& to_write) { - std::ofstream out(file_path, std::ios::binary); - out.write(reinterpret_cast(to_write.data()), to_write.size()); + return write_chunk(to_write); +} + +size_t FsPlainOutput::write_chunk(const raw_data& to_write) { + if(!file_stream.is_open()) + file_stream.open(file_path, std::ios::binary); + file_stream.write(reinterpret_cast(to_write.data()), + to_write.size()); return to_write.size(); } @@ -79,18 +106,40 @@ FsCryptoOutput::FsCryptoOutput(const std::string& folder, const int horcrux_num, size_t FsCryptoOutput::write(const raw_data& to_write) { size = to_write.size() / num; - std::ofstream f; created_files.clear(); + return write_chunk(to_write); +} + +size_t FsCryptoOutput::write_chunk(const raw_data& to_write) { + // horcrux file size must be known in advance + if (size <= 0) { + throw std::invalid_argument("Invalid horcrux size"); + } size_t data_written{0}; - for (int i = 0; i < num; ++i) { - std::string name = base_name + '.' + std::to_string(i) + ".enc"; - created_files.emplace_back(folder_path / name); - f = std::ofstream(created_files.back(), - std::ios::binary | std::ios::trunc); - auto chunk_size = i + 1 != num ? size : size + to_write.size() % size; - f.write(reinterpret_cast(to_write.data()) + data_written, - chunk_size); - data_written += chunk_size; + while (data_written < to_write.size()){ + if(!file_stream.is_open()){ + // create a new file + std::string name = base_name + + '.' + std::to_string(created_files.size()) + ".enc"; + created_files.emplace_back(folder_path / name); + file_stream.open(created_files.back(), + std::ios::binary | std::ios::trunc); + } + // write up to fill the file size or consume the input buffer + auto writable = std::min(size - file_stream.tellp(), + to_write.size() - data_written); + // the last file may contain more data (size % num) + if (num == created_files.size()) { + writable = to_write.size() - data_written; + } + file_stream.write( + reinterpret_cast(to_write.data()) + data_written, + writable); + data_written += writable; + if (file_stream.tellp() >= size){ + // close the file if it has been filled + file_stream.close(); + } } return data_written; } diff --git a/src/io.h b/src/io.h index 5d35422..8f151d6 100644 --- a/src/io.h +++ b/src/io.h @@ -14,9 +14,15 @@ namespace horcrux { class Input { public: virtual ~Input() = default; - /** @brief Read from the input + /** @brief Read and consume the whole input * @return new buffer containing all read data */ virtual raw_data read() = 0; + + /** @brief Read from the input + * @param size the amount of data to extract from the input + * @return a new buffer containing at most size bytes */ + virtual raw_data read(size_t size) = 0; + }; /** @brief Base Class for Output @@ -24,10 +30,16 @@ public: class Output { public: virtual ~Output() = default; - /** @brief Write to the output - * @param data to write + + /** @brief Write everything to the output + * @param to_write the one and only data buffer to write * @return written data */ virtual size_t write(const raw_data& data) = 0; + + /** @brief Write single chunks to the output + * @param to_write the data chunk to write + * @return written data */ + virtual size_t write_chunk(const raw_data& data) = 0; }; /** @brief Input that reads data from a regular file */ @@ -48,9 +60,18 @@ public: * @param path input file path, must be a regular file */ explicit FsPlainInput(const std::string& path); - /** @brief Read from the input + /** @brief get the size of the underlying input file + * @return the file size */ + const size_t& get_file_size() { return file_size; }; + + /** @brief Read and consume the whole input * @return new buffer containing all read data */ raw_data read() override; + + /** @brief Read from the input + * @param size the amount of data to extract from the input + * @return a new buffer containing at most size bytes */ + raw_data read(size_t size) override; }; /** @brief Input that reads scattered data from multiple files */ @@ -59,6 +80,10 @@ class FsCryptoInput : public Input{ std::vector> file_paths; /// total size of input files size_t total_size{0}; + /// input file stream associated to the current input file + std::ifstream file_stream; + /// iterator pointing to the current input file + int current_file{0}; /// remove default constuctor FsCryptoInput() = delete; @@ -69,16 +94,23 @@ public: * @param filenames paths of input files. They will be read in this order */ explicit FsCryptoInput(const std::vector& filenames); - /** @brief Read from the input + /** @brief Read and consume the whole input * @details basically concatenate the input files. * @return new buffer containing all read data */ raw_data read() override; + + /** @brief Read from the input + * @param size the amount of data to extract from the input + * @return a new buffer containing at most size bytes */ + raw_data read(size_t size) override; }; /** @brief Output that writes to a single regular file */ class FsPlainOutput : public Output { /// output file path std::filesystem::path file_path; + /// output file stream associated to the output file + std::ofstream file_stream; /// remove default constructor FsPlainOutput() = delete; @@ -89,10 +121,15 @@ public: * @param filename output file. It can be a new file, it must be writable */ explicit FsPlainOutput(const std::string& filename); - /** @brief Write to the output - * @param data to write + /** @brief Write everything to the output + * @param to_write the one and only data buffer to write * @return written data */ size_t write(const raw_data& to_write) override; + + /** @brief Write single chunks to the output + * @param to_write the data chunk to write + * @return written data */ + size_t write_chunk(const raw_data& to_write) override; }; /** @brief Output that splits the data in equal size chunks and writes them to @@ -108,6 +145,8 @@ class FsCryptoOutput : public Output { int num; /// size of each ouptut file size_t size; + /// output file stream associated to the current output file + std::ofstream file_stream; /// remove default constructor FsCryptoOutput() = delete; @@ -126,10 +165,15 @@ public: const std::string& filename = "horcrux", const size_t horcrux_size = 0); - /** @brief Write to the output - * @param data to write + /** @brief Write everything to the output + * @param to_write the one and only data buffer to write * @return written data */ size_t write(const raw_data& to_write) override; + + /** @brief Write single chunks to the output + * @param to_write the data chunk to write + * @return written data */ + size_t write_chunk(const raw_data& to_write) override; }; }; // namespace horcrux #endif // HORCRUX_SRC_IO_H diff --git a/src/main.cpp b/src/main.cpp index d0287fb..602987e 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -2,6 +2,9 @@ #include #include "horcrux.h" +/* arbitrary buffer size */ +constexpr size_t buffer_size = 1024; + int main(const int argc, const char *argv[]) { int ret = -1; std::vector arguments; @@ -11,7 +14,7 @@ int main(const int argc, const char *argv[]) { try { horcrux::Horcrux h(horcrux::parse_arguments(std::move(arguments))); - h.run(); + h.run(buffer_size); if (h.mode() == horcrux::Mode::kEncrypt) { // print encryption key to stdout std::cout << h.key() << std::endl; diff --git a/test/integration-test.cpp b/test/integration-test.cpp index 18517b3..2c69630 100644 --- a/test/integration-test.cpp +++ b/test/integration-test.cpp @@ -62,35 +62,72 @@ TEST(IntegrationTest, HorcruxOptions){ options.base64_key = to_base64(generate_random(32)); EXPECT_NO_THROW(Horcrux {options}); } -TEST(IntegrationTest, HorcruxEncryptiDecrypt) { - HorcruxOptions options { - .mode = Mode::kEncrypt, - .count = 3, - .input = {image}, - .output = {folder} - }; - Horcrux enc{options}; +TEST(IntegrationTest, HorcruxEncryptDecrypt) { + { + HorcruxOptions options { + .mode = Mode::kEncrypt, + .count = 3, + .input = {image}, + .output = {folder} + }; + Horcrux enc{options}; + enc.run(); + auto key = enc.key(); - enc.run(); - auto key = enc.key(); + HorcruxOptions options2 { + .mode = Mode::kDecrypt, + .count = 3, + .input = {get_encrypted_files(folder, image)}, + .output = {noexist}, + .base64_key = key + }; - HorcruxOptions options2 { - .mode = Mode::kDecrypt, - .count = 3, - .input = {get_encrypted_files(folder, image)}, - .output = {noexist}, - .base64_key = key - }; - - Horcrux dec{options2}; - dec.run(); + Horcrux dec{options2}; + dec.run(); + std::for_each(options2.input.begin(), options2.input.end(), [](auto& p){ + std::filesystem::remove(p);}); + } // Close scope so that Output objects get flushed //Compare input and ouput files + EXPECT_EQ(std::filesystem::file_size(image), std::filesystem::file_size(noexist)); EXPECT_EQ(generic_read_file(image), generic_read_file(noexist)); std::filesystem::remove(noexist); - std::for_each(options2.input.begin(), options2.input.end(), [](auto& p){ - std::filesystem::remove(p);}); } + +TEST(IntegrationTest, HorcruxEncryptDecryptBuffer) { + { + HorcruxOptions options { + .mode = Mode::kEncrypt, + .count = 3, + .input = {image}, + .output = {folder} + }; + Horcrux enc{options}; + // read and process 1024 bytes at a time + enc.run(1024); + auto key = enc.key(); + + HorcruxOptions options2 { + .mode = Mode::kDecrypt, + .count = 3, + .input = {get_encrypted_files(folder, image)}, + .output = {noexist}, + .base64_key = key + }; + + Horcrux dec{options2}; + // read and process 2048 bytes at a time + dec.run(2048); + std::for_each(options2.input.begin(), options2.input.end(), [](auto& p){ + std::filesystem::remove(p);}); + } // Close scope so that Output objects get flushed + + //Compare input and ouput files + EXPECT_EQ(std::filesystem::file_size(image), std::filesystem::file_size(noexist)); + EXPECT_EQ(generic_read_file(image), generic_read_file(noexist)); + std::filesystem::remove(noexist); +} + TEST(IntegrationTest, arguments){ EXPECT_THROW( parse_arguments({"horcrux", "create", "-n", "4", "input_file"}), diff --git a/test/io-test.cpp b/test/io-test.cpp index 86195ba..6fa8a65 100644 --- a/test/io-test.cpp +++ b/test/io-test.cpp @@ -2,33 +2,38 @@ #include "io.h" #include "test.h" + TEST(IoTests, FsPlainInput) { EXPECT_THROW(horcrux::FsPlainInput input{noexist}, std::invalid_argument); EXPECT_THROW(horcrux::FsPlainInput input{folder}, std::invalid_argument); EXPECT_THROW(horcrux::FsPlainInput input{empty}, std::invalid_argument); EXPECT_NO_THROW(horcrux::FsPlainInput input{text}); EXPECT_NO_THROW(horcrux::FsPlainInput input{image}); + check_files_status(); } TEST(IoTests, FsPlainInputReadText) { auto input = horcrux::FsPlainInput{text}; auto buf = generic_read_file(text); EXPECT_EQ(buf, input.read()); + check_files_status(); } TEST(IoTests, FsPlainInputReadImg) { auto input = horcrux::FsPlainInput{image}; auto buf = generic_read_file(image); EXPECT_EQ(buf, input.read()); + check_files_status(); } TEST(IoTests, FsCryptoOutput) { EXPECT_THROW(horcrux::FsCryptoOutput output(noexist, 1), std::invalid_argument); - EXPECT_NO_THROW(horcrux::FsCryptoOutput output(folder, 1)); + EXPECT_NO_THROW(horcrux::FsCryptoOutput output(folder, 1, "test_chunk", 1024)); EXPECT_THROW(horcrux::FsCryptoOutput output(folder, 0), std::invalid_argument); EXPECT_THROW(horcrux::FsCryptoOutput output(empty, 1), std::invalid_argument); EXPECT_THROW(horcrux::FsCryptoOutput output(text, 1), std::invalid_argument); EXPECT_THROW(horcrux::FsCryptoOutput output(image, 1), std::invalid_argument); + check_files_status(); } TEST(IoTests, FsCryptoOutputWriteImage){ @@ -41,6 +46,7 @@ TEST(IoTests, FsCryptoOutputWriteImage){ EXPECT_EQ(10, out.created_files.size()); delete_created_files(out); + check_files_status(); } TEST(IoTests, FsCryptoOutputWriteText){ @@ -52,6 +58,7 @@ TEST(IoTests, FsCryptoOutputWriteText){ EXPECT_EQ(2, out.created_files.size()); delete_created_files(out); + check_files_status(); } TEST(IoTest, FsCryptoInput){ @@ -62,6 +69,7 @@ TEST(IoTest, FsCryptoInput){ EXPECT_THROW(horcrux::FsCryptoInput output({empty, empty}), std::invalid_argument); EXPECT_NO_THROW(horcrux::FsCryptoInput output({empty, empty, text})); EXPECT_NO_THROW(horcrux::FsCryptoInput output({image, empty, text})); + check_files_status(); } TEST(IoTest, FsCryptoInputImage){ @@ -77,11 +85,12 @@ TEST(IoTest, FsCryptoInputImage){ EXPECT_EQ(read, buf); files.push_back(text); - in = horcrux::FsCryptoInput(files); - read = in.read(); + auto in2 = horcrux::FsCryptoInput(files); + read = in2.read(); EXPECT_NE(read, buf); delete_created_files(out); + check_files_status(); } TEST(IoTest, FsPlainOuput){ @@ -91,6 +100,7 @@ TEST(IoTest, FsPlainOuput){ EXPECT_NO_THROW(horcrux::FsPlainOutput output{empty}); EXPECT_NO_THROW(horcrux::FsPlainOutput output{text}); EXPECT_NO_THROW(horcrux::FsPlainOutput output{image}); + check_files_status(); } TEST(IoTest, FsPlainOutputWrite){ @@ -101,6 +111,7 @@ TEST(IoTest, FsPlainOutputWrite){ auto buf2 = generic_read_file(noexist); EXPECT_EQ(buf, buf2); std::filesystem::remove(noexist); + check_files_status(); } TEST(IoTest, PlainToPlain){ @@ -115,6 +126,7 @@ TEST(IoTest, PlainToPlain){ auto buf2 = generic_read_file(noexist); EXPECT_EQ(buf, buf2); std::filesystem::remove(noexist); + check_files_status(); } TEST(IoTest, PlainToCrypto){ auto buf = generic_read_file(image); @@ -130,6 +142,7 @@ TEST(IoTest, PlainToCrypto){ auto buf2 = in_test.read(); EXPECT_EQ(buf, buf2); delete_created_files(out); + check_files_status(); } TEST(IoTest, CryptoToCrypto){ @@ -153,6 +166,7 @@ TEST(IoTest, CryptoToCrypto){ delete_created_files(out_test); delete_created_files(out); + check_files_status(); } TEST(IoTest, CryptoToPlain){ @@ -173,4 +187,5 @@ TEST(IoTest, CryptoToPlain){ EXPECT_EQ(buf, buf2); delete_created_files(out_test); std::filesystem::remove(noexist); + check_files_status(); } diff --git a/test/main.cpp b/test/main.cpp index 148880c..568d4b1 100644 --- a/test/main.cpp +++ b/test/main.cpp @@ -36,6 +36,20 @@ std::vector get_encrypted_files(const std::string& folder, return result; } +void check_files_status(){ + const std::filesystem::path p_folder{folder}; + const std::filesystem::path p_noexist{noexist}; + const std::filesystem::path p_empty{empty}; + const std::filesystem::path p_text{text}; + const std::filesystem::path p_image{image}; + + EXPECT_EQ(false, std::filesystem::exists(p_noexist)); + EXPECT_EQ(0, std::filesystem::file_size(p_empty)); + EXPECT_EQ(1388, std::filesystem::file_size(p_text)); + EXPECT_EQ(124106, std::filesystem::file_size(p_image)); + +} + int main(int argc, char *argv[]) { ::testing::InitGoogleTest(&argc, argv); diff --git a/test/test.h b/test/test.h index c16ec7e..415f96c 100644 --- a/test/test.h +++ b/test/test.h @@ -25,6 +25,8 @@ void delete_created_files(const horcrux::FsCryptoOutput& out); std::vector get_encrypted_files(const std::string& folder, const std::string& basename); +void check_files_status(); + /* Crypto Test utils */ /* test command: