/*****************************************************************************
* checkbrute Checksum Bruteforcing Tool
* Copyright (C) 2020 Syping
*
* Redistribution and use in source and binary forms, with or without modification,
* are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice,
* this list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* This software is provided as-is, no warranties are given to you, we are not
* responsible for anything with use of the software, you are self responsible.
*****************************************************************************/

// Qt includes
#include <QCoreApplication>
#include <QTextStream>
#include <QFile>

// checkbrute includes
#include "mainthread.h"

mainthread::mainthread(const QString &bruteforceFile, const QString &checksumsFile, const QString &algorithms, const QString &start, const QString &stop, const QString &lenghts_str, const QString &threads_str, bool strictHash, const std::map<std::string,checksum*> &checksum_map) : checksum_map(checksum_map), strictHash(strictHash)
{
    QFile bruteforceFileHandler(bruteforceFile);
    if (bruteforceFileHandler.open(QIODevice::ReadOnly)) {
        QTextStream(stderr) << "INFO: Reading source file " << bruteforceFile << "..." << Qt::endl;
        fileContent = bruteforceFileHandler.readAll();
        bruteforceFileHandler.close();
        QTextStream(stderr) << "INFO: Source file have " << fileContent.size() << " bytes" << Qt::endl;
    }
    else {
        QTextStream(stderr) << "ERROR: Source file " << bruteforceFile << " can't be read!" << Qt::endl;
        initialised = false;
        return;
    }

    QFile checksumsFileHandler(checksumsFile);
    if (checksumsFileHandler.open(QIODevice::ReadOnly)) {
        QTextStream(stderr) << "INFO: Reading checksums file " << checksumsFile << "..." << Qt::endl;
        QByteArray checksumFileContent = checksumsFileHandler.readAll();
        checksumsFileHandler.close();
        for (const QByteArray &data : checksumFileContent.split('\n')) {
            const QByteArray tdata = data.trimmed();
            if (!tdata.isEmpty()) {
                QByteArray binary = QByteArray::fromHex(tdata);
                if (!binary.isEmpty()) {
                    QTextStream(stderr) << "INFO: Found valid checksum " << tdata.toLower() << Qt::endl;
                    checksums << tdata.toLower();
                }
                else {
                    QTextStream(stderr) << "INFO: Found invalid data " << tdata << Qt::endl;
                }
            }
        }
    }
    else {
        QTextStream(stderr) << "ERROR: Checksums file " << bruteforceFile << " can't be read!" << Qt::endl;
        initialised = false;
        return;
    }
    if (checksums.isEmpty()) {
        QTextStream(stderr) << "ERROR: No checksums are used!" << Qt::endl;
        initialised = false;
        return;
    }

    if (!algorithms.isEmpty()) {
        for (const QString &algorithm : algorithms.split(',')) {
            std::map<std::string,checksum*>::const_iterator it;
            it = checksum_map.find(algorithm.toUpper().toStdString());
            if (it != checksum_map.end()) {
                QTextStream(stderr) << "INFO: Using " << QString::fromStdString(it->first) << " algorithm" << Qt::endl;
                checksum_vector << it->second;
            }
        }
    }
    else {
        std::map<std::string,checksum*>::const_iterator it;
        for (it = checksum_map.begin(); it != checksum_map.end(); it++) {
            const int checksumSize = it->second->checksumSize();
            if (!strictHash) {
                for (const QByteArray &hash : checksums) {
                    if (checksumSize >= hash.count() / 2) {
                        QTextStream(stderr) << "INFO: Using " << QString::fromStdString(it->first) << " algorithm" << Qt::endl;
                        checksum_vector << it->second;
                        break;
                    }
                }
            }
            else {
                for (const QByteArray &hash : checksums) {
                    if (checksumSize == hash.count() / 2) {
                        QTextStream(stderr) << "INFO: Using " << QString::fromStdString(it->first) << " algorithm" << Qt::endl;
                        checksum_vector << it->second;
                        break;
                    }
                }
            }
        }
    }
    if (checksum_vector.isEmpty()) {
        QTextStream(stderr) << "ERROR: No algorithms are used!" << Qt::endl;
        initialised = false;
        return;
    }

    if (!start.isEmpty()) {
        bool startOk;
        begin = start.toULongLong(&startOk);
        if (!startOk) {
            QTextStream(stderr) << "ERROR: Start value is not valid!" << Qt::endl;
            initialised = false;
            return;
        }
    }
    else {
        begin = 0;
    }

    if (!stop.isEmpty()) {
        bool stopOk;
        end = stop.toULongLong(&stopOk);
        if (!stopOk) {
            QTextStream(stderr) << "ERROR: Stop value is not valid!" << Qt::endl;
            initialised = false;
            return;
        }
    }
    else {
        end = (quint64)fileContent.size();
    }

    if (begin >= end) {
        QTextStream(stderr) << "ERROR: Stop value must be placed after Start value!" << Qt::endl;
        initialised = false;
        return;
    }

    // Clear unneeded bytes
    if (begin != 0 || end != (quint64)fileContent.size()) {
        fileContent = fileContent.mid(begin, end - begin);
        fileContent.squeeze();
        QTextStream(stderr) << "INFO: Bruteforce Start " << begin << " End " << end << Qt::endl;
    }

    if (!lenghts_str.isEmpty()) {
        for (const QString &length_str : lenghts_str.split(',')) {
            if (length_str.contains('-')) {
                const QStringList length_list = length_str.split('-');
                if (length_list.count() == 2) {
                    bool lengthOk;
                    quint64 min = length_list.at(0).toULongLong(&lengthOk);
                    if (!lengthOk || min < 1 || min > (quint64)fileContent.size()) {
                        QTextStream(stderr) << "ERROR: Found invalid length \"" << min << "\"!" << Qt::endl;
                        initialised = false;
                        return;
                    }
                    quint64 max = length_list.at(1).toULongLong(&lengthOk);
                    if (!lengthOk || max > (quint64)fileContent.size() || max < 1) {
                        QTextStream(stderr) << "ERROR: Found invalid length \"" << max << "\"!" << Qt::endl;
                        initialised = false;
                        return;
                    }
                    brutelength length;
                    length.min = min;
                    length.max = max;
                    lengths << length;
                }
                else {
                    QTextStream(stderr) << "ERROR: Found invalid length \"" << length_str << "\"!" << Qt::endl;
                    initialised = false;
                    return;
                }
            }
            else {
                bool lengthOk;
                quint64 val = length_str.toULongLong(&lengthOk);
                if (lengthOk && val >= 1) {
                    brutelength length;
                    length.min = val;
                    length.max = val;
                    lengths << length;
                }
                else {
                    QTextStream(stderr) << "ERROR: Found invalid length \"" << length_str << "\"!" << Qt::endl;
                    initialised = false;
                    return;
                }
            }
        }
    }
    else {
        brutelength length;
        length.min = 1;
        length.max = (quint64)fileContent.size();
        lengths << length;
    }
    if (lengths.isEmpty()) {
        QTextStream(stderr) << "ERROR: No lengths are used!" << Qt::endl;
        initialised = false;
        return;
    }

    if (!threads_str.isEmpty()) {
        bool threadsOk;
        quint16 threadCount = threads_str.toUShort(&threadsOk);
        if (!threadsOk || threadCount > 255) {
            QTextStream(stderr) << "ERROR: Threads value is not valid!" << Qt::endl;
            initialised = false;
            return;
        }
        mthreads = (quint8)threadCount;
    }
    else {
        int threadCount = QThread::idealThreadCount();
        if (threadCount < 255) {
            mthreads = (quint8)threadCount;
        }
        else {
            mthreads = 255;
        }
    }

    found = 0;
    threads = 0;
    initialised = true;
    doneThreads = false;
}

void mainthread::threadFinished()
{
    threads--;
    eventLoop.quit();
    sender()->deleteLater();
    if (doneThreads && threads == 0) {
        QCoreApplication::processEvents(QEventLoop::AllEvents);
        QTextStream(stderr) << "INFO: Bruteforce finished, found " << found << " matches!" << Qt::endl;
        emit finished();
    }
}

bool mainthread::isInitialised()
{
    return initialised;
}

void mainthread::matched()
{
    found++;
}

void mainthread::run()
{
    QTextStream(stderr) << "INFO: Set max threads to " << mthreads << Qt::endl;
    QTextStream(stderr) << "INFO: Start bruteforcing..." << Qt::endl;
    QVector<brutelength>::const_iterator it;
    const int lengthsCount = lengths.count();
    for (it = lengths.constBegin(); it != lengths.constEnd(); it++) {
        brutelength length = *it;
        quint64 min = length.min;
        quint64 max = length.max;
        if (min == max) {
            if (threads >= mthreads)
                eventLoop.exec();
            QTextStream(stderr) << "[" << std::distance(lengths.constBegin(), it) + 1 << "/" << lengthsCount << "] [" << min << "/" << max << "] Bruteforcing...\r";
            brutethread *thread = new brutethread(fileContent, min, begin, checksum_vector, checksums, strictHash);
            QObject::connect(thread, &QThread::finished, this, &mainthread::threadFinished);
            QObject::connect(thread, &brutethread::matched, this, &mainthread::matched);
            thread->start(QThread::LowPriority);
            threads++;
        }
        else if (min < max) {
            for (quint64 clength = min; clength <= max; clength++) {
                if (threads >= mthreads)
                    eventLoop.exec();
                QTextStream(stderr) << "[" << std::distance(lengths.constBegin(), it) + 1 << "/" << lengthsCount << "] [" << clength << "/" << max << "] Bruteforcing...\r";
                brutethread *thread = new brutethread(fileContent, clength, begin, checksum_vector, checksums, strictHash);
                QObject::connect(thread, &QThread::finished, this, &mainthread::threadFinished);
                QObject::connect(thread, &brutethread::matched, this, &mainthread::matched);
                thread->start(QThread::LowPriority);
                threads++;
            }
        }
        else {
            for (quint64 clength = min; clength >= max; clength--) {
                if (threads >= mthreads)
                    eventLoop.exec();
                QTextStream(stderr) << "[" << std::distance(lengths.constBegin(), it) + 1 << "/" << lengthsCount << "] [" << clength << "/" << max << "] Bruteforcing...\r";
                brutethread *thread = new brutethread(fileContent, clength, begin, checksum_vector, checksums, strictHash);
                QObject::connect(thread, &QThread::finished, this, &mainthread::threadFinished);
                QObject::connect(thread, &brutethread::matched, this, &mainthread::matched);
                thread->start(QThread::LowPriority);
                threads++;
            }
        }
    }
    doneThreads = true;
}