change Aes256 to accept buffer

include/aes256.hpp

@@ -25,64 +25,60 @@
 #define AES256_HPP
 
 #include <vector>
-
+#include "buffer.hpp"
-typedef std::vector<unsigned char> ByteArray;
 
 #define BLOCK_SIZE 16
 
-class Aes256 {
+class Aes256
+{
 
-    public:
+public:
-        Aes256(const ByteArray& key);
+    Aes256(const std::string &key);
-        ~Aes256();
+    ~Aes256();
 
-        static ByteArray::size_type encrypt(const ByteArray& key, const ByteArray& plain, ByteArray& encrypted);
+    static bool encrypt(const std::string &key, const Buffer &plain, Buffer &encrypted);
-        static ByteArray::size_type encrypt(const ByteArray& key, const unsigned char* plain, const ByteArray::size_type plain_length, ByteArray& encrypted);
+    static bool decrypt(const std::string &key, const Buffer &encrypted, Buffer &plain);
-        static ByteArray::size_type decrypt(const ByteArray& key, const ByteArray& encrypted, ByteArray& plain);
-        static ByteArray::size_type decrypt(const ByteArray& key, const unsigned char* encrypted, const ByteArray::size_type encrypted_length, ByteArray& plain);
 
-        ByteArray::size_type encrypt_start(const ByteArray::size_type plain_length, ByteArray& encrypted);
+    std::size_t encrypt_start(const std::size_t plain_length, Buffer &encrypted);
-        ByteArray::size_type encrypt_continue(const ByteArray& plain, ByteArray& encrypted);
+    std::size_t encrypt_continue(const Buffer &plain, Buffer &encrypted);
-        ByteArray::size_type encrypt_continue(const unsigned char* plain, const ByteArray::size_type plain_length, ByteArray& encrypted);
+    std::size_t encrypt_end(Buffer &encrypted);
-        ByteArray::size_type encrypt_end(ByteArray& encrypted);
 
-        ByteArray::size_type decrypt_start(const ByteArray::size_type encrypted_length);
+    std::size_t decrypt_start(const std::size_t encrypted_length);
-        ByteArray::size_type decrypt_continue(const ByteArray& encrypted, ByteArray& plain);
+    std::size_t decrypt_continue(const Buffer &encrypted, Buffer &plain);
-        ByteArray::size_type decrypt_continue(const unsigned char* encrypted, const ByteArray::size_type encrypted_length, ByteArray& plain);
+    std::size_t decrypt_end(Buffer &plain);
-        ByteArray::size_type decrypt_end(ByteArray& plain);
 
-    private:
+private:
-        ByteArray            m_key;
+    std::vector<unsigned char> m_key;
-        ByteArray            m_salt;
+    std::vector<unsigned char> m_salt;
-        ByteArray            m_rkey;
+    std::vector<unsigned char> m_rkey;
 
-        unsigned char        m_buffer[3 * BLOCK_SIZE];
+    unsigned char m_buffer[3 * BLOCK_SIZE];
-        unsigned char        m_buffer_pos;
+    unsigned char m_buffer_pos;
-        ByteArray::size_type m_remainingLength;
+    std::size_t m_remainingLength;
 
-        bool                 m_decryptInitialized;
+    bool m_decryptInitialized;
 
-        void check_and_encrypt_buffer(ByteArray& encrypted);
+    void check_and_encrypt_buffer(Buffer &encrypted);
-        void check_and_decrypt_buffer(ByteArray& plain);
+    void check_and_decrypt_buffer(Buffer &plain);
 
-        void encrypt(unsigned char *buffer);
+    void encrypt(unsigned char *buffer);
-        void decrypt(unsigned char *buffer);
+    void decrypt(unsigned char *buffer);
 
-        void expand_enc_key(unsigned char *rc);
+    void expand_enc_key(unsigned char *rc);
-        void expand_dec_key(unsigned char *rc);
+    void expand_dec_key(unsigned char *rc);
 
-        void sub_bytes(unsigned char *buffer);
+    void sub_bytes(unsigned char *buffer);
-        void sub_bytes_inv(unsigned char *buffer);
+    void sub_bytes_inv(unsigned char *buffer);
 
-        void copy_key();
+    void copy_key();
 
-        void add_round_key(unsigned char *buffer, const unsigned char round);
+    void add_round_key(unsigned char *buffer, const unsigned char round);
 
-        void shift_rows(unsigned char *buffer);
+    void shift_rows(unsigned char *buffer);
-        void shift_rows_inv(unsigned char *buffer);
+    void shift_rows_inv(unsigned char *buffer);
 
-        void mix_columns(unsigned char *buffer);
+    void mix_columns(unsigned char *buffer);
-        void mix_columns_inv(unsigned char *buffer);
+    void mix_columns_inv(unsigned char *buffer);
 };
 
 #endif /* AES256_HPP */

include/buffer.hpp

@@ -9,7 +9,7 @@ class Buffer
 public:
 	uint8_t *buffer = nullptr;
 	size_t taken = 0;
-	size_t size = 0;
+
 	std::string file_path;
 
 	Buffer(size_t size);
@@ -35,6 +35,8 @@ public:
 
 	bool load_from_file();
 	bool load_from_file(std::string file_path);
+private:
+		size_t size = 0;
 };
 
-#endif
+#endif //BUFFER_HPP

include/cryptography.hpp

@@ -11,8 +11,8 @@ public:
 	Cryptography(std::string password);
 	~Cryptography();
 	void change_pass(std::string password);
-	bool encrypt(Buffer *plain, Buffer *encrypted);
+	bool encrypt(Buffer &plain, Buffer &encrypted);
-	bool decrypt(Buffer *encrypted, Buffer *decrypted);
+	bool decrypt(Buffer &encrypted, Buffer &decrypted);
 
 private:
 	std::string key;

main.cpp

@@ -78,8 +78,8 @@ int main(int argc, char **argv)
 
 	Buffer decrypted_buffer;
 	// check if encrypted buffer is empty if not, decrypt it
-	if (encrypted_buffer.size > 0)
+	if (encrypted_buffer.taken > 0)
-		if (!crypto.decrypt(&encrypted_buffer, &decrypted_buffer))
+		if (!crypto.decrypt(encrypted_buffer, decrypted_buffer))
 			return 1;
 
 	// if decrypted buffer is empty, add index

source/aes256.cpp

@@ -101,10 +101,16 @@ const unsigned char sboxinv[256] = {
     0x17, 0x2b, 0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26,
     0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d};
 
-Aes256::Aes256(const ByteArray &key)
+Aes256::Aes256(const std::string &key)
-    : m_key(ByteArray(key.size() > KEY_SIZE ? KEY_SIZE : key.size(), 0)), m_salt(ByteArray(KEY_SIZE - m_key.size(), 0)), m_rkey(ByteArray(KEY_SIZE, 0)), m_buffer_pos(0), m_remainingLength(0), m_decryptInitialized(false)
+    : m_key(std::vector<unsigned char>(key.size() > KEY_SIZE ? KEY_SIZE : key.size(), 0)),
+      m_salt(std::vector<unsigned char>(KEY_SIZE - m_key.size(), 0)),
+      m_rkey(std::vector<unsigned char>(KEY_SIZE, 0)),
+      m_buffer_pos(0),
+      m_remainingLength(0),
+      m_decryptInitialized(false)
 {
-    for (ByteArray::size_type i = 0; i < m_key.size(); ++i)
+
+    for (std::size_t i = 0; i < m_key.size(); ++i)
         m_key[i] = key[i];
 }
 
@@ -112,108 +118,73 @@ Aes256::~Aes256()
 {
 }
 
-ByteArray::size_type Aes256::encrypt(const ByteArray &key, const ByteArray &plain, ByteArray &encrypted)
+bool Aes256::encrypt(const std::string &key, const Buffer &plain, Buffer &encrypted)
 {
     Aes256 aes(key);
 
-    aes.encrypt_start(plain.size(), encrypted);
+    aes.encrypt_start(plain.taken, encrypted);
     aes.encrypt_continue(plain, encrypted);
     aes.encrypt_end(encrypted);
 
-    return encrypted.size();
+    return true;
 }
 
-ByteArray::size_type Aes256::encrypt(const ByteArray &key, const unsigned char *plain, const ByteArray::size_type plain_length, ByteArray &encrypted)
+bool Aes256::decrypt(const std::string &key, const Buffer &encrypted, Buffer &plain)
 {
     Aes256 aes(key);
 
-    aes.encrypt_start(plain_length, encrypted);
+    aes.decrypt_start(encrypted.taken);
-    aes.encrypt_continue(plain, plain_length, encrypted);
-    aes.encrypt_end(encrypted);
-
-    return encrypted.size();
-}
-
-ByteArray::size_type Aes256::decrypt(const ByteArray &key, const ByteArray &encrypted, ByteArray &plain)
-{
-    Aes256 aes(key);
-
-    aes.decrypt_start(encrypted.size());
     aes.decrypt_continue(encrypted, plain);
     aes.decrypt_end(plain);
 
-    return plain.size();
+    return true;
 }
 
-ByteArray::size_type Aes256::decrypt(const ByteArray &key, const unsigned char *encrypted, const ByteArray::size_type encrypted_length, ByteArray &plain)
+std::size_t Aes256::encrypt_start(const std::size_t plain_length, Buffer &encrypted)
-{
-    Aes256 aes(key);
-
-    aes.decrypt_start(encrypted_length);
-    aes.decrypt_continue(encrypted, encrypted_length, plain);
-    aes.decrypt_end(plain);
-
-    return plain.size();
-}
-
-ByteArray::size_type Aes256::encrypt_start(const ByteArray::size_type plain_length, ByteArray &encrypted)
 {
     m_remainingLength = plain_length;
 
     // Generate salt
-    ByteArray::iterator it = m_salt.begin(), itEnd = m_salt.end();
+    std::vector<unsigned char>::iterator it = m_salt.begin(), itEnd = m_salt.end();
     while (it != itEnd)
         *(it++) = (rand() & 0xFF);
 
     // Calculate padding
-    ByteArray::size_type padding = 0;
+    std::size_t padding = 0;
     if (m_remainingLength % BLOCK_SIZE != 0)
         padding = (BLOCK_SIZE - (m_remainingLength % BLOCK_SIZE));
     m_remainingLength += padding;
 
     // Add salt
-    encrypted.insert(encrypted.end(), m_salt.begin(), m_salt.end());
+    encrypted.add_end(m_salt.data(), m_salt.size());
     m_remainingLength += m_salt.size();
 
-    // Add 1 bytes for padding size
+    // Add 1 byte for padding size
-    encrypted.push_back(padding & 0xFF);
+    unsigned char pad = (padding & 0xFF);
+    encrypted.add_end(&pad, 1);
     ++m_remainingLength;
 
     // Reset buffer
     m_buffer_pos = 0;
 
-    return encrypted.size();
+    return encrypted.taken;
 }
 
-ByteArray::size_type Aes256::encrypt_continue(const ByteArray &plain, ByteArray &encrypted)
+std::size_t Aes256::encrypt_continue(const Buffer &plain, Buffer &encrypted)
 {
-    ByteArray::const_iterator it = plain.begin(), itEnd = plain.end();
+    std::size_t i = 0;
 
-    while (it != itEnd)
+    while (i < plain.taken)
     {
-        m_buffer[m_buffer_pos++] = *(it++);
+        m_buffer[m_buffer_pos++] = plain.buffer[i++];
 
         check_and_encrypt_buffer(encrypted);
     }
 
-    return encrypted.size();
+    return encrypted.taken;
 }
 
-ByteArray::size_type Aes256::encrypt_continue(const unsigned char *plain, const ByteArray::size_type plain_length, ByteArray &encrypted)
+void Aes256::check_and_encrypt_buffer(Buffer &encrypted)
-{
-    ByteArray::size_type i = 0;
-
-    while (i < plain_length)
-    {
-        m_buffer[m_buffer_pos++] = plain[i++];
-
-        check_and_encrypt_buffer(encrypted);
-    }
-
-    return encrypted.size();
-}
-
-void Aes256::check_and_encrypt_buffer(ByteArray &encrypted)
 {
     if (m_buffer_pos == BLOCK_SIZE)
     {
@@ -221,7 +192,7 @@ void Aes256::check_and_encrypt_buffer(ByteArray &encrypted)
 
         for (m_buffer_pos = 0; m_buffer_pos < BLOCK_SIZE; ++m_buffer_pos)
         {
-            encrypted.push_back(m_buffer[m_buffer_pos]);
+            encrypted.add_end(&m_buffer[m_buffer_pos], sizeof(unsigned char));
             --m_remainingLength;
         }
 
@@ -229,7 +200,7 @@ void Aes256::check_and_encrypt_buffer(ByteArray &encrypted)
     }
 }
 
-ByteArray::size_type Aes256::encrypt_end(ByteArray &encrypted)
+std::size_t Aes256::encrypt_end(Buffer &encrypted)
 {
     if (m_buffer_pos > 0)
     {
@@ -240,14 +211,14 @@ ByteArray::size_type Aes256::encrypt_end(ByteArray &encrypted)
 
         for (m_buffer_pos = 0; m_buffer_pos < BLOCK_SIZE; ++m_buffer_pos)
         {
-            encrypted.push_back(m_buffer[m_buffer_pos]);
+            encrypted.add_end(&m_buffer[m_buffer_pos], sizeof(unsigned char));
             --m_remainingLength;
         }
 
         m_buffer_pos = 0;
     }
 
-    return encrypted.size();
+    return encrypted.taken;
 }
 
 void Aes256::encrypt(unsigned char *buffer)
@@ -271,7 +242,7 @@ void Aes256::encrypt(unsigned char *buffer)
     add_round_key(buffer, i);
 }
 
-ByteArray::size_type Aes256::decrypt_start(const ByteArray::size_type encrypted_length)
+std::size_t Aes256::decrypt_start(const std::size_t encrypted_length)
 {
     unsigned char j;
 
@@ -290,40 +261,26 @@ ByteArray::size_type Aes256::decrypt_start(const ByteArray::size_type encrypted_
     return m_remainingLength;
 }
 
-ByteArray::size_type Aes256::decrypt_continue(const ByteArray &encrypted, ByteArray &plain)
+std::size_t Aes256::decrypt_continue(const Buffer &encrypted, Buffer &plain)
 {
-    ByteArray::const_iterator it = encrypted.begin(), itEnd = encrypted.end();
+    std::size_t i = 0;
 
-    while (it != itEnd)
+    while (i < encrypted.taken)
     {
-        m_buffer[m_buffer_pos++] = *(it++);
+        m_buffer[m_buffer_pos++] = encrypted.buffer[i++];
 
         check_and_decrypt_buffer(plain);
     }
 
-    return plain.size();
+    return plain.taken;
 }
 
-ByteArray::size_type Aes256::decrypt_continue(const unsigned char *encrypted, const ByteArray::size_type encrypted_length, ByteArray &plain)
+void Aes256::check_and_decrypt_buffer(Buffer &plain)
-{
-    ByteArray::size_type i = 0;
-
-    while (i < encrypted_length)
-    {
-        m_buffer[m_buffer_pos++] = encrypted[i++];
-
-        check_and_decrypt_buffer(plain);
-    }
-
-    return plain.size();
-}
-
-void Aes256::check_and_decrypt_buffer(ByteArray &plain)
 {
     if (!m_decryptInitialized && m_buffer_pos == m_salt.size() + 1)
     {
         unsigned char j;
-        ByteArray::size_type padding;
+        std::size_t padding;
 
         // Get salt
         for (j = 0; j < m_salt.size(); ++j)
@@ -345,7 +302,7 @@ void Aes256::check_and_decrypt_buffer(ByteArray &plain)
         for (m_buffer_pos = 0; m_buffer_pos < BLOCK_SIZE; ++m_buffer_pos)
             if (m_remainingLength > 0)
             {
-                plain.push_back(m_buffer[m_buffer_pos]);
+                plain.add_end(&m_buffer[m_buffer_pos], sizeof(unsigned char));
                 --m_remainingLength;
             }
 
@@ -353,9 +310,9 @@ void Aes256::check_and_decrypt_buffer(ByteArray &plain)
     }
 }
 
-ByteArray::size_type Aes256::decrypt_end(ByteArray &plain)
+std::size_t Aes256::decrypt_end(Buffer &plain)
 {
-    return plain.size();
+    return plain.taken;
 }
 
 void Aes256::decrypt(unsigned char *buffer)
@@ -463,7 +420,7 @@ void Aes256::sub_bytes_inv(unsigned char *buffer)
 
 void Aes256::copy_key()
 {
-    ByteArray::size_type i;
+    std::size_t i;
 
     for (i = 0; i < m_key.size(); ++i)
         m_rkey[i] = m_key[i];

source/arg_func.cpp

@@ -156,7 +156,7 @@ std::optional<LoginInfoPointer> arg_new_password(Buffer &decrypted_buffer, Buffe
 	}
 
 	add_logininfo_to_buffer(decrypted_buffer, name.c_str(), username.c_str(), password.c_str());
-	crypto.encrypt(&decrypted_buffer, &encrypted_buffer);
+	crypto.encrypt(decrypted_buffer, encrypted_buffer);
 	encrypted_buffer.save_to_file();
 	Index *index = (Index *)decrypted_buffer.buffer;
 	return get_logininfo_pointer_from_buffer(decrypted_buffer, index->count - 1);
@@ -181,7 +181,7 @@ void arg_username(Buffer &decrypted_buffer, Buffer &encrypted_buffer, const char
 
 	delete_logininfo_from_buffer(decrypted_buffer, pass);
 	add_logininfo_to_buffer(decrypted_buffer, name.c_str(), username.c_str(), password.c_str());
-	crypto.encrypt(&decrypted_buffer, &encrypted_buffer);
+	crypto.encrypt(decrypted_buffer, encrypted_buffer);
 	encrypted_buffer.save_to_file();
 }
 
@@ -205,7 +205,7 @@ void arg_label_name(Buffer &decrypted_buffer, Buffer &encrypted_buffer, const ch
 
 	delete_logininfo_from_buffer(decrypted_buffer, pass);
 	add_logininfo_to_buffer(decrypted_buffer, name.c_str(), username.c_str(), password.c_str());
-	crypto.encrypt(&decrypted_buffer, &encrypted_buffer);
+	crypto.encrypt(decrypted_buffer, encrypted_buffer);
 	encrypted_buffer.save_to_file();
 }
 
@@ -241,7 +241,7 @@ void arg_delete(Buffer &decrypted_buffer, Buffer &encrypted_buffer, const char *
 
 	delete_logininfo_from_buffer(decrypted_buffer, pass);
 
-	crypto.encrypt(&decrypted_buffer, &encrypted_buffer);
+	crypto.encrypt(decrypted_buffer, encrypted_buffer);
 	encrypted_buffer.save_to_file();
 	printf("Password deleted\n");
 }
@@ -313,7 +313,7 @@ void arg_change(Buffer &decrypted_buffer, Buffer &encrypted_buffer, std::string
 	}
 
 	crypto.change_pass(new_string);
-	crypto.encrypt(&decrypted_buffer, &encrypted_buffer);
+	crypto.encrypt(decrypted_buffer, encrypted_buffer);
 	encrypted_buffer.save_to_file();
 	printf("Password changed\n");
 }
@@ -348,7 +348,7 @@ void arg_file(Buffer &decrypted_buffer, Buffer &encrypted_buffer, const char *la
 	if (decrypted_buffer.taken <= sizeof(Index))
 		return;
 
-	crypto.encrypt(&decrypted_buffer, &encrypted_buffer);
+	crypto.encrypt(decrypted_buffer, encrypted_buffer);
 	encrypted_buffer.save_to_file(save.string());
 }
 

source/cryptography.cpp

@@ -21,24 +21,16 @@ void Cryptography::change_pass(std::string password)
 	key = password;
 }
 
-bool Cryptography::encrypt(Buffer *plain, Buffer *encrypted)
+bool Cryptography::encrypt(Buffer &plain, Buffer &encrypted)
 {
-	ByteArray in(plain->buffer, plain->buffer + plain->taken);
+
-	ByteArray out;
+	Aes256::encrypt(key, plain, encrypted);
-	ByteArray key_b(key.begin(), key.end());
+
-	Aes256::encrypt(key_b, in, out);
-	encrypted->taken = 0;
-	encrypted->add_end(out.data(), out.size());
 	return true;
 }
 
-bool Cryptography::decrypt(Buffer *encrypted, Buffer *decrypted)
+bool Cryptography::decrypt(Buffer &encrypted, Buffer &decrypted)
 {
-	ByteArray in(encrypted->buffer, encrypted->buffer + encrypted->taken);
+	Aes256::decrypt(key, encrypted, decrypted);
-	ByteArray out;
-	ByteArray key_b(key.begin(), key.end());
-	Aes256::decrypt(key_b, in, out);
-	decrypted->taken = 0;
-	decrypted->add_end(out.data(), out.size());
 	return true;
 }