/* TODO: Wrap ssl_session in a class for reuse.
 * see programs/ssl/ssl_client2.c for example of session reuse */
#include "sslconnection.h"

#include <QFile>
#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

QString getErrorMsg(int ret)
{
    char errbuf[255];
    polarssl_strerror(ret, errbuf, 255);
    errbuf[254] = '\0'; /* Just to be sure */
    return QString::fromLatin1(errbuf);
}

SSLConnection::SSLConnection(const QString& url,
                             const QByteArray& certificate):
    mUrl(url),
    mPinnedCert(certificate),
    mInitialized(false),
    mConnected(false),
    mServerFD(-1),
    mErrorState(NoError)
{
    int ret = -1;

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

    if (certificate.isEmpty()) {
        QFile certResource(":certs/kolab.org");
        certResource.open(QFile::ReadOnly);
        mPinnedCert = certResource.readAll();
        certResource.close();
    }

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

int SSLConnection::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;
}

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

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

int SSLConnection::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: " << getErrorMsg(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: " << getErrorMsg(ret);
            mErrorState = SSLHandshakeFailed;
            return ret;
        }
    }

    ret = ssl_get_session(&mSSL, &mSavedSession);
    if (ret != 0) {
        qDebug() << "SSL get session failed: " << getErrorMsg(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 SSLConnection::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";
        ret = -1;
#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 SSLConnection::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: " << getErrorMsg(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 SSLConnection::reset()
{
    int ret = -1;
    ssl_close_notify(&mSSL);

    ret = ssl_session_reset(&mSSL);
    if (ret != 0)
    {
        qDebug() << "SSL Connection reset failed: "
                 << getErrorMsg(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." << getErrorMsg(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: "
                 << getErrorMsg(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 SSLConnection::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: " << getErrorMsg(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;
}