'''
exploit.py

Copyright 2006 Andres Riancho

This file is part of w3af, w3af.sourceforge.net .

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 w3af
import traceback
import core.controllers.w3afCore
import core.controllers.outputManager as om
from core.controllers.w3afException import *
import core.data.kb.knowledgeBase as kb
from core.data.kb.shell import shell
from core.ui.consoleUi.config import *
from core.ui.consoleUi.util import *
from core.ui.consoleUi.callbackMenu import *
from core.ui.consoleUi.menu import *
from core.controllers.misc.exploitAll import exploitAll

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 pluginName in self._w3af.getPluginList( 'attack' ):
            plugin = self._w3af.getPluginInstance( pluginName, 'attack' )
            config = configMenu ( pluginName, self._console, self._w3af, self, plugin )
            self._configs[pluginName] = config
        self._loadHelp('exploit')

        # A list with the proxy and shell objects
        self._exploitResults = []
        
       
    def _exploitWrapper( self, parameters ):
        '''
        This is a simple wrapper for the _exploit method.
        '''
        try:
            self._exploit( parameters )
        except w3afException, w:
            om.out.console( str(w) )
        except w3afMustStopException:
            om.out.error( str(w) )
        return True
    
    def _cmd_exploit( self, params, showList=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._exploitAll( 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 w3afException( 'Plugin name was expected.')
        if len( params ) > 1:
            raise w3afException( 'Unexpected parameters: ' + ','.join(params[1:]) )

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

        return self._configs[pluginName]

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

        vuln_list = kb.kb.getAllVulns()

        if not vuln_list:
            om.out.console('They are no vulnerabilities to exploit.')
        else:
            attackPluginList = self._w3af.getPluginList( 'attack' )
            #Now I create the instances...
            instanceList = []
            for pluginName in attackPluginList:
                instanceList.append( self._w3af.getPluginInstance( pluginName, 'attack' ) )
            
            # Its time to sort...
            def sortfunc(x,y):
                # reverse ordering...
                return cmp( y.getRootProbability(), x.getRootProbability() )
            instanceList.sort( sortfunc )
            
            # To have a nicer console ;)
            not_run = []
            continue_exploiting = True
            
            # Exploit !
            for ap in instanceList:
                
                if not continue_exploiting:
                    break
                
                if not ap.canExploit():
                    # save to report later
                    not_run.append(ap.getName())
                else:
                    # can exploit!
                    msg = 'Executing '+ ap.getName() +'.attack plugin to all vulnerabilities:'
                    om.out.console( msg )
                    
                    for vuln_obj in vuln_list:
                        continue_exploiting = True
                        
                        msg = '- Exploiting vulnerability with id:' + str(vuln_obj.getId())
                        om.out.console( msg )
                        
                        try:
                            self._exploit( ap.getName() , vuln_obj.getId(), showList=False )
                        except w3afException, 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'
            msg += ' previously discovered vulnerabilities: ' + ', '.join(not_run)
            om.out.console( msg )
            om.out.console('')

            if self._exploitResults:
                self._show()
                om.out.console( 'Please use the "interact" command to use the shell objects.' )
        
    
    def _exploit( self , pluginName, params, showList=True):
        '''
        Exploits a vuln. using a single plugin.
        '''
        
        # Did the user indicated what vulnerability to exploit ?
        if len( params ) == 1:
            try:
                vulnToExploit = params[0]
                if vulnToExploit != '*':
                    vulnToExploit = int(params[0])
            except:
                raise w3afException( 'You specified an invalid vulnerability id.' )
        else:
            vulnToExploit = None
        
        if pluginName not in self._configs:
            raise w3afException( 'Unknown plugin. Use the list command to view available plugins.' )
        else:
            self._plugin = plugin = self._w3af.getPluginInstance( pluginName, 'attack' )

            try:
                response = plugin.canExploit( vulnToExploit )
            except w3afException, e:
                raise e
            else:
                if not response:
                    raise w3afException( 'No exploitable vulnerabilities found.' )
                else:
                    try:
                        exploit_result = plugin.exploit( vulnToExploit )
                    except w3afMustStopException,  w3mse:
                        raise w3afException( str(w3mse) )
                    except w3afException,  w3:
                        raise w3
                    else:
                        # everything went ok!
                        if not exploit_result:
                            raise w3afException( 'Failed to exploit vulnerability.')
                        else:                            
                            # Assign a unique identifier to this shell
                            for i in range(len(self._exploitResults), len(exploit_result) ):
                                exploit_result[i].setExploitResultId( i )
                            
                            self._exploitResults.extend( exploit_result )
                            om.out.console( 'Vulnerability successfully exploited. ' , newLine=not showList )
                            if showList:
                                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'])
        else:
            if len(parameters) == 0:
                if self._exploitResults:
                    self._show()
                else:
                    om.out.console('No proxy or shell objects have been created.')
            else:
                try:
                    self._selectedShell = int( parameters[0] )
                except:
                    self._cmd_help(['interact'])
                else:
                    
                    # Check that the user selected a valid shell
                    if self._selectedShell in range( len( self._exploitResults ) ):
                        
                        # 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._exploitResults[ self._selectedShell ], shell):
                            msg = 'You can only interact with shell objects.'
                            msg += ' To interact with proxy objects, please use your browser.'
                            om.out.error( msg )
                        else:
                            # Everything ok!
                            prompt = self._exploitResults[ self._selectedShell ].getName()
                            
                            msg = 'Execute "exit" to get out of the remote shell.'
                            msg += ' Commands typed in this menu will be runned through the'
                            msg += ' ' + prompt + ' shell'
                            om.out.console( msg )
                            prompt = prompt+'-'+str(self._selectedShell)
                            return callbackMenu(prompt, self._console, self._w3af, self, self._callback)
    #                        console = self._console.fork()
    #                        console.sh( prompt, self._callback )
                    elif len(self._exploitResults) == 0:
                        om.out.console('No shell objects available; please use the exploit plugins to create them.')
                    elif len(self._exploitResults) == 1:
                        om.out.error('You can only interact with shell object with id 0; no other shells available.')
                    else:
                        om.out.error('Please select a shell in range [0-'+str(len(self._exploitResults)-1)+'].')
        
    def _show( self ):
        '''
        Show a list of available shells.
        '''
        om.out.console( 'This is a list of available shells and proxies:' )
        for e, n in zip(self._exploitResults, range(len(self._exploitResults))):
            om.out.console('- [' + str(n) + '] ' + str(e) )
        
    def _callback( self, command ):
        shell = self._exploitResults[ self._selectedShell ]
        
        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 )
            except w3afException, w3:
                raise
            except Exception, e:
                om.out.error('The '+ self._plugin.getName() +' plugin failed to execute the user command, exception: ' + str(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_fastexploit( self , parameters, showList=True):
        '''
        Performs fast exploiting based on the parameters provided by the user, and
        the previous plugin configuration.
        '''
        # I need this to have logging!
        self._w3af.initPlugins()
        
        if not len( parameters ):
            om.out.console( 'Incorrect call to fastexploit, please see the help:' )
            self._cmd_help( ['fastexploit'] )
        else:
            pluginName = parameters[0]
            if pluginName not in self._w3af.getPluginList('attack'):
                om.out.console( 'Unknown plugin. Use the list command to view available plugins.' )
            else:
                self._plugin = plugin = self._w3af.getPluginInstance( pluginName, 'attack' )

                try:
                    exploit_result = plugin.fastExploit()
                except Exception, e:
#                    traceback.print_exc() # TODO
                    raise e

                # Assign a unique identifier to this shell
                for i in range(len(self._exploitResults), len(exploit_result) ):
                    exploit_result[i].setExploitResultId( i )

                if not exploit_result:
                    raise w3afException( 'Failed to exploit vulnerability.')
                else:
                    self._exploitResults.extend( exploit_result )                    
                    om.out.console( 'Vulnerability successfully exploited. ' , newLine=not showList )
                    if showList:
                        self._show()
                        om.out.console( 'Please use the interact command to interact with the shell objects.' )
        
#        return True
        
    def _cmd_list(self , parameters):
        '''
        Lists all available exploit plugins.
        '''
        table = [['Plugin', 'Description'], []]
        for plugin in self._configs:
            desc = self._w3af.getPluginInstance( plugin, 'attack' ).getDesc()
            table.append([plugin, desc])

        self._console.drawTable(table)
        
    def _back( self, parameters ):
        return False
    
