import os
import json
import csv
import requests
import time
import sys
import git

class PyExploitDb:
    def __init__(self):
        self.cveToExploitMap = {}
        self.currentPath = os.path.dirname(os.path.abspath(__file__))
        self.edbidToCveFile = self.currentPath + "/edbidToCve.json"
        self.cveToEdbidFile = self.currentPath + "/cveToEdbid.json"
        self.exploitDbPath = self.currentPath + "/exploit-database"
        self.requestCoolOffTime = 1
        self.debug = False
        self.autoUpdate = True
        pass


    def openFile(self, exploitMap = "cveToEdbid.json", encoding="utf-8"):
        if not os.path.isdir(self.exploitDbPath):
            print("Cloning exploit-database repository")
            git.Repo.clone_from("https://github.com/offensive-security/exploit-database.git", self.exploitDbPath)
            print("Updating db...")
            self.updateDb()
        else:
            if self.autoUpdate == True:
                print("Pulling exploit-database updates...")
                git.Git(self.exploitDbPath).pull('origin', 'master')
                print("Updating db...")
                self.updateDb()
            print("Loading database...")
            with open(self.currentPath + "/" + exploitMap, encoding="utf-8") as fileData:
                cveToExploitMap = json.load(fileData)
                self.cveToExploitMap = cveToExploitMap
                if self.debug == True:
                    print(self.cveToExploitMap)


    def getCveDetails(self, cveSearch):
        files = open(self.currentPath + "/exploit-database/files_exploits.csv", encoding="utf-8")
        reader = csv.reader(files)
        next(reader)
        result = {}
        found = False
        for row in reader:
            edb, fileName, description, date, author, platform, exploitType, port = tuple(row)
            if edb in self.cveToExploitMap[cveSearch]:
                found = True
                result['edbid'] = edb
                result['exploit'] = self.exploitDbPath + "/" + fileName
                result['date'] = date
                result['author'] = author
                result['platform'] = platform
                result['type'] = exploitType
                if self.debug == True:
                    print("Exploit DB Id: {0}".format(edb))
                    print("File: {0}".format(self.exploitDbPath + "/" + fileName))
                    print("Date: {0}".format(date))
                    print("Author: {0}".format(author))
                    print("Platform: {0}".format(platform))
                    print("Type: {0}".format(exploitType))
                if port != "0":
                    result['port'] = port
                    if self.debug == True:
                        print("Port: {0}".format(port))
        if not found:
            if self.debug == True:
                print("ERROR - No EDB Id found")
        files.close()
        return result


    def searchCve(self, cveSearch):
        if not cveSearch:
            return []
        cveSearch = cveSearch.upper()
        print(cveSearch)
        if cveSearch in self.cveToExploitMap:
            if self.debug == True:  
                print("Found")
            cveData = self.getCveDetails(cveSearch)
            if cveData:
                return cveData
            else:
                return cveSearch
        return []


    def updateDb(self):
        data = {}
        if not os.path.exists(self.edbidToCveFile):
            os.system("touch {0}".format(self.edbidToCveFile))
            data = {}
        else: 
            with open(self.edbidToCveFile, encoding="utf-8") as fileData:
                try:
                    data = json.load(fileData)
                except:
                    print("Possibly corrupt or empty: {0}".format(self.edbidToCveFile))
                    os.system("rm -f {0}".format(self.edbidToCveFile))
                    os.system("touch {0}".format(self.edbidToCveFile))
                    data = {}
        files = open(self.exploitDbPath + "/files_exploits.csv", encoding="utf-8")
        reader = csv.reader(files)
        next(reader)
        reader = list(reader)
        edbCount = len(reader)
        headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36'}
        def locations_of_substring(string, substring):
           import re
           results = [m.start() for m in re.finditer(substring, string)]
           return results
        print("Refreshing EDBID-CVE mapping. This may take a long time.")
        for i in range(edbCount):
            percentDone = (i / edbCount) * 100
            if self.debug == True:
                print("Processing...{0}%".format(str(percentDone)))
            edb = tuple(reader[i])[0]
            if edb in data:
                if self.debug == True:
                    print("Skipping {0}".format(str(edb)))
                pass
            else:
                content = ""
                while True:
                    try:
                        requestUri = "https://www.exploit-db.com/exploits/{0}".format(str(edb))
                        if self.debug == True:
                            print("Requesting {0}".format(requestUri))
                        r = requests.get(requestUri, headers = headers, timeout=10)
                        content = r.content.decode("ISO-8859-1")
                    except Exception as e:
                        if self.debug == True:
                            print("Error {0}".format(e))
                        time.sleep(self.requestCoolOffTime)
                        continue
                    finally:
                        break
                indexes = locations_of_substring(content, 'https://nvd.nist.gov/vuln/detail/CVE-')
                used = []
                for pos in indexes:
                      cve = r.content[pos + 33: pos + 33 + 14]
                      if type(cve) == type(bytes()):
                          cve = cve.decode("ISO-8859-1")
                      cve = cve.replace('\"', '')
                      cve = cve.replace('\r', '')
                      cve = cve.replace('\n', '')
                      if cve in used:
                          continue
                      used.append(cve)
                data[edb] = used
        with open(self.edbidToCveFile, "w", encoding="utf-8") as fileData:
            json.dump(data, fileData, indent = 2)
        self.cveToExploitMap = {}
        for k, v in data.items():
            for e in v:
                self.cveToExploitMap[e] = self.cveToExploitMap.get(e, [])
                self.cveToExploitMap[e].append(k)
        with open(self.cveToEdbidFile, "w", encoding="utf-8") as fileData:
            json.dump(self.cveToExploitMap, fileData, indent = 2)


def test():
    pEdb = PyExploitDb()
    pEdb.debug = False
    pEdb.openFile()
    results = pEdb.searchCve("CVE-2018-14592")
    if results:
        print(results)
        print('PASS')
    else:
        print('FAIL')
        sys.exit(1)


if __name__ == "__main__":
    test()
