#/usr/bin/env python
import os
import sys
import magic
import subprocess

# customize these as needed. if a tool already builds a launcher, include that tool's name in the `existing_launchers` list
launcher_path = '/usr/local/bin/'
tool_path = '/opt'
existing_launchers = ['metasploit-framework', 'msfpro', 'sslscan', 'pixiewps', 'reaver', 'hostapd-wpe']

# supported file types
tool_file_types = ['ruby', 'python', 'perl', 'binary']

# launcher contents for each file type. perl launcher is identical to python
python_launcher = """#!/bin/bash

cd %(tool_path)s/%(tool_name)s; python %(tool_file)s "$@"

"""
perl_launcher = python_launcher.replace('python', 'perl')

ruby_launcher = """#!/bin/bash

source /etc/profile.d/rvm.sh

[ -x %(tool_path)s/%(tool_name)s/%(tool_file)s ] || chmod +x %(tool_path)s/%(tool_name)s/%(tool_file)s

cd %(tool_path)s/%(tool_name)s; ruby %(tool_file)s "$@"

"""

def identify_toolfile(files, tool_path, tool, tool_lang=None):
    """ Take in a list of files for a particular tool and semi-intelligently attempt to identify the tool's entry point along with the type of script it is."""
    if not tool_lang:
        # if `tool_lang` is not provided, attempt to identify both the file type and tool entry point
        found_file_type = get_filetype(
            '%s/%s/%s' % (tool_path, tool, tool), mime=True)
        if tool in files and len(filter(None, [file_type in found_file_type for file_type in tool_file_types])):
            # if there is a file in files matching the tool name and its file type matches a known type of file (python, perl, etc.)
            # return that file type and tool name
            file_type = [
                file_type for file_type in tool_file_types if file_type in found_file_type][0]
            return (file_type, tool)

    for file in files:
        file_type = get_filetype('%s/%s/%s' %
                                 (tool_path, tool, file), mime=True)
        if file.lower().startswith(tool.lower()):
            return (tool_lang, file)
        if file_type and tool_lang in file_type and file.startswith(tool.lower()):
            return (tool_lang, file)


def check_existing(tool_name):
    """ Use `which` to see if a tool already exists or alternatively see if a launcher exists in `launcher_path`."""
    if tool_name in existing_launchers:
        return True
    try:
        return True if subprocess.check_output('which %s' % tool_name, shell=True) else False
    except:
        return os.path.isfile('%s%s' % (launcher_path, tool_name))


def get_filetype(path, mime=True):
    """ A wrapper for magic.from_file() to get a file's type"""
    file_type = None
    try:
        file_type = magic.from_file(path, mime=mime)
    except:
        pass
    return file_type


def build(tool_path, tool_name, tool_file, tool_lang):
    """ Take all params and build an executable launcher pointing to a tool's entry point"""
    launcher = ''
    if tool_lang == 'ruby':
        launcher = ruby_launcher % {
            'tool_path': tool_path,
            'tool_name': tool_name,
            'tool_file': tool_file,
        }
    elif tool_lang == 'python':
        launcher = python_launcher % {
            'tool_path': tool_path,
            'tool_name': tool_name,
            'tool_file': tool_file,
        }
    elif tool_lang == 'perl':
        launcher = perl_launcher % {
            'tool_path': tool_path,
            'tool_name': tool_name,
            'tool_file': tool_file,
        }
    elif tool_lang == 'binary':
        # should probably todo this, but maybe check make options for compiled binaries
        # subprocess.check_output('ln -s %s/%s %s%s' % (tool_path, tool_name, launcher_path, tool_name), shell=True)
        pass

    try:
        with open('%s/%s' % (launcher_path, tool_name), 'wb') as fh:
            fh.write(launcher)
        subprocess.check_output('chmod +x %s%s' %
                                (launcher_path, tool_name), shell=True)
    except Exception as e:
        print 'build failed: %s' % e


def parse_tools():
    for tool in os.listdir(tool_path):
        tool_lang = None
        tool_file = None

        print 'checking tool: %s' % tool
        if check_existing(tool):
            print 'launcher for tool: %s exists' % tool
            continue

        files = os.listdir('%s/%s' % (tool_path, tool))

        #  identify file types if possible
        if 'Gemfile' in files or len(filter(None, [file.endswith('rb') for file in files])):
            tool_lang = tool_file_types[0]
        elif len(filter(None, [file.endswith('py') for file in files])):
            tool_lang = tool_file_types[1]
        elif len(filter(None, [file.endswith('pl') for file in files])):
            tool_lang = tool_file_types[2]
        else:
            tool_lang = tool_file_types[3]

        tool_lang, tool_file = identify_toolfile(
            files, tool_path, tool, tool_lang)
        if tool_path and tool and tool_file:
            print 'building launcher for tool: %s' % tool
            build(tool_path, tool, tool_file, tool_lang)


if __name__ == '__main__':
    parse_tools()
    sys.exit(0)
