#include <boost/filesystem.hpp>
#include <boost/algorithm/string/trim.hpp>
#include <boost/scoped_array.hpp>

#include <string>
#include <fstream>
#include <iomanip>
#include <iostream>

#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/engine.h>


using namespace std;

enum KeySize
{
    Key_1024,
    Key_2048
};

string getKeySize(const KeySize keySize)
{
    switch (keySize)
    {
        case Key_1024: return "1024";
        case Key_2048: return "2048";
    }
}

string privateKey(const KeySize size, const int keyIndex);
string publicKey(const KeySize size, const int keyIndex);

RSA* createPrivateRSA(const string& key);
RSA* createPublicRSA(const string& key);

bool RSASign(RSA *rsa,
             const unsigned char *Msg,
             size_t MsgLen,
             unsigned char **EncMsg,
             size_t *MsgLenEnc);

bool RSAVerifySignature(RSA *rsa,
                        unsigned char *MsgHash,
                        size_t MsgHashLen,
                        const char *Msg,
                        size_t MsgLen,
                        bool *Authentic);

bool testSignVerify(const string& privateKey, const string& publicKey, const string& plainText);

void runTest(const string& text, const KeySize keySize, const int keyIndex)
{
    bool success = testSignVerify(privateKey(keySize, keyIndex), publicKey(keySize, keyIndex), text);
    cout << getKeySize(keySize) << " bit [" << keyIndex << "] -  Text: " << text << " - " << (success ? "Success" : "Failure") << endl;
}

int main (int argc, char** argv)
{
    // Run simple tests using first 2048 bit keys
    runTest("A", Key_2048, 1);
    runTest("AA", Key_2048, 1);
    runTest("AAA", Key_2048, 1);

    // Run simple tests using second 2048 bit keys
    runTest("A", Key_2048, 2);
    runTest("AA", Key_2048, 2);
    runTest("AAA", Key_2048, 2);

    // 1024 bit key seg faults
//    runTest("A", Key_1024, 2);

    return 0;
}

RSA* createPrivateRSA(const string& key)
{
    RSA *rsa = NULL;
    const char *c_string = key.c_str();
    BIO *keybio = BIO_new_mem_buf((void *) c_string, -1);
    if (keybio == NULL)
    {
        return 0;
    }
    rsa = PEM_read_bio_RSAPrivateKey(keybio, &rsa, NULL, NULL);
    BIO_free_all(keybio);
    return rsa;
}

RSA* createPublicRSA(const string& key)
{
    RSA *rsa = NULL;
    BIO *keybio;
    const char *c_string = key.c_str();
    keybio = BIO_new_mem_buf((void *) c_string, -1);
    if (keybio == NULL)
    {
        return 0;
    }
    rsa = PEM_read_bio_RSA_PUBKEY(keybio, &rsa, NULL, NULL);
    BIO_free_all(keybio);
    return rsa;
}

bool testSignVerify(const string& privateKey, const string& publicKey, const string& plainText)
{
    // Create an RSA instance of the private key to be used in the hash generation
    RSA *privateRSA = createPrivateRSA(privateKey);

    boost::scoped_array<unsigned char> pHash;
    unsigned char* hash;
    size_t hashLength;

    // Generate a hash using <hashAlgorithm> and assign dynamically allocated buffer to scoped pointer
    RSASign(privateRSA,
            (unsigned char *) plainText.c_str(), plainText.length(),
            &hash, &hashLength);
    pHash.reset(hash);

    string hashString((char*)pHash.get());
    hashString = hashString.substr(0, hashLength);


    // Verify with public key
    RSA *publicRSA = createPublicRSA(publicKey);

    if (publicRSA == NULL)
    {
        cerr << "Invalid public key" << endl;
        return false;
    }

    bool authentic;
    bool result = RSAVerifySignature(publicRSA,
                                     (unsigned char*)hashString.c_str(), hashLength,
                                     plainText.c_str(), plainText.length(),
                                     &authentic);

    return result & authentic;
}


bool RSASign(RSA *rsa,
             const unsigned char *Msg,
             size_t MsgLen,
             unsigned char **EncMsg,
             size_t *MsgLenEnc)
{
    EVP_MD_CTX *m_RSASignCtx = EVP_MD_CTX_create();
    EVP_PKEY *priKey = EVP_PKEY_new();
    EVP_PKEY_assign_RSA(priKey, rsa);

    EVP_PKEY_CTX *pctx = NULL;

    int ret = EVP_DigestSignInit(m_RSASignCtx, &pctx, EVP_sha512() , NULL, priKey);
    if (ret <= 0)
    {
        EVP_MD_CTX_destroy(m_RSASignCtx);
        EVP_PKEY_free(priKey);
        cerr << "EVP_DigestSignInit failed: " << ret << endl;
        return false;
    }

    ret = EVP_PKEY_CTX_set_rsa_padding(pctx, RSA_PKCS1_PADDING);
    if (ret <= 0)
    {
        EVP_MD_CTX_destroy(m_RSASignCtx);
        EVP_PKEY_free(priKey);
        cerr << "EVP_PKEY_CTX_set_rsa_padding failed: " << ret << endl;
        return false;
    }
    ret = EVP_DigestSignUpdate(m_RSASignCtx, Msg, MsgLen);
    if (ret <= 0)
    {
        EVP_MD_CTX_destroy(m_RSASignCtx);
        EVP_PKEY_free(priKey);
        cerr << "EVP_DigestSignUpdate failed: " << ret << endl;
        return false;
    }
    if (EVP_DigestSignFinal(m_RSASignCtx, NULL, MsgLenEnc) <= 0)
    {
        EVP_MD_CTX_destroy(m_RSASignCtx);
        EVP_PKEY_free(priKey);
        cerr << "EVP_DigestSignFinal failed" << endl;
        return false;
    }

    *EncMsg = new unsigned char[*MsgLenEnc];

    if (EVP_DigestSignFinal(m_RSASignCtx, *EncMsg, MsgLenEnc) <= 0)
    {
        EVP_MD_CTX_destroy(m_RSASignCtx);
        EVP_PKEY_free(priKey);
        cerr << "EVP_DigestSignFinal failed" << endl;
        return false;
    }

    EVP_MD_CTX_destroy(m_RSASignCtx);
    EVP_PKEY_free(priKey);

    return true;
}

bool RSAVerifySignature(RSA *rsa,
                        unsigned char *MsgHash,
                        size_t MsgHashLen,
                        const char *Msg,
                        size_t MsgLen,
                        bool *Authentic)
{
    *Authentic = false;
    EVP_PKEY *pubKey = EVP_PKEY_new();

    EVP_PKEY_assign_RSA(pubKey, rsa);
    EVP_MD_CTX *m_RSAVerifyCtx = EVP_MD_CTX_create();

    EVP_PKEY_CTX *pctx = NULL;

    int ret = EVP_DigestVerifyInit(m_RSAVerifyCtx, &pctx, EVP_sha512(), NULL, pubKey);
    if (ret <= 0)
    {
        cerr << "EVP_DigestVerifyInit failed: " << ret << endl;
        return false;
    }
    ret = EVP_PKEY_CTX_set_rsa_padding(pctx, RSA_PKCS1_PADDING);
    if (ret <= 0)
    {
        cerr << "EVP_PKEY_CTX_set_rsa_padding failed: " << ret << endl;
        return false;
    }
    ret = EVP_DigestVerifyUpdate(m_RSAVerifyCtx, Msg, MsgLen);
    if (ret <= 0)
    {
        cerr << "EVP_DigestVerifyUpdate failed: " << ret << endl;
        return false;
    }
    int AuthStatus = EVP_DigestVerifyFinal(m_RSAVerifyCtx, MsgHash, MsgHashLen);

    EVP_MD_CTX_destroy(m_RSAVerifyCtx);
    EVP_PKEY_free(pubKey);

    if (AuthStatus == 1)
    {
        *Authentic = true;
        return true;
    }
    else if (AuthStatus == 0)
    {
        const char* file = NULL;
        int line;
        const char* data = NULL;
        int flags;

        // Log errors
        unsigned long ret = ERR_get_error_line_data(&file, &line, &data, &flags);
        cerr << "Authentication failure: " << ret << ", " << file << ", " << line << ", " << data << ", " << flags << " " << __LINE__ << endl;

        char errBuf[1000];
        cerr << "Error string: " << ERR_error_string(ret, errBuf) << endl;

        *Authentic = false;
        return true;
    }
    else
    {
        cout << "Authentication failure: " << ERR_get_error() << " " << __LINE__ << endl;
        *Authentic = false;
        return false;
    }
}

string privateKey(const KeySize keySize, const int keyIndex)
{
    switch (keySize)
    {
        // 1024 bit keys
        case Key_1024:
            switch (keyIndex)
            {
                case 1:
                    return "-----BEGIN RSA PRIVATE KEY-----\n"
                            "MIICWwIBAAKBgQCh+J7AEuTVemaND4PraLV4KaRS8p4dYoCeWGZ7teCjumMmSwGR\n"
                            "kjnu3B9vzNO/59z5KXii3dY2EIlb1dTn0Cwam696TWMvSfKspVKrgsPhJ3p5TzKA\n"
                            "9W2iP2bk7u9bNKCOfey0SQRbqVzYA8NL1CDC8ony/gKJ/F/yPkr9dDMLUQIDAQAB\n"
                            "AoGAU5+09p3Kn50uTS3xghzxdDSrrles4qb0mMwzPTPN4zpb7m1TRgEk//mOEX7C\n"
                            "5JWlNM8AhbXa3tby1AM2hOipm80IcV+XF3uNdul2AqXZKOOCDUdzEiV8uxaIaSuI\n"
                            "bYIRl98aaVLz2awG0zTj+Jf6lYUVmB+JOGJr8x5WDY48wUECQQD3Dtv5RBPE1ipf\n"
                            "H/vrBU0SfTXouRD6TiTRlkPhxcL8VsmHfEaGK5XocrVnaE3cirOD1ioj+y/inqqg\n"
                            "JS5UJw71AkEAp9VhSV/ZMT5ir+bf7MGlBkoJQKRmUpfnka5feHcXhSJhuRPZXkWY\n"
                            "gqiXCX4GrvWul2Ifyl9MdFTdIg7W/dTZbQJAHXjOUPwIB7EsuNBN8wlc2lnJdout\n"
                            "NM7hoYb1bhdzXTrVK65tk0/bwDD+5ukQcMCzDYUc3dH3MmePHy+QSUWniQJAHlgU\n"
                            "+9gpg5685wK+D6c/N/Dg1WAj4+D1DEDToAsgJID/3B97mHDKzyy+CghQGwfuZg5m\n"
                            "owcx4oKJbKP0OPZnzQJAdTYkS84V3LhIy5zz8ZvF1pi/8xVsgjClHINEmTzX/zat\n"
                            "MQ34nMNULKVvZ3JCpxfc/aA91a";

                default:
                    return "";
            }

        // 2048 bit keys
        case Key_2048:
            switch (keyIndex)
            {
                case 1:
                    return "-----BEGIN RSA PRIVATE KEY-----\n"
                            "MIIEogIBAAKCAQEAsuTQfBSGr1Ru8hcbGV7uRwB7P3MA3ErB3SErNtok0PAmbVfC\n"
                            "G99uk3nm4mS5SaB60aWbUzf4R8naexZcyKqAaqnMrGsd3eMjzEY8XR7aGTzM00LK\n"
                            "umG6lyaYr1z7rqNRso9IQmebIu9W3jpvKgw9wvmf4W2tuaviTZ0f/tHVB3VW4Wby\n"
                            "xcgch04hhDkHX1dp1CSFZXRDSGp2MJQ3FAHB9rX+3UGjy2sjWLQqcoIs7KztI5/j\n"
                            "gaOoQGSQj7gBM0+LNFoUXFDju/u7MYeob1KEBWvmQxa8XZlHcr6sGwIruC12cBjv\n"
                            "7jxMlrjILxB2Vr+UQb3PDAlsESuahw6DTooi4QIDAQABAoIBAGgIaLW/uhN8Joji\n"
                            "pv+6FRd1UDc1yepSDytpfpgHhCF9httPjbJn3QYmff5qVNZcXRBOHmXmxpv45vVt\n"
                            "q8eGX+enlz8ti9DdppPExJXiuzKyCOtSjFKXQstjxqc0i6ijkX8llAAKAE3wobPa\n"
                            "h4sMU15vjGqsPbscGh7IoTKvFBOhUEr+6LmdfbdfmQXNGsi7GGLqiuvZa8yFBUTd\n"
                            "bG80Qwm+r+NVesz0WubjY7QK0GOC4eZH3SDgTCH/3326dq18lRMmjm+Ia3ofXs6/\n"
                            "RvSusCERtoAbxxShkYHeUrO9hRCwAvkNZXE0dNGWas7p+53uHxUkTup4phXdkRFf\n"
                            "8zEf4AECgYEA20ksmk2UfxW8tpGj8mU16c0DsEE/yAj5dzyj00RsVmYOBm0IQn9+\n"
                            "fzLv/RKpuITrQFGhtA4UedetkTH92/ed5URDp+MW8qRNUScvRP/8HDCGC+Taiawg\n"
                            "MPCaEzzbHXZXPwc8axhZR8hHhve17faaJHsajK5JfXn62h02IQNVVXECgYEA0Nhk\n"
                            "/LzxQU6MyzT1lqhN/H6eRIw1nL1JU5MyDAc/9tzzPn87+NEWt71QdyiBzUvmLfDh\n"
                            "JRNH8myA0j5ivZN6LNFhNs3C5UDgwknEXDUoKPYzuEl6RNp4L9lN8YTXW/jUJl44\n"
                            "GIMsfVmmzOUa1pADfszZLuv+J6sn5GyGdYYwLHECgYBaNFSruMQgLwHtHCYT3Tx7\n"
                            "S8F7rKufDsoIGaLOknP6aorUa9+c8zNHAd4k/ZthcxQO5kWderlDkcwtUSI7bTW/\n"
                            "jQBDvK/A6oN/oZ3ZRXAg7HkHxYF8ZpjraN4Ikd0U1ox3OfPJnj+Dvshlpdpp5Cpu\n"
                            "XVGDk2JZwmbFeIS+uc9NkQKBgHJyQ1UlDWhmbU1mZjb7YFE/6edSloagTDCad1CU\n"
                            "IH9dsYbMHsdJfZOPEStbbVZdg3QljpvMpx64SxfrzBxrVwZ8bzOjKmyLUx/lYM/O\n"
                            "6YUY9lrFGVXmCvrdsJLr40CpaUjlVQRJzfWbTYb2/fAa12+DPnkyJgMDDkdlGYIJ\n"
                            "2RhhAoGAezeE0TOk/U1g7I3E/pYP3mtXE0Ok4Dy+3QxpA0pB9Wl0YfljceskFbq3\n"
                            "1rtG+2Nnle3x1Y0dmr8/5JvCORBQqRVxOaLdnmZBlCjd+/zM+b6yzlHiAZkrNKij\n"
                            "87fydFOTSAVGYcyd522CiMKFsKNMv4FL6DX/RZaKI+7m4WNaBYg=\n"
                            "-----END RSA PRIVATE KEY-----";

                case 2:
                    return "-----BEGIN RSA PRIVATE KEY-----\n"
                            "MIIEogIBAAKCAQBd9JlX4B1X6j6FMimSYyW7uMgjC3VtYPgvSp2++1jPJU0OXU7C\n"
                            "kv6sKfSS1bUslTD1Vk6bmhacrTmVFdpLMWYhgUnSTTX9jcb1PJGvvLttHDsQooNC\n"
                            "fTDTrE7W0u0bpwD6J882aEHVm+WER3nvYwKz3mLMA+J1ajiT8RmmSuHrUy1u/0iy\n"
                            "OWTEgT/0mx0H3Pl2XktPsRM94XZTxj6Jf8RWLT830uicNkfjMGC221I+Civ0rO0k\n"
                            "y270dWrcG74kVeR10jEmMvQNtlIKUscMQFhTu+57mMuiCeD5VnK1LHaW+N23bzks\n"
                            "ccMKPeGlOFw586l/jvwM7FygooL18N0yOw5zAgMBAAECggEAE3c71fL+4VoMvffY\n"
                            "mRWlmbD89hXHrHOst8+Twx+8FJ67ykAWs0VVMkI3A7z3HleCnzHUc3SHPsq85WBz\n"
                            "q6R+bm6k5sdMak3XbLOx02Ecy8xPiHDDFSFu3kwSwXtN2K7PCwT4b2NGUIVOFgvg\n"
                            "/XHb+3wVqAMeBgmN00jecY+OIv0GHoVgfOseIEZ4EhrdqqPBsOcKSSelNhto4Zc2\n"
                            "jJ6QGAtrYqlEbIth2ohpx5vKltQlnrQH6a6jPg8qc/7WTD4DiF0gDpJWeDwH/jbM\n"
                            "2Htz5STO8qo1DGMwU7dOgJeP1KHZjH3M0QpIUJLOQ32JnP1PS3dh7mprRIR+eOGB\n"
                            "0Ea8IQKBgQCqXbMX2R84rJauZi2BisM9Yo3tMo2U17dRaXw2qb6YxuweI2tyw4Xa\n"
                            "B6FafEDDlJiqAoJFtl6IS0NrMo/zjJfNO7L7Jsi24rJe/t9rGPXRM5viDtORXBo6\n"
                            "ceCRKiSsLDrNyuLGT5Q2yL3ixuconAIKqjobyYjISnjaxU3sH9rUmwKBgQCNLpGN\n"
                            "+zxHQTIE0a6B20S3M1yeg/hF/hQkSwIxu69IeF/yi+N/17ZhJtGmCQaWThAZZzgR\n"
                            "BPIRhZdZKMlTTUNkmMtgw62Q3n8P2Vv/d00pCFsY/dloJxYjD9HYy/ZmORZCCBZ5\n"
                            "JOegymWVi1NUhrLBeO8dA8TVoRfXEqiyq0WPCQKBgQCWXCZBaQZuLZy7hUCqcT7J\n"
                            "gZ8DdreQYybO+ot/bybhjcc8/kPJwRj0igF1a+5470Ah5B73yfH61npFNRXpfO9J\n"
                            "WmQn114UU1XcE5Q54Iirv6HXzQAa7a/UX0c0t34wWsT3dhXZssh75YX5jVTc+Xuv\n"
                            "z/eDKMbJBzlM1BBm0uJ++wKBgEnJaPFXuS1UeZvjKt1biBJnmaSPPJb4L/a37TGo\n"
                            "8fdiUmMlkYUjbIJiE3g/h13O9sMQG1Fp0EowE+XZVz9eRjaxVLoCCxiH8MiLxXVD\n"
                            "GOS1JUduhM3Wf89bWUWBdAoBFAz2xSgTjauPnIAIDUlEDPDpD6+OZLVyOE+DxP98\n"
                            "XCJhAoGACLEIDSzHlJmyNa1Ejciidsb7x816hBE7N56SDMj8HcTveU6W2usYmW+x\n"
                            "v53dvftct3nIsNM6zvrxIT7TWraiC4yahR1ro8fDlgYerJSB5CXS7HHR+AqIgjuR\n"
                            "TQLc9h404JCmoowuvXxCPVhS+OrTt8R2AXVBMQtCj2Nl55EbUwA=\n"
                            "-----END RSA PRIVATE KEY-----";

                default:
                    return "";
            }
    }
}

string publicKey(const KeySize keySize, const int keyIndex)
{
    switch (keySize)
    {
        // 1024 bit keys
        case Key_1024:
            switch (keyIndex)
            {
                case 1:
                    return "-----BEGIN PUBLIC KEY-----\n"
                            "MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCh+J7AEuTVemaND4PraLV4KaRS\n"
                            "8p4dYoCeWGZ7teCjumMmSwGRkjnu3B9vzNO/59z5KXii3dY2EIlb1dTn0Cwam696\n"
                            "TWMvSfKspVKrgsPhJ3p5TzKA9W2iP2bk7u9bNKCOfey0SQRbqVzYA8NL1CDC8ony\n"
                            "/gKJ/F/yPkr9dDMLUQIDAQAB\n"
                            "-----END PUBLIC KEY-----";
            }

        // 2048 bit keys
        case Key_2048:
            switch (keyIndex)
            {
                // 2048 bit - 1
                case 1:
                    return "-----BEGIN PUBLIC KEY-----\n"
                            "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAsuTQfBSGr1Ru8hcbGV7u\n"
                            "RwB7P3MA3ErB3SErNtok0PAmbVfCG99uk3nm4mS5SaB60aWbUzf4R8naexZcyKqA\n"
                            "aqnMrGsd3eMjzEY8XR7aGTzM00LKumG6lyaYr1z7rqNRso9IQmebIu9W3jpvKgw9\n"
                            "wvmf4W2tuaviTZ0f/tHVB3VW4Wbyxcgch04hhDkHX1dp1CSFZXRDSGp2MJQ3FAHB\n"
                            "9rX+3UGjy2sjWLQqcoIs7KztI5/jgaOoQGSQj7gBM0+LNFoUXFDju/u7MYeob1KE\n"
                            "BWvmQxa8XZlHcr6sGwIruC12cBjv7jxMlrjILxB2Vr+UQb3PDAlsESuahw6DTooi\n"
                            "4QIDAQAB\n"
                            "-----END PUBLIC KEY-----";

                    // 2048 bit - 2
                case 2:
                    return "-----BEGIN PUBLIC KEY-----\n"
                            "MIIBITANBgkqhkiG9w0BAQEFAAOCAQ4AMIIBCQKCAQBd9JlX4B1X6j6FMimSYyW7\n"
                            "uMgjC3VtYPgvSp2++1jPJU0OXU7Ckv6sKfSS1bUslTD1Vk6bmhacrTmVFdpLMWYh\n"
                            "gUnSTTX9jcb1PJGvvLttHDsQooNCfTDTrE7W0u0bpwD6J882aEHVm+WER3nvYwKz\n"
                            "3mLMA+J1ajiT8RmmSuHrUy1u/0iyOWTEgT/0mx0H3Pl2XktPsRM94XZTxj6Jf8RW\n"
                            "LT830uicNkfjMGC221I+Civ0rO0ky270dWrcG74kVeR10jEmMvQNtlIKUscMQFhT\n"
                            "u+57mMuiCeD5VnK1LHaW+N23bzksccMKPeGlOFw586l/jvwM7FygooL18N0yOw5z\n"
                            "AgMBAAE=\n"
                            "-----END PUBLIC KEY-----";

                default:
                    return "";
            }
    }
}