#include "downloader.h"

#ifndef DOWNLOAD_SERVER
#define DOWNLOAD_SERVER "https://www.intevation.de"
#endif

#include <QFile>
#include <QDir>
#include <QDebug>
#include <QStandardPaths>
#include <QUuid>
#include <QApplication>
#include <QTextStream>
#include <QLocale>
#include <QSaveFile>

#include <polarssl/net.h>
#include <polarssl/ssl.h>
#include <polarssl/entropy.h>
#include <polarssl/ctr_drbg.h>
#include <polarssl/error.h>
#include <polarssl/certs.h>

#define MAX_SW_SIZE 10485760
#define MAX_LIST_SIZE 1048576
#define MAX_IO_TRIES 10

#define LIST_RESOURCE "/incoming/aheinecke/test"
#define SW_RESOURCE "/incoming/aheinecke/test"

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

#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);
}

Downloader::Downloader(QObject* parent, const QString& url,
                       const QByteArray& certificate,
                       const QDateTime& newestSW,
                       const QDateTime& newestList):
    QThread(parent),
    mUrl(url),
    mPinnedCert(certificate),
    mLastModSW(newestSW),
    mLastModList(newestList),
    mErrorState(NoError),
    mInitialized(false),
    mServerFD(-1)
{
    int ret = -1;

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

    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 Downloader::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;
}

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

QString Downloader::getDataDirectory()
{
    QString candidate =
        QStandardPaths::writableLocation(QStandardPaths::DataLocation);

    if (candidate.isEmpty()) {
        qDebug() << "Could not find writeable locaction for me";
        return QString();
    }

    QDir cDir(candidate);

    if (!cDir.exists()) {
        if (!cDir.mkpath(candidate)) {
            qDebug() << "Could not create path to: " << candidate;
            return QString();
        }
    }
    return cDir.absolutePath();
}

int Downloader::establishSSLConnection() {
    int ret = -1;
    const x509_crt *peerCert;

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

    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;
        }
    }

    /* 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;
}

/* Helper around polarssl bare bone api */
int polarSSLWrite (ssl_context *ssl, 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();

    qDebug() << "Seinding 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(ssl, 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) {
            return ret;
        }
        tries++;
        net_usleep(100000); /* sleep 100ms to give the socket a chance
                               to clean up. */
    } while (tries < MAX_IO_TRIES);

    return ret;
}

/* Helper around polarssl bare bone api read at most len bytes
 * and return them as a byte array returns a NULL byte array on error*/
QByteArray polarSSLRead (ssl_context *ssl, size_t len)
{
    unsigned char buf[len];
    QByteArray retval("");
    int ret = -1;

    do {
        memset (buf, 0, sizeof(buf));
        ret = ssl_read(ssl, buf, len);
        if (ret == 0 ||
            ret == POLARSSL_ERR_SSL_CONN_EOF) {
            /* EOF */
            return retval;
        }
        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, len);
    } while (len > 0);

    return retval;
}


QDateTime Downloader::getLastModifiedHeader(const QString &resource) {
    int ret = -1;
    QByteArray response;
    QTextStream responseStream(&response);
    QLocale cLocale = QLocale::c();
    QString headRequest =
        QString::fromLatin1("HEAD %1 HTTP/1.1\r\n"
                "Connection: Keep-Alive\r\n"
                "\r\n\r\n").arg(resource);

    ret = polarSSLWrite (&mSSL, headRequest.toUtf8());
    if (ret != 0) {
        mErrorState = ConnectionLost;
        return QDateTime();
    }

    response = polarSSLRead(&mSSL, 1024);

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

    while (1) {
        QString line = responseStream.readLine();
        if (line.isNull()) {
            break;
        }
        if (line.startsWith("Last-Modified:")) {
            QDateTime candidate = cLocale.toDateTime(line, "'Last-Modified: 'ddd, dd MMM yyyy HH:mm:ss' GMT'");
            qDebug() << "Parsed line : " << line << " to  " << candidate;
            if (candidate.isValid()) {
                return candidate;
            }
        }
    }

    mErrorState = InvalidResponse;
    return QDateTime();
}

bool Downloader::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.1\r\n\r\n").arg(resource);

    QSaveFile outputFile(fileName);

    ret = polarSSLWrite (&mSSL, getRequest.toUtf8());

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

    if (ret != 0) {
        mErrorState = ConnectionLost;
        return false;
    }

    do {
        QByteArray response = polarSSLRead(&mSSL, 8192);
        if (response.isNull()) {
            qDebug() << "Error reading response";
            mErrorState = ConnectionLost;
            return false;
        }
        if (response.isEmpty()) {
            /* We have read everything there is to read */
            break;
        }

        outputFile.write(response);
        bytesRead += response.size();
    } while (bytesRead < maxSize);

    return outputFile.commit();
}

void Downloader::run() {
    int ret;
    QDateTime remoteModList;
    QDateTime remoteModSW;

    if (!mInitialized) {
        emit error(tr("Failed to initialize SSL Module."), ErrUnknown);
        return;
    }

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

    if (ret != 0) {
        mErrorState = NoConnection;
        emit error(tr("Failed to connect to %1.").arg(mUrl.host()),
                mErrorState);
        return;
    }

    emit progress(tr("Connected"), 1, -1);

    ret = establishSSLConnection();
    if (ret != 0) {
        qDebug() << "SSL conncetion failed: " << getErrorMsg(ret);
        emit error(tr("Failed to connect to %1.").arg(mUrl.host()),
                mErrorState);
        return;
    }

    remoteModSW = getLastModifiedHeader(QString::fromLatin1("/incoming/aheinecke/test"));
    remoteModList = getLastModifiedHeader(QString::fromLatin1("/incoming/aheinecke/test"));

    if (!remoteModSW.isValid() || !remoteModList.isValid()) {
        qDebug() << "Could not read headers";
        return;
    }

    if (!mLastModSW.isValid() || remoteModSW > mLastModSW) {
        QString dataDirectory = getDataDirectory();

        if (dataDirectory.isEmpty()) {
            qDebug() << "Failed to get data directory";
            return;
        }

        QString fileName = dataDirectory.append("/SW-")
            .append(remoteModSW.toString("yyyymmddHHmmss"))
            .append(".exe");

        qDebug() << "fileName: " << fileName;

        if (!downloadFile(QString::fromLatin1(SW_RESOURCE),
                   fileName, MAX_SW_SIZE)) {
            return;
        }

        emit newSoftwareAvailable(fileName, remoteModSW);
    } else if (!mLastModList.isValid() || remoteModList > mLastModList) {
        QString dataDirectory = getDataDirectory();

        if (dataDirectory.isEmpty()) {
            qDebug() << "Failed to get data directory";
            return;
        }

        QString fileName = dataDirectory.append("/list-")
            .append(remoteModSW.toString("yyyymmddHHmmss"))
            .append(".txt");

        qDebug() << "fileName: " << fileName;

        if (!downloadFile(QString::fromLatin1(LIST_RESOURCE),
                   fileName, MAX_LIST_SIZE)) {
            return;
        }

        emit newListAvailable(fileName, remoteModList);
    }

    emit progress(tr("Closing"), 1, -1);
    ssl_close_notify (&mSSL);

    if (mServerFD != -1) {
        net_close(mServerFD);
        mServerFD = -1;
    }
}