Database.cpp

#include <filesystem>
#include <utility>
#include <fstream>
#include "Database.h"
#include <iterator>
#include <chrono>

#define DEBUG_ENCRYPTION_MANAGER 0
#define LOG_READ_VALUES 0
#define LOG_WRITE_VALUES 0

#if LOG_READ_VALUES == 1 || LOG_WRITE_VALUES == 1 || DEBUG_ENCRYPTION_MANAGER == 1
#include <iostream>
#endif

using std::ifstream, std::shared_ptr, std::ofstream, std::unique_ptr;

Database::Database(fs::path path, string password) {
    this->path = path;
    this->password = password;
}

#if LOG_WRITE_VALUES == 1 || LOG_READ_VALUES == 1
void printSharedPtr(shared_ptr<unsigned char[]> ptr, int size) {
    std::cout << std::hex;
    for (int i = 0; i < size; ++i) {
        std::cout << (int)ptr.get()[i];
    }
    std::cout << std::dec << std::endl;
}
#endif

void Database::load() {
    categories.clear();
    entries.clear();

    // open file in binary mode
    ifstream file(this->path, std::ios::binary);

    // stop eating new lines
    file.unsetf(std::ios::skipws);

    // read database password salt
    shared_ptr<unsigned char[]> salt(new unsigned char[EncryptionManager::KEY_SIZE]);
    file.read((char*)salt.get(), EncryptionManager::KEY_SIZE);
#if LOG_READ_VALUES == 1
    std::cout << "Read salt:\n";
    printSharedPtr(salt, EncryptionManager::KEY_SIZE);
#endif

    // read initialization vector
    shared_ptr<unsigned char[]> initializationVector(new unsigned char[EncryptionManager::BLOCK_SIZE]);
    file.read((char*)initializationVector.get(), EncryptionManager::BLOCK_SIZE);
#if LOG_READ_VALUES == 1
    std::cout << "Read init vector:\n";
    printSharedPtr(initializationVector, EncryptionManager::BLOCK_SIZE);
#endif

    if (this->encryptionManager == nullptr) {
        this->encryptionManager = std::make_unique<EncryptionManager>(salt, initializationVector, password);
#if DEBUG_ENCRYPTION_MANAGER == 1
        std::cout << "Created encryption manager for reading\n";
#endif
    }

    // read timestamp
//    shared_ptr<unsigned char[]> timestamp(new unsigned char[sizeof(long)]);
//    file.read((char*)initializationVector.get(), sizeof(long));
//    this->lastReadTimestamp = *((long*)timestamp.get());
    long* timestamp = new long;
    file.read((char*)timestamp, sizeof(long));
    lastReadTimestamp = *timestamp;
    delete timestamp;
#if LOG_READ_VALUES == 1
    std::cout << "Read timestamp: " << lastReadTimestamp << std::endl;
#endif

    // read size of entries
    unique_ptr<char[]> sizeBuffer(new char[sizeof(unsigned long)]);
    file.read(sizeBuffer.get(), sizeof(unsigned long));
    unsigned long size = *((long*)sizeBuffer.get());
#if LOG_READ_VALUES == 1
    std::cout << "Size of password entries: " << size << std::endl;
#endif

    // allocate memory for entries
    vector<unsigned char> bytes(size);

    // read and decrypt entries
    file.read((char*)&bytes[0], size);
#if LOG_READ_VALUES == 1
    std::cout << "Size of read entries: " << bytes.size() << std::endl;
#endif
    this->encryptionManager->decrypt(bytes);
    char* rangeFromPointer = (char*)&bytes[0];
    for (int i = 0; i < bytes.size(); ++i) {
        char* ptr = (char*)&bytes[0] + i;
        if (*ptr == '\n') {
            auto distance = ptr - rangeFromPointer;
            string formattedEntry(rangeFromPointer, distance);
            rangeFromPointer = ptr + 1; // skip current delimiter

            PasswordEntry passwordEntry(formattedEntry);
            entries.push_back(passwordEntry);
        }
    }

    bytes.clear();
    unsigned char temp;
    while (file >> temp) {
        bytes.push_back(temp);
    }

#if LOG_READ_VALUES == 1
    std::cout << "Size of read categories: " << bytes.size() << std::endl;
#endif
    this->encryptionManager->decrypt(bytes);

#if LOG_READ_VALUES == 1
    std::cout << "Size of read categories after decryption: " << bytes.size() << std::endl;
#endif

    rangeFromPointer = (char*)bytes.data();
    for (int i = 0; i < bytes.size(); ++i) {
        char* ptr = (char*)bytes.data() + i;
        if (*ptr == '\n') {
            auto distance = ptr - rangeFromPointer; // -1 to skip the delimiter

            string category(rangeFromPointer, distance);
#if LOG_READ_VALUES == 1
            std::cout << "Read category: " << category << std::endl;
#endif
            categories.push_back(category);

            rangeFromPointer = ptr + 1;
        }
    }
}

void Database::save() {
    ofstream file(this->path, std::ios::out | std::ios::binary);

    // create encryption manager if it wasn't created before
    // generate salt and initialization vector
    if (this->encryptionManager == nullptr) {
        shared_ptr<unsigned char[]> salt(new unsigned char[EncryptionManager::SALT_SIZE]);
        for (int i = 0; i < EncryptionManager::SALT_SIZE; ++i) {
            salt[i] = (unsigned char) rand();
        }

        shared_ptr<unsigned char[]> initializationVector(new unsigned char[EncryptionManager::BLOCK_SIZE]);
        for (int i = 0; i < EncryptionManager::BLOCK_SIZE; ++i) {
            salt[i] = (unsigned char) rand();
        }

        this->encryptionManager = std::make_unique<EncryptionManager>(salt, initializationVector, password);
#if DEBUG_ENCRYPTION_MANAGER == 1
        std::cout << "Created encryption manager for saving\n";
#endif
    }

    // write salt
    shared_ptr<unsigned char[]> salt = this->encryptionManager->salt;
    file.write((char*)salt.get(), EncryptionManager::SALT_SIZE);
#if LOG_WRITE_VALUES == 1
    std::cout << "Written salt:\n";
    printSharedPtr(salt, EncryptionManager::SALT_SIZE);
#endif

    // write initialization vector
    shared_ptr<unsigned char[]> initializationVector = this->encryptionManager->initializationVector;
    file.write((char*)initializationVector.get(), EncryptionManager::BLOCK_SIZE);
#if LOG_WRITE_VALUES == 1
    std::cout << "Written init vector:\n";
    printSharedPtr(initializationVector, EncryptionManager::BLOCK_SIZE);
#endif

    // write timestamp
    const auto now = std::chrono::system_clock::now();
    long timestamp = std::chrono::duration_cast<std::chrono::seconds>(now.time_since_epoch()).count();
    file.write((char*)&timestamp, sizeof(long));
#if LOG_WRITE_VALUES == 1
    std::cout << "Written timestamp: " << timestamp << std::endl;
#endif

    // encrypt and write entries
    vector<unsigned char> bytes;
    for (int i = 0; i < entries.size(); ++i) {
        const string formattedEntry = entries[i].getFormattedEntry();
        for (int j = 0; j < formattedEntry.length(); ++j) {
            bytes.push_back((unsigned char)formattedEntry.c_str()[j]);
        }
        bytes.push_back((unsigned char)'\n');
    }
    this->encryptionManager->encrypt(bytes);
    unsigned long size = bytes.size();
#if LOG_WRITE_VALUES == 1
    std::cout << "Written size of entries: " << size << std::endl;
#endif
    file.write((char*)&size, sizeof(unsigned long)); // write size of the encrypted entries
    file.write((char*)bytes.data(), bytes.size());

    // encrypt and write categories
    bytes.clear();
    for (int i = 0; i < categories.size(); ++i) {
        for (int j = 0; j < categories[i].length(); ++j) {
            bytes.push_back((unsigned char)categories[i].c_str()[j]);
        }
        bytes.push_back((unsigned char)'\n');
    }
    this->encryptionManager->encrypt(bytes);
#if LOG_WRITE_VALUES == 1
    std::cout << "Size of written categories: " << bytes.size() << "\n";
#endif
    file.write((char*)bytes.data(), bytes.size());
}

void Database::addPassword(const PasswordEntry& entry) {
    if (entries.empty() || std::find(entries.begin(), entries.end(), entry) == entries.end()) {
        bool categoryExists = false;
        for (int i = 0; i < categories.size(); ++i) {
            if (categories[i] == entry.getCategory()) {
                categoryExists = true;
                break;
            }
        }
        if (!categoryExists) {
            throw std::runtime_error("Tried to add a password to a category that doesn't exist");
        }

        needsSaving = true;
        entries.push_back(entry);
    } else {
        throw std::runtime_error("Tried to add a password that already exists");
    }
}

void Database::addCategory(string &category) {
    if (categories.empty() || std::find(categories.begin(), categories.end(), category) == categories.end()) {
        needsSaving = true;
        categories.push_back(category);
    } else {
        string message = "Tried to add category that already exists: [" + category + "]";
        throw std::runtime_error(message);
    }
}

vector<PasswordEntry> Database::getPasswords() {
    return this->entries;
}

vector<string> Database::getCategories() {
    return this->categories;
}

void Database::editPassword(const int& entryIndex, const PasswordEntry &newEntry) {
    if (entryIndex >= 0 && entryIndex < entries.size()) {
        entries[entryIndex] = newEntry;
        needsSaving = true;
    } else {
        throw std::runtime_error("Invalid entry index");
    }
}

void Database::deletePassword(const int& entryIndex) {
    if (entryIndex >= 0 && entryIndex < entries.size()) {
        needsSaving = true;
        entries.erase(entries.begin() + entryIndex);
    } else {
        throw std::runtime_error("Invalid entry index");
    }
}

void Database::deleteCategoryAndItsPasswords(const int& categoryIndex) {
    if (categoryIndex < 0 || categoryIndex >= categories.size()) {
        throw std::runtime_error("Invalid category index");
    }

    needsSaving = true;
    string category = categories[categoryIndex];
    categories.erase(categories.begin() + categoryIndex);
    vector<int> entriesToDelete;
    for (int i = 0; i < entries.size(); ++i) {
        PasswordEntry entry = entries[i];
        if (entry.getCategory() == category) {
            entriesToDelete.push_back(i);
        }
    }

    for (int i = 0; i < entriesToDelete.size(); ++i) {
        deletePassword(entriesToDelete[i]);
    }
}

bool Database::getNeedsSaving() {
    return this->needsSaving;
}