#include "hybridcrypt.h"

HybridCrypt::HybridCrypt()
{
    // Initialisiere Nutzerschlüssel mit NULL
    userKeypair = NULL;

    // Lade die menschenlesbaren Fehlerstrings für LibCrypto
    ERR_load_crypto_strings();

    // Lade alle Hash- und  Verschlüsselungsalgorithmen
    OpenSSL_add_all_algorithms();

    // Lade Konfigurationsdatei und andere wichtige Initialisierungen
    OPENSSL_config(NULL);

    // Zeige ob der Zufallszahlengenerator initialisiert wurde
    qDebug() << "Zufallszahlengenerator erfolgreich initialisiert: " <<
             isCsprngSeeded();
}

HybridCrypt::~HybridCrypt()
{
    // Räume Nutzerschlüssel auf
    freeEvpKey(&userKeypair);

    // Räume OpenSSL auf
    // Entferne alle Hash- und Verschlüsselungsalgorithmen
    EVP_cleanup();

    // Falls das nächste Ausgelassen wird, könnte ein Speicherleck auftreten,
    // wenn die BIO Api verwendet wurde (Base64 Transformationen)
    CRYPTO_cleanup_all_ex_data();

    // Lösche den CSPRNG sicher
    RAND_cleanup();

    // Entferne Fehlerstrings
    ERR_free_strings();
}

void HybridCrypt::encrypt(QString infileName, QString outfileName,
                          QVector<QString> recipientKeyfileNames)
{
    if (recipientKeyfileNames.length() < 1)
    {
        throw CryptException("Keine Empfänger wurden angegeben.",
                             CryptException::NoRecipients);
    }

    throwExceptionIfCsprngIsNotSeeded();

    QFile infile(infileName);
    QFile outfile(outfileName);

    if (!infile.open(QFile::ReadOnly))
    {
        throw CryptException("Datei nicht gefunden: " + infileName.toStdString(),
                             CryptException::FileNotFound);
    }

    if (!outfile.open(QFile::ReadWrite | QFile::Truncate))
    {
        throw CryptException("Konnte Datei nicht schreiben: " +
                             outfileName.toStdString(), CryptException::FileNotWritable);
    }

    QDataStream infileStream(&infile);
    QDataStream outfileStream(&outfile);

    QByteArray aesKey = getCsprngBytes(EVP_CIPHER_key_length(EVP_aes_256_cbc()));
    QByteArray aesIv = getCsprngBytes(EVP_CIPHER_iv_length(EVP_aes_256_cbc()));

    int keyCount = recipientKeyfileNames.length();

    // Nur wenn ein Nutzerschlüssel existiert wird er auch für die Verschlüsselung verwendet
    if (userKeypair != NULL)
    {
        ++keyCount;
    }

    // Schreibe die Anzahl der PublicKeys als erstes in die Datei
    outfileStream << (qint32) keyCount;

    EVP_PKEY *tmpKey;

    for (int i = 0; i < recipientKeyfileNames.length(); i++)
    {
        tmpKey = readRecipientKey(recipientKeyfileNames[i]);

        // Schreibe den Header in die Datei
        writeRsaHeader(&outfileStream, tmpKey, aesKey + aesIv);

        // Räume Schlüssel auf
        freeEvpKey(&tmpKey);
    }

    // Damit der Nutzer selbst die Datei auch entschlüsseln kann, falls er einen Schlüssel hat
    if (userKeypair != NULL)
    {
        writeRsaHeader(&outfileStream, userKeypair, aesKey + aesIv);
    }

    // Erzeuge den symmetrischen Kontext
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();

    // Überprüfe den Schlüssel
    if (!ctx)
    {
        throwOpenSslException();
    }

    // Setze den Verschlüsselungsalgorithmus, sowie den Schlüssel und den IV
    if (1 != EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL,
                                (unsigned char *) aesKey.data(), (unsigned char *) aesIv.data()))
    {
        // Räume den symmetrischen Schlüsselkontext ab
        freeCipherCtx(&ctx);
        throwOpenSslException();
    }

    QByteArray inBuf(EVP_CIPHER_block_size(EVP_aes_256_cbc()), 0);
    QByteArray outBuf(EVP_CIPHER_block_size(EVP_aes_256_cbc()), 0);
    int bufRead;
    int len;

    while (!infileStream.atEnd())
    {
        // Hole jeweils einen Block an Daten aus der Datei
        bufRead = infileStream.readRawData(inBuf.data(), inBuf.size());

        // Verschlüssele den Block mit AES
        if (1 != EVP_EncryptUpdate(ctx, (unsigned char *) outBuf.data(), &len,
                                   (unsigned char *) inBuf.data(), bufRead))
        {
            // Räume den symmetrischen Schlüsselkontext ab
            freeCipherCtx(&ctx);
            throwOpenSslException();
        }

        // Schreibe den Block in die Ausgabedatei
        outfileStream.writeRawData(outBuf.data(), len);
    }

    // Nun bekommen wir den letzten Block, falls die Datei nicht ein Vielfaches der AES Blöckgröße war
    if (1 != EVP_EncryptFinal_ex(ctx, (unsigned char *) outBuf.data(), &len))
    {
        // Räume den symmetrischen Schlüsselkontext ab
        freeCipherCtx(&ctx);
        throwOpenSslException();
    }

    outfileStream.writeRawData(outBuf.data(), len);

    // Räume den symmetrischen Schlüsselkontext ab
    freeCipherCtx(&ctx);

}

void HybridCrypt::decrypt(QString infileName, QString outfileName)
{
    throwExceptionIfUserKeyIsNull();

    QFile infile(infileName);
    QFile outfile(outfileName);

    if (!infile.open(QFile::ReadOnly))
    {
        throw CryptException("Datei nicht gefunden: " + infileName.toStdString(),
                             CryptException::FileNotFound);
    }

    if (!outfile.open(QFile::ReadWrite | QFile::Truncate))
    {
        throw CryptException("Konnte Datei nicht schreiben: " +
                             outfileName.toStdString(), CryptException::FileNotWritable);
    }

    QDataStream infileStream(&infile);
    QDataStream outfileStream(&outfile);

    qint32 keyCount = 0;

    // Lese die Anzahl der Schlüssel aus der Datei
    infileStream >> keyCount;

    qint32 tmpDataLength = 0;
    QByteArray tmpData;
    QByteArray aesKey(EVP_CIPHER_key_length(EVP_aes_256_cbc()), 0);
    QByteArray aesIv(EVP_CIPHER_iv_length(EVP_aes_256_cbc()), 0);

    // Versuche alle Header zu entschlüsseln
    for (qint32 i = 0; i < keyCount && !infileStream.atEnd(); i++)
    {
        infileStream >> tmpDataLength;

        // Finde den Anfang der AES verschlüsselten Daten,
        // wenn der AES Key und IV gefunden wurden
        if (tmpData.length() != 0)
        {
            infileStream.skipRawData(tmpDataLength);
        }
        else
        {
            tmpData.resize(tmpDataLength);
            infileStream.readRawData(tmpData.data(), tmpData.length());

            // Versuche nur den RSA Block zu entschlüsseln,
            // wenn die Modulogröße des Nutzerschlüssels auf die
            // des zum Verschlüsseln genutzten Schlüssels passt.
            if (tmpDataLength == RSA_size(userKeypair->pkey.rsa))
            {
                tmpData = decryptAesData(userKeypair, tmpData);
            }
            else
            {
                // Um in die nächste Entschlüsselungsrunde zu kommen
                tmpData.resize(0);
            }
        }
    }

    // Wenn kein Header entschlüsselt werden konnte, wirf eine Excption
    if (tmpData.length() == 0)
    {
        throw CryptException("Fehler beim Entschlüsseln des RSA Blockes",
                             CryptException::DecryptionErrorRsa);
    }

    // Ließ die AES Daten aus
    aesKey = tmpData.left(aesKey.length());
    aesIv = tmpData.mid(aesKey.length(), aesIv.length());

    // Erzeuge den symmetrischen Kontext
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();

    // Überprüfe den Schlüssel
    if (!ctx)
    {
        throwOpenSslException();
    }

    // Setze den Verschlüsselungsalgorithmus, sowie den Schlüssel und den IV
    if (1 != EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL,
                                (unsigned char *) aesKey.data(), (unsigned char *) aesIv.data()))
    {
        // Räume den symmetrischen Schlüsselkontext ab
        freeCipherCtx(&ctx);
        throwOpenSslException();
    }

    QByteArray inBuf(EVP_CIPHER_block_size(EVP_aes_256_cbc()), 0);
    QByteArray outBuf(EVP_CIPHER_block_size(EVP_aes_256_cbc()), 0);
    int bufRead;
    int len;

    while (!infileStream.atEnd())
    {
        // Hole jeweils einen Block an Daten aus der Datei
        bufRead = infileStream.readRawData(inBuf.data(), inBuf.size());

        // Entschlüssele den Block mit AES
        if (1 != EVP_DecryptUpdate(ctx, (unsigned char *) outBuf.data(), &len,
                                   (unsigned char *) inBuf.data(), bufRead))
        {
            // Räume den symmetrischen Schlüsselkontext ab
            freeCipherCtx(&ctx);
            throwOpenSslException();
        }

        // Schreibe den Block in die Ausgabedatei
        outfileStream.writeRawData(outBuf.data(), len);
    }

    // Wirft Fehler, falls Padding falsch oder der letzte Block unvollständig war
    if (1 != EVP_DecryptFinal_ex(ctx, (unsigned char *) outBuf.data(), &len))
    {
        // Räume den symmetrischen Schlüsselkontext ab
        freeCipherCtx(&ctx);
        throwOpenSslException();
    }

    outfileStream.writeRawData(outBuf.data(), len);

    // Räume den symmetrischen Schlüsselkontext ab
    freeCipherCtx(&ctx);
}

void HybridCrypt::createKeypair()
{
    throwExceptionIfCsprngIsNotSeeded();

    // Räume den alten Schlüssel vorher ab
    freeEvpKey(&userKeypair);

    // Lege Schlüsselkontextvariable an
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, NULL);

    // Erzeuge den Schlüsselkontext
    if (!ctx)
    {
        throwOpenSslException();
    }

    // Initialisiere den Schlüsselgenerator
    if (EVP_PKEY_keygen_init(ctx) <= 0)
    {
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Lege den Schlüssel mit 2048 Bit an
    if (EVP_PKEY_CTX_set_rsa_keygen_bits(ctx, 2048) <= 0)
    {
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Erzeuge den Schlüssel
    if (EVP_PKEY_keygen(ctx, &userKeypair) <= 0)
    {
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Räume den Schlüsselkontext ab
    freePkeyCtx(&ctx);

}

void HybridCrypt::importUserKeypair(QString keyfileName, QString password)
{
    // Räume den alten Schlüssel vorher ab
    freeEvpKey(&userKeypair);

    FILE *keyfile = fopen(keyfileName.toStdString().c_str(), "r");

    // Datei existiert nicht
    if (keyfile == NULL)
    {
        throw CryptException("Datei nicht gefunden: " + keyfileName.toStdString(),
                             CryptException::FileNotFound);
    }

    // Ließ den Schlüssel des Nutzers anhand der Erweiterung ein
    QString keyfileExtension = keyfileName.right(3).toLower();

    if (keyfileExtension == "pem")
    {
        userKeypair = PEM_read_PrivateKey(keyfile, NULL, NULL,
                                          (void *) password.toStdString().c_str());
    }
    else if (keyfileExtension == "der")
    {
        userKeypair = d2i_PrivateKey_fp(keyfile, NULL);
    }
    else
    {
        // Räume die Ressourcen auf
        fclose(keyfile);
        throw CryptException("Dateiformat wird nicht unterstützt.",
                             CryptException::FormatNotSupported);
    }

    // Räume die Ressourcen auf
    fclose(keyfile);

    // Wirf einen Fehler, falls OpenSSL ein Problem hatte
    if (userKeypair == NULL)
    {
        throwOpenSslException();
    }

    // Wirf Exception, falls kein RSA Schlüssel importiert wurde
    throwExceptionIfEvpKeyIsNotRsa(&userKeypair);
}

void HybridCrypt::exportUserKeypair(QString keyfileName, QString password)
{
    throwExceptionIfCsprngIsNotSeeded();

    throwExceptionIfUserKeyIsNull();

    FILE *keyfile = fopen(keyfileName.toStdString().c_str(), "w+");

    // Teste, ob die Datei zum Schreiben geöffnet werden konnte
    if (keyfile == NULL)
    {
        throw CryptException("Konnte Datei nicht schreiben: " +
                             keyfileName.toStdString(), CryptException::FileNotWritable);
    }

    // Wenn password ein Leerstring ist, verschlüssele den privaten Schlüssel nicht
    // Wird nur im PEM-Format exportiert
    int opensslReturnError = PEM_write_PKCS8PrivateKey(keyfile, userKeypair,
                             (password.isEmpty()) ? NULL : EVP_aes_256_cbc(), NULL, 0, NULL,
                             (void *) password.toStdString().c_str());

    // Räume die Ressourcen auf
    fclose(keyfile);

    // Wirf einen Fehler, falls OpenSSL ein Problem hatte
    if (!opensslReturnError)
    {
        throwOpenSslException();
    }
}

void HybridCrypt::exportPublicUserKey(QString keyfileName)
{
    throwExceptionIfUserKeyIsNull();

    FILE *keyfile = fopen(keyfileName.toStdString().c_str(), "w+");

    // Teste, ob die Datei zum Schreiben geöffnet werden konnte
    if (keyfile == NULL)
    {
        throw CryptException("Konnte Datei nicht schreiben: " +
                             keyfileName.toStdString(), CryptException::FileNotWritable);
    }

    int opensslReturnError;

    // Wähle Funktion zum Schreiben des öffentlichen Schlüssels anhand der Dateierweiterung aus
    QString keyfileExtension = keyfileName.right(3).toLower();

    if (keyfileExtension == "pem")
    {
        opensslReturnError = PEM_write_PUBKEY(keyfile, userKeypair);
    }
    else if (keyfileExtension == "der")
    {
        opensslReturnError = i2d_PUBKEY_fp(keyfile, userKeypair);
    }
    else
    {
        // Räume die Ressourcen auf
        fclose(keyfile);
        throw CryptException("Dateiformat wird nicht unterstützt.",
                             CryptException::FormatNotSupported);
    }

    // Räume die Ressourcen auf
    fclose(keyfile);

    // Wirft einen Fehler, falls OpenSSL ein Problem hatte
    if (!opensslReturnError)
    {
        throwOpenSslException();
    }
}

bool HybridCrypt::isUserKeyInitialised()
{
    return userKeypair != NULL;
}

/*
 *  Private Funktionen
 */

QByteArray HybridCrypt::encryptAesData(EVP_PKEY *pkey, QByteArray data)
{
    // Das QByteArray das zurück gegeben wird
    QByteArray out;

    // Wird benutzt um die Länge von out zu bestimmen
    size_t outlen;

    // Lege Schlüsselkontext an
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pkey, NULL);

    // Überprüfe erflogreiche Schlüsselkontexterzeugung
    if (!ctx)
    {
        throwOpenSslException();
    }

    // Initialisiere den Schlüsselkontext
    if (EVP_PKEY_encrypt_init(ctx) <= 0)
    {
        // Räume den Schlüsselkontext ab
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Setze für den Schlüsselkontext das RSA Padding
    if (EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_OAEP_PADDING) <= 0)
    {
        // Räume den Schlüsselkontext ab
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Finde die Blockgröße für RSA raus
    if (EVP_PKEY_encrypt(ctx, NULL, &outlen, (unsigned char *) data.data(),
                         data.length()) <= 0)
    {
        // Räume den Schlüsselkontext ab
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Setze out auf RSA Blockgröße
    out.resize(outlen);

    // Verschlüssele den AES Key und den IV
    if (EVP_PKEY_encrypt(ctx, (unsigned char *) out.data(), &outlen,
                         (unsigned char *) data.data(), data.length()) <= 0)
    {
        // Räume den Schlüsselkontext ab
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Räume den Schlüsselkontext ab
    freePkeyCtx(&ctx);

    return out;
}

QByteArray HybridCrypt::decryptAesData(EVP_PKEY *pkey, QByteArray data)
{
    // Das QByteArray das zurück gegeben wird
    QByteArray out;

    // Wird benutzt um die Länge von out zu bestimmen
    size_t outlen;

    // Lege Schlüsselkontext an
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pkey, NULL);

    // Überprüfe erflogreiche Schlüsselkontexterzeugung
    if (!ctx)
    {
        throwOpenSslException();
    }

    // Initialisiere den Schlüsselkontext
    if (EVP_PKEY_decrypt_init(ctx) <= 0)
    {
        // Räume den Schlüsselkontext ab
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Setze für den Schlüsselkontext das RSA Padding
    if (EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_OAEP_PADDING) <= 0)
    {
        // Räume den Schlüsselkontext ab
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Finde die Blockgröße für RSA raus
    if (EVP_PKEY_decrypt(ctx, NULL, &outlen, (unsigned char *) data.data(),
                         data.length()) <= 0)
    {
        // Räume den Schlüsselkontext ab
        freePkeyCtx(&ctx);
        throwOpenSslException();
    }

    // Setze out auf RSA Blockgröße
    out.resize(outlen);

    // Entschlüssele den AES Key und den IV
    if (EVP_PKEY_decrypt(ctx, (unsigned char *) out.data(), &outlen,
                         (unsigned char *) data.data(), data.length()) <= 0)
    {
        // Wenn der Fehler vom Padding kommt, ignoriere fürs erste
        if (ERR_GET_LIB(ERR_peek_error()) == ERR_LIB_RSA &&
            (ERR_GET_REASON(ERR_peek_error()) == RSA_R_OAEP_DECODING_ERROR ||
             ERR_GET_REASON(ERR_peek_error()) == RSA_R_PADDING_CHECK_FAILED))
        {
            // Lösche Error aus der Schlange
            ERR_get_error();
            out.resize(0);
        }
        else
        {
            // Räume den Schlüsselkontext ab
            freePkeyCtx(&ctx);
            throwOpenSslException();
        }
    }

    // Räume den Schlüsselkontext ab
    freePkeyCtx(&ctx);

    return out;
}

void HybridCrypt::writeRsaHeader(QDataStream *outfileStream, EVP_PKEY *pkey,
                                 QByteArray aesData)
{
    QByteArray data = encryptAesData(pkey, aesData);

    if (data.length() > 0)
    {
        *outfileStream << (qint32) data.length();
        outfileStream->writeRawData(data.data(), data.length());
    }
}

EVP_PKEY *HybridCrypt::readRecipientKey(QString keyfileName)
{
    FILE *tmpKeyfile = fopen(keyfileName.toStdString().c_str(), "r");

    // Datei existiert nicht
    if (tmpKeyfile == NULL)
    {
        throw CryptException("Datei nicht gefunden: " +
                             keyfileName.toStdString(), CryptException::FileNotFound);
    }

    EVP_PKEY *tmpKey;

    // Entscheide welche Funktion zu benutzen ist
    QString keyfileExtension = keyfileName.right(3).toLower();

    if (keyfileExtension == "pem")
    {
        tmpKey = PEM_read_PUBKEY(tmpKeyfile, NULL, NULL, NULL);
    }
    else if (keyfileExtension == "der")
    {
        tmpKey = d2i_PUBKEY_fp(tmpKeyfile, NULL);
    }
    else
    {
        // Räume die Ressourcen auf
        fclose(tmpKeyfile);
        throw CryptException("Dateiformat wird nicht unterstützt.",
                             CryptException::FormatNotSupported);
    }

    // Räume die Ressourcen auf
    fclose(tmpKeyfile);

    // Wirf Fehler, falls der Schlüssel nicht gelesen werden konnte
    if (tmpKey == NULL)
    {
        throwOpenSslException();
    }

    throwExceptionIfEvpKeyIsNotRsa(&tmpKey);

    return tmpKey;
}

bool HybridCrypt::isCsprngSeeded()
{
    return RAND_status() == 1;
}

QByteArray HybridCrypt::getCsprngBytes(int count)
{
    QByteArray out(count, 0);

    if (1 != RAND_bytes((unsigned char *) out.data(), count))
    {
        throwOpenSslException();
    }

    return out;
}

bool HybridCrypt::isKeyRsa(EVP_PKEY *key)
{
    return EVP_PKEY_type(key->type) == EVP_PKEY_RSA;
}

void HybridCrypt::throwOpenSslException()
{
    QString errorMsg("OpenSSL Fehler. Fehlermeldung: ");
    errorMsg.append(ERR_error_string(ERR_get_error(), NULL));
    throw CryptException(errorMsg.toStdString(), CryptException::OpenSslError);
}

void HybridCrypt::throwExceptionIfEvpKeyIsNotRsa(EVP_PKEY **key)
{
    if (!isKeyRsa(*key))
    {
        freeEvpKey(key);
        throw CryptException("Nur RSA Schlüssel werden unterstüzt.",
                             CryptException::KeyNotRsa);
    }
}

void HybridCrypt::throwExceptionIfCsprngIsNotSeeded()
{
    if (!isCsprngSeeded())
    {
        throw new CryptException("Zufallszahlengenerator ist nicht initialisiert",
                                 CryptException::CsprngNotSeeded);
    }
}

void HybridCrypt::throwExceptionIfUserKeyIsNull()
{
    if (userKeypair == NULL)
    {
        throw CryptException("Kein Schlüssel wurde angelegt",
                             CryptException::NoUserKeyCreated);
    }
}

void HybridCrypt::freeEvpKey(EVP_PKEY **key)
{
    EVP_PKEY_free(*key);
    *key = NULL;
}

void HybridCrypt::freeCipherCtx(EVP_CIPHER_CTX **ctx)
{
    EVP_CIPHER_CTX_free(*ctx);
    *ctx = NULL;
}

void HybridCrypt::freePkeyCtx(EVP_PKEY_CTX **ctx)
{
    EVP_PKEY_CTX_free(*ctx);
    *ctx = NULL;
}