/* Copyright (C) 2014 by Bundesamt für Sicherheit in der Informationstechnik
 * Software engineering by Intevation GmbH
 *
 * This file is Free Software under the GNU GPL (v>=2)
 * and comes with ABSOLUTELY NO WARRANTY!
 * See LICENSE.txt for details.
 */
/* TODO: Wrap ssl_session in a class for reuse.
 * see programs/ssl/ssl_client2.c for example of session reuse */
#include "sslconnection_bare.h"
#include "sslhelp.h"

#include <QFile>
#include <QSaveFile>
#include <QUuid>
#include <QApplication>

#define MAX_IO_TRIES 10
#define MAX_RESETS 10

#ifdef CONNECTION_DEBUG
static void my_debug(void *ctx, int level, const char *str)
{
    fprintf((FILE *) ctx, "%s", str);
    fflush((FILE *) ctx);
}
#endif

SSLConnectionBare::SSLConnectionBare(const QString& url,
                             const QByteArray& certificate):
    SSLConnection (url, certificate)
{
    int ret = -1;

    memset(&mSSL, 0, sizeof(ssl_context));
    memset(&mSavedSession, 0, sizeof( ssl_session ) );

    ret = init();
    if (ret == 0) {
        mInitialized = true;
    } else {
        qDebug() << "Initialization error: " + getPolarSSLErrorMsg(ret);
    }
}

int SSLConnectionBare::init()
{
    int ret = -1;
    QUuid uuid = QUuid::createUuid();
    QString personalString = QApplication::applicationName() + uuid.toString();
    QByteArray personalBa = personalString.toLocal8Bit();

    x509_crt_init(&mX509PinnedCert);
    entropy_init(&mEntropy);

    ret = ssl_init(&mSSL);
    if (ret != 0) {
        /* The only documented error is malloc failed */
        mErrorState = ErrUnknown;
        return ret;
    }

    /*
     * Initialize random generator.
     * Personalisation string, does not need to be random but
     * should be unique according to documentation.
     *
     * the ctr_drbg structure does not need to be freed explicitly.
     */
    ret = ctr_drbg_init(&mCtr_drbg, entropy_func, &mEntropy,
                        (const unsigned char*) personalBa.constData(),
                        personalBa.size());
    if (ret != 0) {
        ssl_free(&mSSL);
        mErrorState = ErrUnknown;
        return ret;
    }

    ret = x509_crt_parse(&mX509PinnedCert,
                         (const unsigned char*) mPinnedCert.constData(),
                         mPinnedCert.size());
    if (ret != 0){
        ssl_free(&mSSL);
        mErrorState = InvalidPinnedCertificate;
        return ret;
    }

    ssl_set_endpoint(&mSSL, SSL_IS_CLIENT);
    ssl_set_authmode(&mSSL, SSL_VERIFY_OPTIONAL);
    ssl_set_ca_chain(&mSSL, &mX509PinnedCert, NULL, NULL);
    ssl_set_renegotiation(&mSSL, SSL_RENEGOTIATION_DISABLED);
    ssl_set_rng(&mSSL, ctr_drbg_random, &mCtr_drbg);
#ifdef RELEASE_BUILD
    ssl_set_min_version(&mSSL, SSL_MAJOR_VERSION_3, SSL_MINOR_VERSION_3);
#endif

#ifdef CONNECTION_DEBUG
    ssl_set_dbg(&mSSL, my_debug, stdout);
#endif

    return 0;
}

SSLConnectionBare::~SSLConnectionBare() {
    disconnect();
    x509_crt_free(&mX509PinnedCert);
    entropy_free(&mEntropy);
    if (mInitialized) {
        ssl_free(&mSSL);
    }
}

void SSLConnectionBare::disconnect() {
    if (mConnected) {
        ssl_close_notify(&mSSL);
        if (mServerFD != -1) {
            net_close(mServerFD);
            mServerFD = -1;
        }
        ssl_session_free(&mSavedSession);
        mConnected = false;
    }
}

int SSLConnectionBare::connect() {
    int ret = -1;

    if (!mInitialized) {
        mErrorState = ErrUnknown;
        return -1;
    }

    ret = net_connect(&mServerFD, mUrl.host().toLatin1().constData(),
                  mUrl.port(443));

    if (ret != 0) {
        qDebug() << "Connect failed: " << getPolarSSLErrorMsg(ret);
        mErrorState = NoConnection;
        return ret;
    }

    ssl_set_bio(&mSSL, net_recv, &mServerFD,
                       net_send, &mServerFD);

    while ((ret = ssl_handshake(&mSSL)) != 0) {
        if (ret != POLARSSL_ERR_NET_WANT_READ &&
                ret != POLARSSL_ERR_NET_WANT_WRITE) {
            qDebug() << "SSL Handshake failed: " << getPolarSSLErrorMsg(ret);
            mErrorState = SSLHandshakeFailed;
            return ret;
        }
    }

    ret = ssl_get_session(&mSSL, &mSavedSession);
    if (ret != 0) {
        qDebug() << "SSL get session failed: " << getPolarSSLErrorMsg(ret);

        mErrorState = NoConnection;
        return ret;
    }
    printf( " ok\n    [ Ciphersuite is %s ]\n",
            ssl_get_ciphersuite( &mSSL) );
    ret = validateCertificate();

    if (ret == 0) {
        mConnected = true;
    }
    return ret;
}

int SSLConnectionBare::validateCertificate()
{
    int ret = -1;
    const x509_crt *peerCert = NULL;

    /* we might want to set the verify function
     * with ssl_set_verify before to archive the
     * certificate pinning. */

    ret = ssl_get_verify_result(&mSSL);

    if (ret != 0 ) {
        if((ret & BADCERT_EXPIRED) != 0)
           qDebug() << "server certificate has expired";
        if((ret & BADCERT_REVOKED) != 0)
           qDebug() << "server certificate has been revoked";
        if((ret & BADCERT_CN_MISMATCH) != 0)
           qDebug() << "CN mismatch";
        if((ret & BADCERT_NOT_TRUSTED) != 0)
            qDebug() << "self-signed or not signed by a trusted CA";
#ifdef RELEASE_BUILD
        mErrorState = InvalidCertificate;
        return -1;
#endif
    }

    peerCert = ssl_get_peer_cert(&mSSL);

    if (!peerCert) {
        mErrorState = InvalidCertificate;
        qDebug() << "Failed to get peer cert";
        return -1;
    }

    if (peerCert->raw.len == 0 ||
        peerCert->raw.len != mX509PinnedCert.raw.len) {
        mErrorState = InvalidCertificate;
        qDebug() << "Certificate length mismatch";
        return -1;
    }

    /* You can never be sure what those c++ operators do..
    if (mPinnedCert != QByteArray::fromRawData(
                (const char*) peerCert->raw.p,
            peerCert->raw.len)) {
        qDebug() << "Certificate content mismatch";
    }
    */

    for (unsigned int i = 0; i < peerCert->raw.len; i++) {
        if (peerCert->raw.p[i] != mX509PinnedCert.raw.p[i]) {
            qDebug() << "Certificate content mismatch";
            mErrorState = InvalidCertificate;
            return -1;
        }
    }
    return 0;
}

int SSLConnectionBare::write (const QByteArray& request)
{
    unsigned int tries = 0;
    int ret = -1;

    const unsigned char *buf = (const unsigned char *) request.constData();
    size_t len = (size_t) request.size();

    if (mNeedsReset) {
        ret = reset();
        if (ret != 0) {
            qDebug() << "Reset failed: " << getPolarSSLErrorMsg(ret);
            return ret;
        }
    }

    qDebug() << "Sending request: " << request;
    /* According to doc for ssl_write:
     *
     * When this function returns POLARSSL_ERR_NET_WANT_WRITE,
     * it must be called later with the same arguments,
     * until it returns a positive value.
     */
    do {
        ret = ssl_write(&mSSL, buf, len);
        if (ret >= 0) {
            if ((unsigned int) ret == len) {
                return 0;
            } else {
                qDebug() << "Write failed to write everything";
                return -1;
            }
        }
        if (ret != POLARSSL_ERR_NET_WANT_WRITE &&
                ret != POLARSSL_ERR_NET_WANT_READ) {
            return ret;
        }
        tries++;
        net_usleep(100000); /* sleep 100ms to give the socket a chance
                               to clean up. */
    } while (tries < MAX_IO_TRIES);

    return ret;
}


int SSLConnectionBare::reset()
{
    int ret = -1;
    ssl_close_notify(&mSSL);

    ret = ssl_session_reset(&mSSL);
    if (ret != 0)
    {
        qDebug() << "SSL Connection reset failed: "
                 << getPolarSSLErrorMsg(ret);
        return ret;
    }

    ssl_set_session(&mSSL, &mSavedSession);

    ret = net_connect(&mServerFD, mUrl.host().toLatin1().constData(),
                  mUrl.port(443));

    if (ret != 0) {
        mErrorState = NoConnection;
        qDebug() << "Connection failed." << getPolarSSLErrorMsg(ret);
        return ret;
    }

    while ((ret = ssl_handshake(&mSSL)) != 0) {
        if (ret != POLARSSL_ERR_NET_WANT_READ &&
                ret != POLARSSL_ERR_NET_WANT_WRITE) {
            qDebug() << "SSL Handshake failed: "
                 << getPolarSSLErrorMsg(ret);
            mErrorState = SSLHandshakeFailed;
            return ret;
        }
    }

    qDebug() << "Reset connection. ";
    /* Validation should not be necessary as we reused a saved
     * session. But just to be sure. */
    return validateCertificate();
}

QByteArray SSLConnectionBare::read(size_t len)
{
    unsigned char buf[len];
    QByteArray retval("");
    int ret = -1;
    unsigned int tries = 0;

    mNeedsReset = true;
    do {
        memset (buf, 0, sizeof(buf));
        ret = ssl_read(&mSSL, buf, len);
        if (ret == 0 ||
            ret == POLARSSL_ERR_SSL_CONN_EOF ||
            ret == POLARSSL_ERR_SSL_PEER_CLOSE_NOTIFY) {
            /* EOF */
            return retval;
        }
        if (ret == POLARSSL_ERR_NET_WANT_WRITE ||
                ret == POLARSSL_ERR_NET_WANT_READ) {
            net_usleep(100000); /* sleep 100ms to give the socket a chance
                                   to recover */
            tries++;
        }
        if (ret <= 0) {
            qDebug() << "Read failed: " << getPolarSSLErrorMsg(ret);
            return QByteArray();
        }
        if (len < (len - (unsigned int) ret)) {
            /* Should never happen if ssl_read behaves */
            qDebug() << "integer overflow in polarSSLRead";
            return QByteArray();
        }
        len -= (unsigned int) ret;
        retval.append((const char *)buf, ret);
    } while (len > 0 && tries < MAX_IO_TRIES);

    return retval;
}

QMap<QString, QString> SSLConnectionBare::parseHeaders(QByteArray *data)
{
    int bodyStart = data->indexOf("\r\n\r\n");
    QMap<QString, QString> retval;
    QByteArray headers;
    QString response(*data);
    if (bodyStart == -1) {
        qDebug() << "Could not find header end.";
        return retval;
    }

    /* Take the headers with one additional line break */
    headers = data->left(bodyStart + 2);
    /* Chop off the head */

    foreach (const QString& line, response.split("\r\n")) {
        int sepPos = -1;
        sepPos = line.indexOf(": ");
        if (sepPos == -1) {
            continue;
        }
        QString key = line.left(sepPos);
        QString value = line.right(line.size() - sepPos - 2);

        retval.insert(key, value);
    }

    *data = data->right(data->size() - bodyStart - 4);
    return retval;
}

QDateTime SSLConnectionBare::getLastModifiedHeader(const QString &resource) {
    int ret = -1;
    QByteArray response;
    QLocale cLocale = QLocale::c();
    QMap<QString, QString> headers;
    QString headRequest =
        QString::fromLatin1("HEAD %1 HTTP/1.0\r\n\r\n").arg(resource);

    ret = write(headRequest.toUtf8());
    if (ret != 0) {
        return QDateTime();
    }

    response = read(1024);

    qDebug() << "Response from server was: " << response;

    if (response.isNull()) {
        qDebug() << "No response";
        return QDateTime();
    }

    headers = parseHeaders(&response);
    const QString lastModified = headers.value("Last-Modified");
    if (!lastModified.isEmpty()) {
        QDateTime candidate = cLocale.toDateTime(lastModified,
                "ddd, dd MMM yyyy HH:mm:ss' GMT'");
        if (candidate.isValid()) {
            return candidate;
        }
    }
    return QDateTime();
}

bool SSLConnectionBare::downloadFile(const QString &resource,
                              const QString &fileName,
                              size_t maxSize)
{
    int ret = -1;
    size_t bytesRead = 0;
    QString getRequest =
        QString::fromLatin1("GET %1 HTTP/1.0\r\n\r\n").arg(resource);

    QSaveFile outputFile(fileName);

    ret = write(getRequest.toUtf8());

    if (ret != 0) {
        qDebug() << "Failed to send request.";
        return false;
    }

    // Open / Create the file to write to.
    if (!outputFile.open(QIODevice::WriteOnly)) {
        qDebug() << "Failed to open file";
        return false;
    }

    bool inBody = false;
    QMap <QString, QString> headers;
    do {
        /* Read the response in 8KiB chunks */
        int responseSize = 0;
        QByteArray response = read(8192);
        if (response.isNull()) {
            qDebug() << "Error reading response";
            return false;
        }
        responseSize = response.size();
        if (!inBody) {
            headers = parseHeaders(&response);
            inBody = true;
        }
        outputFile.write(response);
        bytesRead += responseSize;
        if (responseSize < 8192) {
            /* Nothing more to read */
            break;
        }
    } while (bytesRead < maxSize);

    return outputFile.commit();
}

void SSLConnectionBare::setCiphersuites(int ciphers[]) {
    if (mInitialized) {
        ssl_set_ciphersuites(&mSSL, ciphers);
    }
}