"""
exploit.py

Copyright 2006 Andres Riancho

This file is part of w3af, http://w3af.org/ .

w3af 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 version 2 of the License.

w3af 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 w3af; if not, write to the Free Software
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

"""
import shlex

import w3af.core.controllers.output_manager as om
import w3af.core.data.kb.knowledge_base as kb

from w3af.core.controllers.exceptions import (BaseFrameworkException,
                                              NoVulnerabilityFoundException,
                                              ExploitFailedException)
from w3af.core.data.kb.shell import Shell
from w3af.core.ui.console.config import ConfigMenu
from w3af.core.ui.console.callbackMenu import callbackMenu
from w3af.core.ui.console.menu import menu, suggest


class exploit(menu):
    """
    This is the exploit configuration menu for the console.

    :author: Andres Riancho (andres.riancho@gmail.com)
    """
    def __init__(self, name, console, w3af, parent=None):
        menu.__init__(self, name, console, w3af, parent)

        self._configs = {}
        for plugin_name in self._w3af.plugins.get_plugin_list('attack'):
            plugin = self._w3af.plugins.get_plugin_inst('attack', plugin_name)
            config = ConfigMenu(
                plugin_name, self._console, self._w3af, self, plugin)
            self._configs[plugin_name] = config
        self._load_help('exploit')

        # A list with the proxy and shell objects
        self._exploit_results = []

    def _cmd_exploit(self, params, show_list=True):
        """
        Exploit a vuln
        """

        if not len(params):
            self._cmd_help(['exploit'])
            return

        subcmd, params = params[0], params[1:]
        if subcmd == 'config':
            return self._configExploit(params)
        elif subcmd == '*':
            return self._exploit_all(params)
        else:
            return self._exploit(subcmd, params)

    def _para_exploit(self, params, part):
        if len(params) == 0:
            arr = ['*', 'config'] + self._configs.keys()
            return suggest(arr, part)

        if len(params) == 1:
            arr = []
            if params[0] == 'config':
                arr = self._configs.keys()
            if params[0] == '*':
                arr = ['stopOnFirst']

            return suggest(arr, part)

        return []

    def _configExploit(self, params):
        if len(params) == 0:
            raise BaseFrameworkException('Plugin name was expected.')
        if len(params) > 1:
            raise BaseFrameworkException(
                'Unexpected parameters: ' + ','.join(params[1:]))

        plugin_name = params[0]
        if plugin_name not in self._configs:
            raise BaseFrameworkException("Unknown plugin " + plugin_name)

        return self._configs[plugin_name]

    def _exploit_all(self, params):
        lp = len(params)
        stopOnFirst = len(params) > 0 and params[0] == 'stopOnFirst'
        maxLen = int(stopOnFirst)
        if len(params) > maxLen:
            raise BaseFrameworkException('Unexpected parameters: ' +
                                ','.join(params[maxLen:]))

        vuln_list = kb.kb.get_all_vulns()

        if not vuln_list:
            om.out.console('They are no vulnerabilities to exploit.')
        else:
            attackPluginList = self._w3af.plugins.get_plugin_list('attack')
            #Now I create the instances...
            instance_list = []
            for plugin_name in attackPluginList:
                plugin_inst = self._w3af.plugins.get_plugin_inst('attack',
                                                                 plugin_name)
                instance_list.append(plugin_inst)

            # Its time to sort...
            def sortfunc(x, y):
                # reverse ordering...
                return cmp(y.get_root_probability(), x.get_root_probability())
            instance_list.sort(sortfunc)

            # To have a nicer console ;)
            not_run = []
            continue_exploiting = True

            # Exploit !
            for ap in instance_list:

                if not continue_exploiting:
                    break

                if not ap.can_exploit():
                    # save to report later
                    not_run.append(ap.get_name())
                else:
                    # can exploit!
                    msg = 'Applying %s.attack plugin to all vulnerabilities:'
                    om.out.console(msg % ap.get_name())

                    for vuln_obj in vuln_list:
                        continue_exploiting = True

                        msg = '- Exploiting vulnerability with id:' + \
                            str(vuln_obj.get_id())
                        om.out.console(msg)

                        try:
                            self._exploit(ap.get_name(
                            ), vuln_obj.get_id(), show_list=False)
                        except BaseFrameworkException, w:
                            continue_exploiting = True
                            om.out.console(str(w))
                        else:
                            # We get here when the exploit was successful
                            if stopOnFirst:
                                continue_exploiting = False
                                break

                    om.out.console('')

            msg = 'The following plugins weren\'t run because they can\'t'\
                  ' exploit any of the previously discovered vulnerabilities:'\
                  ' %s.'
            om.out.console(msg % ', '.join(not_run))
            om.out.console('')

            if self._exploit_results:
                self._show()
                om.out.console('Please use the "interact" command to use the'
                               ' shell objects.')

    def _exploit(self, plugin_name, params, show_list=True):
        """
        Exploits a vuln. using a single plugin.
        """
        # Did the user indicated what vulnerability to exploit ?
        if len(params) == 1:
            try:
                vuln_to_exploit = params[0]
                if vuln_to_exploit != '*':
                    vuln_to_exploit = int(params[0])
            except:
                raise BaseFrameworkException(
                    'You specified an invalid vulnerability id.')
        else:
            vuln_to_exploit = None

        if plugin_name not in self._configs:
            raise BaseFrameworkException('Unknown plugin. Use the list command'
                                         ' to view available plugins.')

        # Get the exploit plugin
        self._plugin = plugin = self._w3af.plugins.get_plugin_inst('attack',
                                                                   plugin_name)

        try:
            response = plugin.can_exploit(vuln_to_exploit)
        except BaseFrameworkException, e:
            raise e
        else:
            if not response:
                msg = 'No exploitable vulnerabilities found.'
                raise BaseFrameworkException(msg)
            else:
                try:
                    exploit_result = plugin.exploit(vuln_to_exploit)
                except NoVulnerabilityFoundException:
                    raise
                except ExploitFailedException:
                    raise
                else:
                    # everything went ok!
                    if not exploit_result:
                        msg = 'Failed to exploit vulnerability.'
                        raise BaseFrameworkException(msg)
                    else:
                        # Assign a unique identifier to this shell
                        for i in range(len(self._exploit_results), len(exploit_result)):
                            exploit_result[i].set_exploit_result_id(i)

                        self._exploit_results.extend(exploit_result)
                        om.out.console('Vulnerability successfully exploited. ',
                                       new_line=not show_list)
                        if show_list:
                            self._show()
                            om.out.console('Please use the interact command'
                                           ' to interact with the shell'
                                           ' objects.')

    def _cmd_interact(self, parameters):
        """
        Show the available shells and interact with them.
        """
        if len(parameters) not in (0, 1):
            self._cmd_help(['interact'])
            return

        if len(parameters) == 0:
            om.out.console('No shell objects available; please use the exploit'
                           ' plugins to create them.')
            self._show()
            return
        
        try:
            self._selected_shell = int(parameters[0])
        except ValueError:
            self._cmd_help(['interact'])
            return

        # Do we have shells?
        if not len(self._exploit_results):
            msg = 'No shells available.'
            om.out.error(msg)
            return
        
        # Check that the user selected a valid shell
        if self._selected_shell not in range(len(self._exploit_results)):
            msg = 'Please select a shell in range [0-%s].'
            srange = len(self._exploit_results) - 1
            om.out.error(msg % srange)
            return

        # And now check that the "shell" is actually a shell
        # because maybe the user want's to interact with a proxy :S
        # TODO: I should use different variables to store shells and proxies
        if not isinstance(self._exploit_results[self._selected_shell], Shell):
            msg = 'You can only interact with shell objects.'\
                  ' To interact with proxy objects, please use'\
                  ' your browser.'
            om.out.error(msg)
            return
        
        # Everything ok!
        prompt = self._exploit_results[self._selected_shell].get_name()

        msg = 'Execute "exit" to get out of the remote shell.'\
              ' Commands typed in this menu will be run through the %s shell.'
        om.out.console(msg % prompt)
        prompt = prompt + '-' + str(self._selected_shell)
        return callbackMenu(prompt, self._console, self._w3af, self, self._callback)

    def _show(self):
        """
        Show a list of available shells.
        """
        if not self._exploit_results:
            msg = 'No proxy or shell objects have been created.'
            om.out.console(msg)
            return
        
        om.out.console('This is a list of available shells and proxies:')
        for e, n in zip(self._exploit_results, range(len(self._exploit_results))):
            om.out.console('- [' + str(n) + '] ' + str(e))

    def _callback(self, command):
        """
        :param command: The command as typed by the user in the console
        """
        try:
            parsed_command = shlex.split(command)
        except ValueError, ve:
            om.out.console('%s' % ve)
        else:

            shell = self._exploit_results[self._selected_shell]

            command = parsed_command[0]
            params = parsed_command[1:]

            if command == 'exit':
                # We "ask" the shell if we can end it or not
                # after all, the shell has the control right now
                end_it = shell.end_interaction()
                if end_it:
                    return self._console.back
                else:
                    return None
            else:
                try:
                    response = shell.generic_user_input(command, params)
                except BaseFrameworkException:
                    raise
                except Exception, e:
                    msg = 'The "%s" plugin failed to execute the user command,'\
                          ' exception: "%s".'
                    om.out.error(msg % (self._plugin.get_name(), e))
                    return False
                else:
                    # In some cases I just want this callback to print nothing.
                    if response is None:
                        return None
                    else:
                        om.out.console(response)
                return None

    def _cmd_list(self, parameters):
        """
        Lists all available exploit plugins.
        """
        table = [['Plugin', 'Description'], []]
        for plugin in self._configs:
            desc = self._w3af.plugins.get_plugin_inst(
                'attack', plugin).get_desc()
            table.append([plugin, desc])

        self._console.draw_table(table)

    def _back(self, parameters):
        return False
