getExploit.py

#!/usr/bin/env python2

# Made by Giovanny A. Gongora G. (gioyik)

# define
__app__ = 'getExploit'
__author__ = 'Giovanny Andres Gongora Granada'
__date__ = '2015-01-02'
__license__ = 'MIT'

# imports
import cmd
import csv
import os
import re
import shutil
import sys
import zipfile
from copy import copy
from collections import defaultdict
from urllib2 import urlopen

MYSELF = os.path.realpath(__file__)
if os.path.islink(MYSELF): MYSELF = os.readlink(MYSELF)
CURRENT_DIR = os.path.realpath(os.path.dirname(MYSELF))
EXPLOITS_DIR = os.path.join(CURRENT_DIR, 'exploits')
EXPLOITS_CSV = os.path.join(EXPLOITS_DIR, 'exploit-database-master/files.csv')
ARCHIVE_PATH = os.path.join(CURRENT_DIR, 'archive.zip')
ARCHIVE_URL = "https://github.com/offensive-security/exploit-database/archive/master.zip?raw"


class ExploitSearch(cmd.Cmd):
    prompt = '\033[1;31mgetExploit\033[0m\033[1;32m>\033[0m '
    intro = (
        '\n'
        '#[                    getExploit.py                    ]#\n'
        '#[       Search exploits from exploit-db database      ]#\n'
    )
    fields = ('id', 'file', 'description', 'date', 'author', 'platform', 'type', 'port')
    search_regexes = {
        'plain': [
            re.compile(r'(\w+):(?!r[\'"])([^ \'"]+)'),
            re.compile(r'(\w+):(?:\'|")([^"\']+)')
        ],
        'regex': [
            re.compile(r'(\w+):r(?:\'|")([^"\']+)')
        ]
    }
    unquoted_search_re = re.compile(r'(\w+):(?!r[\'"])([^ \'"]+)')
    quoted_search_re = re.compile(r'(\w+):(?:\'|")([^"\']+)')
    regex_search_re = re.compile(r'(\w+):r(?:\'|")([^"\']+)')
    highlighted_fields_map = {
        'id': ['id'],
        'description': ['description'],
        'file': ['type', 'platform']
    }

    def __init__(self, csv_file=EXPLOITS_CSV):
        self.csv_file = csv_file
        self.exploits = []
        self.load_csv()
        self.fields_value_completion = {
            'platform': set([e['platform'] for e in self.exploits]),
            'type': set([e['type'] for e in self.exploits]),
            'port': set([e['port'] for e in self.exploits])
        }
        cmd.Cmd.__init__(self)

    def download_archive(self):
        resp = urlopen(ARCHIVE_URL)
        file_size = int(resp.info().getheaders("Content-Length")[0])
        downloaded_size = 0
        block_size = 4096
        with open(ARCHIVE_PATH, 'wb') as outfile:
            buff = resp.read(block_size)
            while buff:
                outfile.write(buff)
                downloaded_size += len(buff)
                downloaded_part = float(downloaded_size) / file_size
                progress_size = int(downloaded_part * 50)
                status = '[{0}{1}] {2:.2%}'.format(
                    '#' * progress_size,
                    ' ' * (50 - progress_size),
                    downloaded_part
                )
                sys.stdout.write(status)
                sys.stdout.write('\b' * (len(status)+1))
                buff = resp.read(block_size)
            sys.stdout.write('\n')

    def parse_args(self, args):
        search_args = {
            'plain': {},
            'regex': {}
        }
        if ':' not in args:
            search_args['plain']['description'] = args
        else:
            for search_type, regexes in self.search_regexes.iteritems():
                for regex in regexes:
                    result = regex.findall(args)
                    if result is not None:
                        for field_name, pattern in result:
                            search_args[search_type][field_name] = pattern
        return search_args

    def load_csv(self):
        if not os.path.exists(self.csv_file):
            print "Database not found, updating now!\n"
            self.updatedb()
        with open(self.csv_file) as infile:
            reader = csv.reader(infile)
            header = reader.next()
            for entry in reader:
                exploit = dict(zip(header, entry))
                if exploit['port'] == '0':
                    exploit['port'] = 'n/a'
                if not exploit['platform']:
                    exploit['platform'] = 'n/a'
                if '//' in exploit['file']:
                    exploit['file'] = exploit['file'].replace('//', '/')
                self.exploits.append(exploit)

    def search(self, search_params):
        matches = []
        args = self.parse_args(search_params)
        for exploit in self.exploits:
            matching = True
            for search_type, search_args in args.iteritems():
                if search_type == 'plain':
                    for field_name, pattern in search_args.iteritems():
                        if pattern.lower() not in exploit[field_name].lower():
                            matching = False
                            break
                    if not matching:
                        break
                elif search_type == 'regex':
                    for field_name, pattern in search_args.iteritems():
                        if re.search(pattern, exploit[field_name], flags=re.I) is None:
                            matching = False
                            break
                    if not matching:
                        break
            if matching:
                matches.append((exploit, args))
        return matches

    def do_search(self, line):
        """
        search - search database for exploits
        Usage: search field:pattern [field:pattern, ...]
        """
        results = self.search(line)
        for result, args in results:
            flattened_args = {}
            for search_type in ('plain', 'regex'):
                for k, v in args[search_type].iteritems():
                    flattened_args[k] = v
            result = copy(result)
            for field_name, search_vals in self.highlighted_fields_map.iteritems():
                for search_val in search_vals:
                    if search_val in flattened_args:
                        pattern = flattened_args[search_val]
                        result[field_name] = re.sub(
                            pattern if (field_name in args['regex']) else re.escape(pattern),
                            lambda matchobj: '\033[1;33m' + matchobj.group(0) + '\033[0m',
                            result[field_name],
                            flags=re.I
                        )
            result_str = "[%s] %s - %s" % (result['id'], result['description'], result['file'])
            print result_str
        print ''

    def complete_search(self, text, line, begidx, endidx):
        last_arg = line.split()[-1]
        if last_arg:
            if ':' in last_arg:
                field_name, pattern = last_arg.split(':')
                if pattern:
                    return [
                        f for f in self.fields_value_completion[field_name]
                        if f.startswith(pattern)
                    ]
                return [f for f in self.fields_value_completion[field_name]]
            return [f+':' for f in self.fields if f.startswith(last_arg)]
        return [f + ':' for f in self.fields]

    def info(self, exploit_id):
        for exploit in self.exploits:
            if exploit['id'] == exploit_id:
                return exploit
        return None

    def do_info(self, line):
        """
        info - get details about given exploit
        Usage: info [exploitID]
        """
        result = self.info(line)
        if result is None:
            print "No exploits in the database with this ID: %s\n" % line
            return
        desc_len = len(result['description'])
        fstring = "{:-<13} | {:-<%d}" % (desc_len+2)
        print ("{:=^%d}" % (desc_len+18)).format(' #%s ' % result['id'])
        print fstring.format('Filename ', result['file']+' ')
        print fstring.format('Description ', result['description']+' ')
        print fstring.format('Date ', result['date']+' ')
        print fstring.format('Author ', result['author']+' ')
        print fstring.format('Platform ', result['platform']+' ')
        print fstring.format('Type ', result['type']+' ')
        print fstring.format('Port ', result['port']+' ')
        print (desc_len+18)*'='+'\n'

    def complete_info(self, text, line, begidx, endidx):
        if not text:
            return [e['id'] for e in self.exploits]
        else:
            return [e['id'] for e in self.exploits if e['id'].startswith(text)]

    def updatedb(self):
        print "Downloading latest exploits archive..."
        self.download_archive()
        print "Extracting files..."
        if os.path.exists(EXPLOITS_DIR):
            shutil.rmtree(EXPLOITS_DIR)
            os.mkdir(EXPLOITS_DIR)
        with zipfile.ZipFile(ARCHIVE_PATH, 'r') as infile:
            infile.extractall(path=EXPLOITS_DIR)
        os.remove(ARCHIVE_PATH)
        os.chmod(EXPLOITS_CSV, 0644)
        self.exploits = []
        self.load_csv()
        print "OK\n"

    def do_updatedb(self, line):
        """
        updatedb - update local exploits database
        Usage: updatedb
        """
        self.updatedb()

    def do_EOF(self, line):
        return True


def main(restarted=False):
    es = ExploitSearch()
    if restarted:
        es.intro = '\n'
    try:
        es.cmdloop()
    except KeyboardInterrupt:
        main(True)


if __name__ == '__main__':
    main()