#!/usr/bin/python
import subprocess
import os
import hashlib
import string
import re
import stat
from debug import Debug
from data_alchemy import *
from sqlalchemy import and_

rootMessage = """
Root privileges are required to run this program.
"""

longOptions = ['update-db']

debug = Debug()
debug.level = 0

def githash(data):
    s = hashlib.sha1()
    s.update("blob %u\0" %len(data))
    s.update(data)
    return s.hexdigest()

source_windows_base = "./exploitdb/platforms/windows/"
destination_windows_base = "./exploits/windows/"
source_linux_base = "./exploitdb/platforms/linux/"
destination_linux_base = "./exploits/linux/"

script_extensions = ['py', 'pl', 'php']

def convert_dos_file(source_file, destination_file):
    if debug.level > 0:
        debug.msg(source_file + destination_file)

    if source_file != None and destination_file != None:
        source_file_handle = open(source_file)
        destination_file_handle = open(destination_file, 'w')
        destination_file_lines = []
        for line in source_file_handle:
            regExp = re.compile("\r\n|\n|\r")
            destination_file_lines.append(regExp.sub("\n", line))

        destination_file_handle.writelines(destination_file_lines)

def add_exploit_parameter():
    pass

def compile_exploits(os_family, service_name, exploit_type, exploit_source):
    if os_family.lower() == "windows":
        source_base = source_windows_base
        destination_base = destination_windows_base
    elif os_family.lower() == "linux":
        source_base = source_linux_base
        destination_base = destination_linux_base
    
    tmp_out = subprocess.Popen(["grep", "-Rl", service_name, source_base + exploit_type +"/"], stdout=subprocess.PIPE)
    service_out = subprocess.Popen(["grep", "\.c$"], stdin=tmp_out.stdout, stdout=subprocess.PIPE)
    service_files, service_err = service_out.communicate()

    for service_file in service_files.splitlines():
        if service_file != None and os.path.exists(service_file):
            source_filename = os.path.basename(service_file)
            compiled_filename, extension = source_filename.split(".")
            compiled_full_path = destination_base + service_name + "/" + compiled_filename + ".o"
            subprocess.call(["gcc", service_file, "-o", compiled_full_path])
            if os.path.exists(compiled_full_path):
                exploit_path = compiled_full_path
                exploit_file = open(exploit_path, 'r')
                exploit_file_contents = exploit_file.read()
                exploit_githash = githash(exploit_file_contents)
                exploit_sha1 = hashlib.sha1(exploit_file_contents).hexdigest()
                exploit_file.close()
                exploits = session.query(Exploit).filter(Exploit.exploit_sha1 == exploit_sha1).all()
                
                if debug.level > 0:
                    debug.msg(exploits)
                
                if len(exploits) == 0:
                    service_exploit_list = []
                    exploit = Exploit(exploit_sha1 = exploit_sha1, exploit_githash = exploit_githash, source_file = service_file)
                    service_exploit = Service_exploit(exploit_sha1 = exploit_sha1, exploit_path = exploit_path, exploit_source = exploit_source, 
                            os_family = os_family, service_name = service_name)
                    service_exploit_list.append(service_exploit)
                    exploit.service_exploit = service_exploit_list
                    session.add(exploit)
                    session.commit()

                # This is to make the file executable
                os.chmod(exploit_path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)

def copy_exploits(os_family, service_name, exploit_type, exploit_source):
    if debug.level > 0:
        debug.msg(os_family + service_name + exploit_type + exploit_source)

    if os_family.lower() == "windows":
        source_base = source_windows_base
        destination_base = destination_windows_base
    elif os_family.lower() == "linux":
        source_base = source_linux_base
        destination_base = destination_linux_base
    
    tmp_out = subprocess.Popen(["grep", "-Rl", service_name, source_base + exploit_type +"/"], stdout=subprocess.PIPE)
    for script_extension in script_extensions:
        service_out = subprocess.Popen(["grep", "\." + script_extension + "$"], stdin=tmp_out.stdout, stdout=subprocess.PIPE)
        service_files, service_err = service_out.communicate()
        
        if debug.level > 0:
            debug.msg(service_files)
        
        for service_file in service_files.splitlines():
            if service_file != None and os.path.exists(service_file):
                source_filename = os.path.basename(service_file)
                destination_path = destination_base + service_name + "/" + source_filename
                if not os.path.exists(destination_base):
                    os.mkdir(destination_base)
                if not os.path.exists(destination_base + service_name):
                    os.mkdir(destination_base + service_name)

                convert_dos_file(service_file, destination_path)
                if os.path.exists(destination_path):
                    exploit_path = destination_path
                    exploit_file = open(exploit_path)
                    exploit_file_contents = exploit_file.read()
                    exploit_githash = githash(exploit_file_contents)
                    exploit_sha1 = hashlib.sha1(exploit_file_contents).hexdigest()
                    exploit_file.close()
                    exploits = session.query(Exploit).filter(Exploit.exploit_sha1 == exploit_sha1).all()
                    
                    if debug.level > 0:
                        debug.msg(exploits)
                    
                    if len(exploits) == 0:
                        service_exploit_list = []
                        exploit = Exploit(exploit_sha1 = exploit_sha1, exploit_githash = exploit_githash, source_file = service_file)
                        service_exploit = Service_exploit(exploit_sha1 = exploit_sha1, exploit_path = exploit_path, exploit_source = exploit_source, 
                                os_family = os_family, service_name = service_name)
                        service_exploit_list.append(service_exploit)
                        exploit.service_exploit = service_exploit_list
                        session.add(exploit)
                        session.commit()

                    # This is to make the file executable
                    os.chmod(exploit_path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)

def get_remote_ip(service_exploit):
    return service_exploit.ip

def main():
    os_families = ['Windows', 'linux']
    service_names = ['smb', 'ssh', 'smtp', 'ipp', 'postgres', 'mysql', 'ftp', 'pptp', 'oracle', 'snmp', 'pop3', 'rdp', 'postgresql', 'microsoft-ds', 'msrpc', 'netbios-ssn', 'microsoft-rdp'] 
    exploit_types = ['remote', 'local']
    exploit_sources = ['exploitdb']
    for os_family in os_families: 
        for service_name in service_names:
            for exploit_type in exploit_types:
                for exploit_source in exploit_sources:
                    compile_exploits(os_family, service_name, exploit_type, exploit_source)
                    copy_exploits(os_family, service_name, exploit_type, exploit_source)

if __name__ == "__main__":
    if os.getuid() == 0:
        main()
    else:
        print rootMessage
        sys.exit(1)
