/*
Copyright 2011-2017  Francesco Cecconi <francesco.cecconi@gmail.com>

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License as
published by the Free Software Foundation; either version 2 of
the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "vulnerability.h"
#include "mainwindow.h"

VulnerabilityWidget::VulnerabilityWidget(QWidget* parent): QWidget(parent)
{
    setupUi(this);
}

Vulnerability::Vulnerability(MainWindow* parent)
    : QObject(parent), m_ui(parent), m_completerVuln(0), m_vulnModel(0)
{

    m_vulnerabilityWidget = new VulnerabilityWidget(m_ui);

    m_vulnerabilityWidget->tWresult->setTabsClosable(true);
    m_vulnerabilityWidget->tWresult->removeTab(0);
    m_vulnerabilityWidget->tWresult->setVisible(false);
    updateCompleter();
    loadSearchUrlIntoCombo();

    m_vulnerabilityWidget->comboVulnRis->setStyleSheet(QString::fromUtf8("color: rgb(153, 153, 153);"));
    m_vulnerabilityWidget->comboVulnRis->insertItem(0, tr("Search Vulnerabilities"));

    connect(m_vulnerabilityWidget->tWresult, &QTabWidget::tabCloseRequested,
            this, &Vulnerability::closeVulnTab);
    connect(m_vulnerabilityWidget->comboVulnRis->lineEdit(), &QLineEdit::returnPressed,
            this, &Vulnerability::searchVulnerabilityFromCombo);
    connect(m_vulnerabilityWidget->comboVulnRis->lineEdit(), &QLineEdit::cursorPositionChanged,
            this, &Vulnerability::updateComboServicesProperties);
    connect(m_vulnerabilityWidget->comboVuln, static_cast<void (QComboBox::*)(const QString&)>(&QComboBox::currentIndexChanged),
            this, &Vulnerability::updateComboVuln);
    connect(m_vulnerabilityWidget->treeVulnNseRecovered, &QTreeWidget::itemClicked,
            this, &Vulnerability::searchVulnFromTreeWidget);

    m_mainHorizontalLeftSplitter = new QSplitter(m_ui);

    m_mainHorizontalLeftSplitter->setOrientation(Qt::Horizontal);
    m_mainHorizontalLeftSplitter->addWidget(m_vulnerabilityWidget->treeVulnNseRecovered);
    m_mainHorizontalLeftSplitter->addWidget(m_vulnerabilityWidget->frameVuln);
    m_mainHorizontalLeftSplitter->addWidget(m_vulnerabilityWidget->frameOptions);

    m_vulnerabilityWidget->layout()->addWidget(m_mainHorizontalLeftSplitter);

    QSettings settings("nmapsi4", "nmapsi4");
    if (!settings.value("vulnTreeSplitterHorizontal").toByteArray().isEmpty()) {
        m_mainHorizontalLeftSplitter->restoreState(settings.value("vulnTreeSplitterHorizontal").toByteArray());
    } else {
        // set a default width
        QList<int> size = m_mainHorizontalLeftSplitter->sizes();
        size[0] = 100; // ratio
        size[1] = 250;
        size[2] = 100;
        m_mainHorizontalLeftSplitter->setSizes(size);
    }

    m_vulnerabilityWidget->comboWebV->setCurrentIndex(settings.value("vulnUrlComboIndex", 0).toInt());

    QSpacerItem *verticalSpacer = new QSpacerItem(20, 163, QSizePolicy::Minimum, QSizePolicy::Expanding);
    m_welcomeQml = new QQuickView;
    QVBoxLayout *qmlWelcomeLayout = new QVBoxLayout(QWidget::createWindowContainer(m_welcomeQml));
    m_welcomeQml->setResizeMode(QQuickView::SizeRootObjectToView);
    qmlWelcomeLayout->addItem(verticalSpacer);
    m_welcomeQml->setSource(Package::qmlPath("vulnerabilityWelcome"));

    m_vulnerabilityWidget->frameVuln->layout()->addWidget(QWidget::createWindowContainer(m_welcomeQml));
    m_vulnerabilityWidget->treeVulnNseRecovered->setVisible(false);
}

Vulnerability::~Vulnerability()
{
}

void Vulnerability::syncSettings()
{
    QSettings settings("nmapsi4", "nmapsi4");
    settings.setValue("vulnTreeSplitterHorizontal", m_mainHorizontalLeftSplitter->saveState());
    settings.setValue("vulnUrlComboIndex", m_vulnerabilityWidget->comboWebV->currentIndex());
}

void Vulnerability::updateCompleter()
{
    // load vulnerability services string model
    if (!m_ui->m_bookmark->isBookmarkServicesListEmpty()) {
        if (m_completerVuln) {
            QStringListModel *newModel = qobject_cast<QStringListModel*>(m_completerVuln->model());
            newModel->setStringList(m_ui->m_bookmark->getServicesListFromBookmark());
        } else if (!m_vulnModel) {
            m_vulnModel = new QStringListModel(m_ui->m_bookmark->getServicesListFromBookmark(), this);
        }
    }
}

const QList< QPair<QString, QString> > Vulnerability::getDefaultUrlList()
{
    QList< QPair<QString, QString> > urlListModel;
    QPair<QString, QString> urlModel;

    // set default name+value url
    urlModel.first = "cve.mitre.org";
    urlModel.second = "http://cve.mitre.org/cgi-bin/cvekey.cgi?keyword=";
    urlListModel.push_back(urlModel);

    urlModel.first = "www.securityfocus.com";
    urlModel.second = "http://www.securityfocus.com/swsearch?sbm=bid&submit=Search%21&metaname=alldoc&sort=swishrank&query=";
    urlListModel.push_back(urlModel);

    urlModel.first = "secunia.com";
    urlModel.second = "http://secunia.com/advisories/search/?search=";
    urlListModel.push_back(urlModel);

    urlModel.first = "web.nvd.nist.gov";
    urlModel.second = "http://web.nvd.nist.gov/view/vuln/search-results?search_type=all&cves=on&query=";
    urlListModel.push_back(urlModel);

    urlModel.first = "osvdb.org";
    urlModel.second = "http://osvdb.org/search?search[vuln_title]=";
    urlListModel.push_back(urlModel);

    return urlListModel;
}

void Vulnerability::searchVulnerabilityFromCombo()
{
    if (m_vulnerabilityWidget->comboVulnRis->currentText().isEmpty()) {
        return;
    }

    // save current search on bookmark
    m_ui->m_bookmark->saveServiceToBookmark(m_vulnerabilityWidget->comboVulnRis->currentText(), m_ui->m_hostCache);
    //update completer cache
    updateCompleter();

    QString finalUrlString;
    QString tmpSearchString;
    // search string from QComboBox
    tmpSearchString = m_vulnerabilityWidget->comboVulnRis->currentText();
    // space to + for correct search string format
    tmpSearchString.replace(QString(" "), QString("+"));

    // create QHash for default profile
    QList< QPair<QString, QString> > urlListModel = getDefaultUrlList();

    // selection url from key for search
    if (m_vulnerabilityWidget->comboWebV->currentIndex() < urlListModel.size()) {
        // It is a default url, stored in QList
        QListIterator< QPair<QString, QString> > i(urlListModel);
        while (i.hasNext()) {
            QPair<QString, QString> urlModel = i.next();
            if (urlModel.first == m_vulnerabilityWidget->comboWebV->currentText()) {
                finalUrlString = urlModel.second;
            }
        }
    } else {
        // It isn't a default url, but an url from treeWidget
        QList<QTreeWidgetItem *> resultList_ = m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->findItems(m_vulnerabilityWidget->comboWebV->currentText(), Qt::MatchExactly, 0);
        finalUrlString = resultList_[0]->text(1);
    }

    finalUrlString.append(tmpSearchString);

    QUrl urlFinal(finalUrlString);
    openTab(urlFinal, m_vulnerabilityWidget->comboVulnRis->lineEdit()->text());
}

void Vulnerability::openTab(const QUrl& address, const QString& tabName)
{
    // default action settings
    m_ui->m_collections->m_collectionsVulnerability.value("stop-act")->setEnabled(true);
    m_ui->m_collections->m_collectionsVulnerability.value("back-act")->setEnabled(true);
    m_ui->m_collections->m_collectionsVulnerability.value("forward-act")->setEnabled(true);

    m_welcomeQml->setVisible(false);
    // make tabBar visible for the first search
    if (!m_webViewList.size()) {
        m_vulnerabilityWidget->tWresult->setVisible(true);
    }

    // create new tab with a QWebEngineView
    QWebEngineView *page = new QWebEngineView(m_ui);
    page->setContextMenuPolicy(Qt::NoContextMenu);
    m_webViewList.push_back(page);
    page->load(address);
    m_vulnerabilityWidget->tWresult->addTab(page, tabName);
    connect(page, &QWebEngineView::loadProgress, m_vulnerabilityWidget->progressWeb, &QProgressBar::setValue);
    connect(page, &QWebEngineView::loadFinished, this, &Vulnerability::vulnPostScan);
}

void Vulnerability::searchVulnFromTreeWidget()
{
    if (!m_vulnerabilityWidget->treeVulnNseRecovered->selectedItems().size()) {
        return;
    }

    QString addressName = m_ui->m_scanWidget->treeMain->selectedItems()[0]->text(0).split(' ', QString::SkipEmptyParts)[0];

    openTab(QUrl::fromUserInput(m_vulnerabilityWidget->treeVulnNseRecovered->selectedItems()[0]->text(0)),
                        addressName
                        + "::"
                        + static_cast<QUrl>(m_vulnerabilityWidget->treeVulnNseRecovered->selectedItems()[0]->text(0)).authority());
}

void Vulnerability::closeVulnTab(int index)
{
    QWebEngineView *page = m_webViewList.takeAt(index);
    m_vulnerabilityWidget->tWresult->removeTab(index);
    delete page;

    if (!m_webViewList.size()) {
        m_vulnerabilityWidget->tWresult->setVisible(false);
        m_welcomeQml->setVisible(true);
        // disable search action in vulnerability toolBar
        m_ui->m_collections->m_collectionsVulnerability.value("search-act")->setEnabled(false);
    }
}

void Vulnerability::tabWebBack()
{
    m_webViewList[m_vulnerabilityWidget->tWresult->currentIndex()]->triggerPageAction(QWebEnginePage::Back);
}

void Vulnerability::tabWebForward()
{
    m_webViewList[m_vulnerabilityWidget->tWresult->currentIndex()]->triggerPageAction(QWebEnginePage::Forward);
}

void Vulnerability::tabWebStop()
{
    m_webViewList[m_vulnerabilityWidget->tWresult->currentIndex()]->triggerPageAction(QWebEnginePage::Stop);
}

void Vulnerability::vulnPostScan()
{
    m_ui->m_collections->m_collectionsVulnerability.value("stop-act")->setEnabled(false);
}

void Vulnerability::callVulnCheck()
{
    m_vulnerabilityWidget->comboVulnRis->clear();
    updateComboServicesProperties();
    m_vulnerabilityWidget->comboVulnRis->insertItem(0, m_ui->m_bookmark->m_vulnBookmarkWidget->treeBookVuln->currentItem()->text(0));
    searchVulnerabilityFromCombo();
}

void Vulnerability::linkCompleterToServices()
{
    if (!m_ui->m_collections->m_collectionsScanSection.value("bookmarkAddService-action")->isEnabled()) {
        m_ui->m_collections->m_collectionsScanSection.value("bookmarkAddService-action")->setEnabled(true);
    }

    if (!m_ui->m_collections->m_collectionsVulnerability.value("search-act")->isEnabled()) {
        m_ui->m_collections->m_collectionsVulnerability.value("search-act")->setEnabled(true);
    }

    if (m_vulnerabilityWidget->comboVulnRis->lineEdit()->text().isEmpty()
        && !m_vulnerabilityWidget->comboVuln->currentIndex()) {
        m_ui->m_collections->m_collectionsVulnerability.value("search-act")->setEnabled(false);
    }


    if (!m_vulnModel) {
        return;
    }

    if (!m_completerVuln) {
        m_completerVuln = new QCompleter(m_vulnModel, m_ui);
        m_completerVuln->setCompletionRole(QCompleter::InlineCompletion);
        m_completerVuln->setWrapAround(false);
        m_completerVuln->setCaseSensitivity(Qt::CaseInsensitive);
        m_vulnerabilityWidget->comboVulnRis->setCompleter(m_completerVuln);
    }
}

void Vulnerability::updateComboVuln(const QString& value)
{
    if (m_vulnerabilityWidget->comboVuln->currentIndex()) {
        m_vulnerabilityWidget->comboVulnRis->clear();
        m_vulnerabilityWidget->comboVulnRis->insertItem(0, value);
        m_ui->m_collections->m_collectionsVulnerability.value("search-act")->setEnabled(true);
    } else {

        if (m_vulnerabilityWidget->comboVulnRis->lineEdit()->text().isEmpty()) {
            m_ui->m_collections->m_collectionsVulnerability.value("search-act")->setEnabled(false);
        }
    }
}

void Vulnerability::checkVulnerabilitiesFromPortsTree()
{
    updateComboServicesProperties();

    int notEmpty = 0;

    foreach (QTreeWidgetItem* item, m_ui->m_scanWidget->listWscan->selectedItems()) {
        if (!item->text(3).isEmpty()) {
            if (m_vulnerabilityWidget->comboVulnRis->itemText(0).isEmpty()) {
                m_vulnerabilityWidget->comboVulnRis->addItem(item->text(3));
                searchVulnerabilityFromCombo();
            } else {
                m_vulnerabilityWidget->comboVulnRis->setItemText(0, item->text(3));
                searchVulnerabilityFromCombo();
            }
            notEmpty++;
        }
    }

    if (notEmpty) {
        m_ui->updateVulnerabilitySection();
    }
}

void Vulnerability::loadSearchUrlIntoCombo()
{
    // clear comboWeb
    m_vulnerabilityWidget->comboWebV->clear();

    // insert default static url in comboWeb
    QListIterator< QPair<QString, QString> > i(getDefaultUrlList());
    while (i.hasNext()) {
        m_vulnerabilityWidget->comboWebV->insertItem(m_vulnerabilityWidget->comboWebV->count() + 1, i.next().first);
    }

    m_vulnerabilityWidget->comboWebV->insertSeparator(m_vulnerabilityWidget->comboWebV->count() + 1);

    // value from treeWidget url
    for (int index = 0; index < m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->topLevelItemCount(); index++) {
        m_vulnerabilityWidget->comboWebV->insertItem(m_vulnerabilityWidget->comboWebV->count() + 1,
                m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->topLevelItem(index)->text(0));
    }
}

void Vulnerability::showAddUrlUi()
{
    QPointer<AddVulnerabilityUrl> dialogUrl = new AddVulnerabilityUrl(m_ui);

    connect(dialogUrl, &AddVulnerabilityUrl::doneUrl,
            this, &Vulnerability::addUrlToBookmarks);

    dialogUrl->exec();

    if (dialogUrl) {
        delete dialogUrl;
    }
}

void Vulnerability::addUrlToBookmarks(const QString urlName, const QString urlAddr)
{
    m_ui->m_bookmark->saveAddressToBookmark(urlName, urlAddr);
    // reload comboVuln with website
    loadSearchUrlIntoCombo();
}

void Vulnerability::removeUrlToBookmarks()
{
    if (m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->selectedItems().isEmpty()) {
        return;
    }

    //remove address from bookmark
    m_ui->m_bookmark->deleteAddressFromBookmark(m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->selectedItems()[0]->text(0));
    loadSearchUrlIntoCombo();
}

void Vulnerability::updateComboServicesProperties()
{
    m_vulnerabilityWidget->comboVulnRis->clear();
    m_ui->m_collections->m_collectionsVulnerability.value("search-act")->setEnabled(true);
    m_vulnerabilityWidget->comboVulnRis->setStyleSheet(QString::fromUtf8(""));
    bool signalState = m_vulnerabilityWidget->comboVulnRis->lineEdit()->disconnect(SIGNAL(cursorPositionChanged(int,int)));

    if (!signalState)
        return;

    connect(m_vulnerabilityWidget->comboVulnRis, &QComboBox::editTextChanged,
            this, &Vulnerability::linkCompleterToServices);
}

void Vulnerability::openUrlFromScanPortsTree()
{
    QString address;
    foreach (QTreeWidgetItem* item, m_ui->m_scanWidget->listWscan->selectedItems()) {
        if (item->text(2).contains("http") && !item->text(2).contains("ssl")) {
            address.append("http://");
            address.append(m_ui->m_scanWidget->treeMain->currentItem()->text(0).split(' ')[0]);
            address.append(':' + item->text(0).split('/')[0]);
            // open tab for the address
            openTab(QUrl(address), address);
        }

        address.clear();
    }

    m_ui->updateVulnerabilitySection();
}
