#!/usr/bin/env python
# -*- coding: utf-8 -*-

__license__ = """
GoLismero 2.0 - The web knife - Copyright (C) 2011-2014

Golismero project site: https://github.com/golismero
Golismero project mail: contact@golismero-project.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, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
"""

from golismero.api.data.vulnerability import Vulnerability
from golismero.api.data.vulnerability.vuln_utils import \
    convert_references_to_vuln_ids, convert_vuln_ids_to_references
from golismero.api.logger import Logger
from golismero.api.plugin import TestingPlugin

import re
import requests


#------------------------------------------------------------------------------
class ExploitDBPlugin(TestingPlugin):
    """
    This plugin searches for public exploits for any vulnerability found.
    """


    #--------------------------------------------------------------------------
    def get_accepted_types(self):
        return [Vulnerability]


    #--------------------------------------------------------------------------
    def run(self, info):

        # Skip if the vulnerability doesn't have CVE or OSVDB tags.
        if not info.cve and not info.osvdb:
            Logger.log_more_verbose("No CVE or OSVDB tags found, skipped.")
            return

        # Search for public exploits in the exploit-db database.
        ids = self.query_exploitdb(info.cve, info.osvdb)

        # Add any previously existing IDs.
        ids.update(info.edb)

        # Are there any new IDs?
        new_ids = ids.difference(info.edb)
        if new_ids:

            # Log the new IDs.
            msg = "\n".join(
                "    " + x
                for x in sorted( convert_vuln_ids_to_references(new_ids) )
            )
            Logger.log("Public exploits found:\n" + msg)

            # Update the Vulnerability object.
            info.edb = ids

            # Return the updated object.
            return info


    #--------------------------------------------------------------------------
    # Unofficial Exploit-DB search API.

    URL_CVE = (
        "http://www.exploit-db.com/search/?action=search&filter_page=%d&"
        "filter_cve=%s"
    )
    URL_OSVDB = (
        "http://www.exploit-db.com/search/?action=search&filter_page=%d&"
        "filter_osvdb=%s"
    )
    REGEXP = re.compile(
        r'href=\"(https?\:\/\/www\.exploit\-db\.com\/exploits\/[0-9]+/?)\"'
    )

    def query_exploitdb(self, cve=(), osvdb=()):

        # We'll store all the EDB-IDs here.
        ids = set()

        # For each CVE value...
        for n in cve:

            # Remove the prefix.
            if n.startswith("CVE-"):
                n = n[4:]

            # Query the database and collect all the Exploit-DB IDs.
            ids.update( self.__query_exploitdb(self.URL_CVE, n) )

        # For each OSVDB value...
        for n in osvdb:

            # Remove the prefix.
            if n.startswith("OSVDB-"):
                n = n[6:]

            # Query the database and collect all the Exploit-DB IDs.
            ids.update( self.__query_exploitdb(self.URL_OSVDB, n) )

        # Return all the EDB-IDs.
        return ids

    def __query_exploitdb(self, url_template, query_string):

        # This is where we'll collect the IDs.
        collected = set()

        # For each page of results...
        page = 1
        while True:

            # Build the request URL.
            url = url_template % (page, query_string)

            # Get the result page. Stop on any errors.
            try:
                result = requests.get(url).content
            except Exception:
                break

            # Extract all URLs from the result page.
            urls = self.REGEXP.findall(result)

            # Extract all Exploit-DB IDs from the URLs.
            ids = convert_references_to_vuln_ids(urls).get("EDB", [])

            # If no more IDs were found, stop.
            if not ids:
                break

            # Collect the IDs.
            collected.update(ids)

            # Go to the next page.
            page += 1

        # Return the collected IDs.
        return collected
