'''
w3afCore.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 core.controllers.outputManager as om

# Before doing anything, check if I have all needed dependencies
from core.controllers.misc.dependencyCheck import dependencyCheck
dependencyCheck()

# Called here to init some variables in the config ( cf.cf.save() )
import core.controllers.miscSettings as miscSettings

import os,sys
from core.controllers.misc.factory import factory
from core.controllers.misc.parseOptions import parseOptions
from core.data.url.xUrllib import xUrllib
import core.data.parsers.urlParser as urlParser
from core.controllers.w3afException import w3afException
from core.controllers.w3afException import w3afRunOnce
from core.controllers.sessionManager import sessionManager
from core.controllers.targetSettings import targetSettings as targetSettings

import traceback
import copy

import core.data.kb.knowledgeBase as kb
import core.data.kb.config as cf
from core.data.request.frFactory import createFuzzableRequests

class w3afCore:
	'''
	This is the core of the framework, it calls all plugins, handles exceptions,
	coordinates all the work, creates threads, etc.
	 
	@author: Andres Riancho ( andres.riancho@gmail.com )
	'''

	def __init__(self ):
		self.uriOpener = xUrllib()
		self._session = sessionManager()
		self._session.saveSession()
		
		# A map with plugin types as keys and a list of plugin names as values
		self._strPlugins = {'audit':[],'grep':[],'bruteforce':[],'discovery':[],\
		'evasion':[], 'mangle':[], 'output':['console']}
		
		# A map with plugin types as keys and a list of plugin instances as values
		self._plugins = {'audit':[],'grep':[],'bruteforce':[],'discovery':[],\
		'evasion':[], 'mangle':[], 'output':[]}

		self._pluginsOptions = {}
		
		self._fuzzableRequestList  = []
		
		self._initialized = False
		self.target = targetSettings()
	
	def resumeSession( self, sessionName ):
		'''
		Resumes a session object.
		
		@parameter sessionName: The name of the session
		@return: None
		'''
		self._session = sessionManager()
		try:
			self._session.loadSession( sessionName )
		except w3afException, w3:
			om.out.error( str(w3) )
		else:
			self._pluginsOptions = self._session.getData('pluginsOptions')
			self._strPlugins = self._session.getData('strPlugins')
			om.out.console('Restoring saved session: ' + sessionName )
			om.out.console('w3af>> start')
			self.start()
		
	def saveSession( self, sessionName ):
		'''
		Creates a session object, to make it ready to write data to it.
		
		@parameter sessionName: The name of the session
		@return: None
		'''
		self._session = sessionManager()
		self._session.saveSession( sessionName )
	
	def _rPlugFactory( self, strReqPlugins, PluginType ):
		'''
		This method creates the requested modules list.
		
		@parameter strReqPlugins: A string list with the requested plugins to be executed.
		@parameter PluginType: [audit|discovery|grep]
		@return: A list with plugins to be executed, this list is ordered using the exec priority.
		'''		
		requestedPluginsList = []
		
		if 'all' in strReqPlugins:
			fileList = [ f for f in os.listdir('plugins' + os.path.sep+ PluginType + os.path.sep ) ]    
			allPlugins = [ os.path.splitext(f)[0] for f in fileList if os.path.splitext(f)[1] == '.py' ]
			allPlugins.remove ( '__init__' )
			
			if len ( strReqPlugins ) != 1:
				# [ 'all', '!sqli' ]
				# I want to run all plugins except sqli
				unwantedPlugins = [ x[1:] for x in strReqPlugins if x[0] =='!' ]
				strReqPlugins = list( set(allPlugins) - set(unwantedPlugins) ) #bleh! v2
			else:
				strReqPlugins = allPlugins
			
			# Update the plugin list
			# This update is usefull for cases where the user selected "all" plugins,
			# the self._strPlugins[PluginType] is useless if it says 'all'.
			self._strPlugins[PluginType] = strReqPlugins
				
		for pluginName in strReqPlugins:
			plugin = factory( 'plugins.' + PluginType + '.' + pluginName )
			
			# Now we are going to check if the plugin dependencies are met
			for dep in plugin.getPluginDeps():
				try:
					depType, depPlugin = dep.split('.')
				except:
					raise w3afException('Plugin dependencies must be indicated using pluginType.pluginName notation.\
					This is an error in ' + pluginName +'.getPluginDeps() .')
				if depType == PluginType:
					if depPlugin not in strReqPlugins:
						if cf.cf.getData('autoDependencies'):
							strReqPlugins.append( depPlugin )
							om.out.information('Auto-enabling plugin: ' + PluginType + '.' + depPlugin)
							# nice recursive call, this solves the "dependency of dependency" problem =)
							return self._rPlugFactory( strReqPlugins, depType )
						else:
							raise w3afException('Plugin '+ pluginName +' depends on plugin ' + dep + ' and ' + dep + ' is not enabled. ')
				else:
					if depPlugin not in self._strPlugins[depType]:
						if cf.cf.getData('autoDependencies'):
							dependObj = factory( 'plugins.' + depType + '.' + depPlugin )
							dependObj.setUrlOpener( self.uriOpener )
							if dependObj not in self._plugins[depType]:
								self._plugins[depType].insert( 0, dependObj )
								self._strPlugins[depType].append( depPlugin )
							om.out.information('Auto-enabling plugin: ' + depType + '.' + depPlugin)
						else:
							raise w3afException('Plugin '+ pluginName +' depends on plugin ' + dep + ' and ' + dep + ' is not enabled. ')
					else:
						# if someone in another planet depends on me... run first
						self._strPlugins[depType].remove( depPlugin )
						self._strPlugins[depType].insert( 0, depPlugin )
			
			# Now we set the plugin options
			if pluginName in self._pluginsOptions.keys():
				plugin.setOptions( self._pluginsOptions[pluginName] )
				
			# This sets the url opener for each module that is called inside the for loop
			plugin.setUrlOpener( self.uriOpener )
			# Append the plugin to the list
			requestedPluginsList.append ( plugin )
			
		# The plugins are all on the requestedPluginsList, now I need to order them
		# based on the module dependencies. For example, if A depends on B , then
		# B must be runned first.
		
		orderedPluginList = []
		for plugin in requestedPluginsList:
			deps = plugin.getPluginDeps()
			if len( deps ) != 0:
				# This plugin has dependencies I should add the plugins in order
				for plugin2 in requestedPluginsList:
					if PluginType+'.'+plugin2.__class__.__name__ in deps:
						orderedPluginList.insert( 1, plugin2)

			# Check if I was added because of a dep
			if plugin not in orderedPluginList:
				orderedPluginList.insert( 100, plugin )
		
		return orderedPluginList
	
	def _initPlugins( self ):
		self._initialized = True
		
		# This is inited before all, to have a full logging facility.
		om.out.setOutputPlugins( self._strPlugins['output'] )
		
		# First, create an instance of each requested plugin and add it to the plugin list
		# Plugins are added taking care of plugin dependencies
		self._plugins['audit'] = self._rPlugFactory( self._strPlugins['audit'] , 'audit')
		
		self._plugins['bruteforce'] = self._rPlugFactory( self._strPlugins['bruteforce'] , 'bruteforce')		
		
		# First, create an instance of each requested module and add it to the module list
		self._plugins['discovery'] = self._rPlugFactory( self._strPlugins['discovery'] , 'discovery')
		
		self._plugins['grep'] = self._rPlugFactory( self._strPlugins['grep'] , 'grep')
		self.uriOpener.setGrepPlugins( self._plugins['grep'] )
		
		self._plugins['mangle'] = self._rPlugFactory( self._strPlugins['mangle'] , 'mangle')
		self.uriOpener.settings.setManglePlugins( self._plugins['mangle'] )
	
	def _createURLList( self ):
		'''
		Creates an URL list in the kb
		'''
		urlList = [ fr.getURL() for fr in self._fuzzableRequestList]
		urlList = list( set( urlList ) )
		kb.kb.save( 'urls', 'urlList' ,  urlList )

		uriList = [ fr.getURI() for fr in self._fuzzableRequestList]
		uriList = list( set( uriList ) )
		kb.kb.save( 'urls', 'uriList' ,  uriList )
	
	def _discoverAndBF( self ):
		'''
		Discovery and bruteforce phases are related, so I have joined them
		here in this method.
		'''
		go = True
		tmpList = copy.deepcopy( self._fuzzableRequestList )
		res = []
		discoveredFrList = []
		
		# this is an identifier to know what call number of _discoverWorker we are working on
		self._count = 0
		
		while go:
			discoveredFrList = self._discover( tmpList )
			successfullyBruteforced = self._bruteforce( discoveredFrList )
			if not successfullyBruteforced:
				# Havent found new credentials
				go = False
				for fr in discoveredFrList:
					if fr not in res:
						res.append( fr )
			else:
				tmp = []
				tmp.extend( discoveredFrList )
				tmp.extend( successfullyBruteforced )
				for fr in tmp:
					if fr not in res:
						res.append( fr )
				
				# So in the next "while go:" loop I can do a discovery
				# using the new credentials I found
				tmpList = successfullyBruteforced
				
				# Now I reconfigure the urllib to use the newly found credentials
				self._reconfigureUrllib()
				
		return res
	
	def _reconfigureUrllib( self ):
		'''
		Configure the main urllib with the newly found credentials.
		'''
		for v in kb.kb.getData( 'basicAuthBrute' , 'auth' ):
			self.uriOpener.settings.setBasicAuth( v.getURL(), v['user'], v['pass'] )
		
		# TODO: Add something here for form auth brute.
		# I should add a header with a cookie or something...
		'''
		for v in kb.kb.getData( 'formAuthBrute' , 'auth' ):
			self.uriOpener.settings.setBasicAuth( v.getURL(), v['user'], v['pass'] )
		'''
	
	def start(self):
		'''
		Starts the work.
		
		@return: No value is returned.
		''' 
		try:
			self._checkParameters()
		except Exception,e:
			om.out.debug('_checkParameters() raised an exception' )
			om.out.error( str(e) )
		else:
			try:
				
				###### This is the main section ######
				
				# Init mangling, output and audit plugins
				self._initPlugins()
				
				# Create the first fuzzableRequestList
				for url in cf.cf.getData('targets'):
					try:
						response = self.uriOpener.GET( url )
						self._fuzzableRequestList.extend( createFuzzableRequests( response ) )
					except KeyboardInterrupt:
						raise
					except:
						om.out.information( 'The target URL: ' + url + ' is unreachable.' )
				
				# Create the initial urlList
				self._createURLList()
			
				self._fuzzableRequestList = self._discoverAndBF()
				
				# Update the list
				self._createURLList()
				
				if len( self._fuzzableRequestList ) == 0:
					om.out.information('No URLs found by discovery.')
				else:
					om.out.information('The list of found URLs is:')
					for url in kb.kb.getData( 'urls', 'urlList'):
						om.out.information('- ' + url )
					msg = 'Found ' + str(len( kb.kb.getData( 'urls', 'urlList') )) + ' URLs and ' + str(len( self._fuzzableRequestList)) + ' different points of injection.'
					om.out.information( msg )
					
					om.out.information('The list of Fuzzable requests is:')
					for fuzzRequest in self._fuzzableRequestList:
						om.out.information( '- ' + str( fuzzRequest) )
				
					self._audit()
					
				self.end()
				###########################
				
			except w3afException, e:
				self.end()
				om.out.error( str(e) )
				raise e
			except KeyboardInterrupt, e:
				# I wont handle this. 
				# The user interface must know what to do with it
				self.end()
				raise e
	
	def end( self ):
		'''
		This method is called when the process ends.
		'''
		for plugin in self._plugins['grep']:
			plugin.end()
		
		om.out.endOutputPlugins()
		
	def _discover( self, toWalk ):
		# The active session doesn't have a discovery result saved to it
		# Init some internal variables
		self._alreadyWalked = toWalk
		self._urls = []
		self._firstDiscovery = True
		
		for fr in toWalk:
			fr.iterationNumber = 0
		
		result = []
		try:
			result = self._discoverWorker( toWalk )
		except KeyboardInterrupt, e:
			om.out.information('The user interrupted the discovery phase, continuing with audit.')
			result = self._alreadyWalked
		
		return result
		
	def _discoverWorker(self, toWalk ):
		om.out.debug('Called _discoverWorker()' )
		
		while len( toWalk ) and self._count < cf.cf.getData('maxDiscoveryLoops'):
		
			# This variable is for session saving and LOOP evasion
			self._count += 1
			
			pluginsToRemoveList = []
			fuzzableRequestList = []
			
			for plugin in self._plugins['discovery']:
				for fr in toWalk:
					
					if fr.iterationNumber > cf.cf.getData('maxDepth'):
						om.out.debug('Avoiding discovery loop in fuzzableRequest: ' + str(fr) )
					else:
						if self._session.getData( (plugin.getName(), fr.dump(), self._count) ) == None:
							# Nothing saved in the session about this tuple
							om.out.debug('Running plugin: ' + plugin.getName() )
							try:
								fuzzableRequestList.extend( plugin.discover( fr ) )
							except w3afException,e:
								om.out.error( str(e) )
							except w3afRunOnce, rO:
								# Some plugins are ment to be run only once
								# that is implemented by raising a w3afRunOnce exception
								pluginsToRemoveList.append( plugin )
							else:
								self._session.save( (plugin.getName(), fr.dump(), self._count), fuzzableRequestList )
							om.out.debug('Ending plugin: ' + plugin.getName() )
						else:
							# The data saved in the session is usefull !
							om.out.debug('Restoring results from session for plugin: ' + plugin.getName() + '; dump:' + fr.dump().replace('\n', '') +' ; count: ' + str(self._count)  )
							fuzzableRequestList.extend( self._session.getData( (plugin.getName(), fr.dump(), self._count) ) )
					#end-if
				#end-for
			#end-for
			
			##
			##	The search has finished, now i'll some mangling with the requests
			##
			newFR = []
			for iFr in fuzzableRequestList:
				# I dont care about fragments ( http://a.com/foo.php#frag ) and I dont really trust plugins
				# so i'll remove fragments here
				iFr.setURL( urlParser.removeFragment( iFr.getURL() ) )
				
				# Increment the iterationNumber !
				iFr.iterationNumber = fr.iterationNumber + 1
				
				if iFr not in self._alreadyWalked and urlParser.baseUrl( iFr.getURL() ) in cf.cf.getData('baseURLs'):
					# Found a new fuzzable request
					newFR.append( iFr )
					self._alreadyWalked.append( iFr )
					if iFr.getURL() not in self._urls:
						om.out.information('New URL found by discovery: ' +  iFr.getURL() )
						self._urls.append( iFr.getURL() )
			
			##
			##	Cleanup!
			##
			
			# This wont be used anymore, here i'm duplicating objects that are already saved
			# in the self._alreadyWalked list.
			del fuzzableRequestList
			try: del iFr
			except: pass
			
			# Get ready for next while loop
			toWalk = newFR
			
			# Remove plugins that dont want to be runned anymore
			for pluginToRemove in pluginsToRemoveList:
				if pluginToRemove in self._plugins['discovery']:
					self._plugins['discovery'].remove( pluginToRemove )
					om.out.debug('The discovery plugin: ' + pluginToRemove.getName() + ' wont be runned anymore.')		
				
		return self._alreadyWalked
	
	def _audit(self):
		om.out.debug('Called _audit()' )
		
		# This two for loops do all the audit magic [KISS]
		for plugin in self._plugins['audit']:
			if not self._session.getData( plugin.getName() ):
				om.out.information('Starting ' + plugin.getName() + ' plugin execution.')
				for fr in self._fuzzableRequestList:
					# Sends each fuzzable request to the plugin
					try:
						plugin.audit( fr )
					except w3afException, e:
						om.out.error( str(e) )
				# Let the plugin know that we are not going to use it anymore
				try:
					plugin.end()
				except w3afException, e:
					om.out.error( str(e) )
				# Save to the session a flag saying : "the plugin xyz was runned"
				self._session.save( plugin.getName(), True )
			else:
				om.out.information('Not running plugin: ' + plugin.getName() + ' , the plugin results are being read from the session information.')
				
	def _bruteforce(self, fuzzableRequestList):
		'''
		@parameter fuzzableRequestList: A list of fr's to be analyzed by the bruteforce plugins
		@return: A list of the URL's that have been successfully bruteforced
		'''
		res = []
		
		om.out.debug('Called _bruteforce()' )
		
		for plugin in self._plugins['bruteforce']:
			if not self._session.getData( plugin.getName() ):
				om.out.information('Starting ' + plugin.getName() + ' plugin execution.')
				for fr in fuzzableRequestList:
					# Sends each url to the plugin
					try:
						frList = plugin.bruteforce( fr )
					except w3afException, e:
						om.out.error( str(e) )
						
					try:
						plugin.end()
					except w3afException, e:
						om.out.error( str(e) )
						
					res.extend( frList )
				
				# Save to the session a flag saying : "the plugin xyz was runned"
				self._session.save( plugin.getName(), True )
		
		return res

	def setPluginOptions(self, pluginName, PluginsOptions ):
		'''
		@parameter PluginsOptions: A tuple with a string and a dictionary with the options for a plugin. For example:\
		{ googlespider_plugin:{'LICENSE_KEY':'AAAA'} }
			
		@return: No value is returned.
		'''
		pluginName, PluginsOptions = parseOptions( pluginName, PluginsOptions ) 		
		self._pluginsOptions[pluginName] = PluginsOptions
		self._session.save('pluginsOptions', self._pluginsOptions )
	
	def getPlugins( self, pluginType ):
		return self._strPlugins[ pluginType ]
	
	def setPlugins( self, pluginNames, pluginType ):
		'''
		This method sets the plugins that w3afCore is going to use. Before this plugin
		existed w3afCore used setDiscoveryPlugins() / setAuditPlugins() / etc , this wasnt
		really extensible and was replaced with a combination of setPlugins and getPluginTypes.
		This way the user interface isnt bound to changes in the plugin types that are added or
		removed.
		
		@parameter pluginNames: A list with the names of the Plugins that will be runned.
		@parameter pluginType: The type of the plugin.
		@return: None
		'''
		setMap = {'discovery':self._setDiscoveryPlugins, 'audit':self._setAuditPlugins,\
		'grep':self._setGrepPlugins, 'evasion':self._setEvasionPlugins, 'output':self._setOutputPlugins\
		, 'mangle': self._setManglePlugins, 'bruteforce': self._setBruteforcePlugins}
		
		func = setMap[ pluginType ]
		func( pluginNames )
		self._session.save( 'strPlugins', self._strPlugins )
		
	def getPluginTypes( self ):
		'''
		@return: A list with all plugin types.
		'''
		pluginTypes = [ f for f in os.listdir('plugins/') if f.count('.py') == 0 ]
		pluginTypes.remove( 'attack' )
		if '.svn' in pluginTypes:
			pluginTypes.remove('.svn')
		return pluginTypes
	
	def _setBruteforcePlugins( self, bruteforcePlugins ):
		'''
		@parameter manglePlugins: A list with the names of output Plugins that will be runned.
		@return: No value is returned.
		'''
		self._strPlugins['bruteforce'] = bruteforcePlugins
	
	def _setManglePlugins( self, manglePlugins ):
		'''
		@parameter manglePlugins: A list with the names of output Plugins that will be runned.
		@return: No value is returned.
		'''
		self._strPlugins['mangle'] = manglePlugins
	
	def _setOutputPlugins( self, outputPlugins ):
		'''
		@parameter outputPlugins: A list with the names of output Plugins that will be runned.
		@return: No value is returned.
		'''
		self._strPlugins['output'] = outputPlugins
		
	def _setDiscoveryPlugins( self, discoveryPlugins ):
		'''
		@parameter discoveryPlugins: A list with the names of Discovery Plugins that will be runned.
		@return: No value is returned.
		''' 		
		self._strPlugins['discovery'] = discoveryPlugins
	
	def _setAuditPlugins( self, AuditPlugins ):
		'''
		@parameter AuditPlugins: A list with the names of Audit Plugins that will be runned.
		@return: No value is returned.
		''' 		
		self._strPlugins['audit'] = AuditPlugins
		
	def _setGrepPlugins( self, GrepPlugins):
		'''
		@parameter GrepPlugins: A list with the names of Grep Plugins that will be used.
		@return: No value is returned.
		''' 	
		self._strPlugins['grep'] = GrepPlugins
		
	def _setEvasionPlugins( self, EvasionPlugins ):
		'''
		@parameter EvasionPlugins: A list with the names of Evasion Plugins that will be used.
		@return: No value is returned.
		''' 	
		self._plugins['evasion'] = self._rPlugFactory( EvasionPlugins , 'evasion')
		self.uriOpener.setEvasionPlugins( self._plugins['evasion'] )

	def _checkParameters(self):
		'''
		Checks if all parameters where configured correctly by the above layer (w3af.py)
		''' 	
		try:
			assert cf.cf.getData('targets')  != None, 'No target URI configured.'
		except AssertionError, ae:
			raise w3afException( str(ae) )
			
		try:
			cry = True
			if len(self._strPlugins['audit']) == 0 and len(self._strPlugins['discovery']) == 0 \
			and len(self._strPlugins['grep']) == 0:
				cry = False
			assert cry , 'No audit, grep or discovery plugins configured to run.'
		except AssertionError, ae:
			raise w3afException( str(ae) )
			
	def getPluginList( self, PluginType ):
		'''
		@return: A string list of the names of all available plugins by type.
		'''
		fileList = [ f for f in os.listdir('plugins/' + PluginType +'/') ]    
		strReqPlugins = [ os.path.splitext(f)[0] for f in fileList if os.path.splitext(f)[1] == '.py' ]
		strReqPlugins.remove ( '__init__' )
		return strReqPlugins
		
	def getPluginInstance( self, pluginName ):
		'''
		@return: An instance of a plugin.
		'''
		pTypes = self.getPluginTypes()
		pTypes.append('attack')
		
		# First we need to find the plugin.
		for pluginType in pTypes:
			fileList = [ f for f in os.listdir('plugins/' + pluginType +'/') ]    
			fileList = [ os.path.splitext(f)[0] for f in fileList if os.path.splitext(f)[1] == '.py' ]
			fileList.remove ( '__init__' )
			if pluginName in fileList:
				ModuleName = 'plugins.'+pluginType+ '.'+ pluginName
				__import__(ModuleName)
				aModule = sys.modules[ModuleName]
				className = ModuleName.split('.')[len(ModuleName.split('.'))-1]
				aClass = getattr( aModule , className )
				plugin = apply(aClass, ())
				# This sets the url opener for each module that is called inside the for loop
				plugin.setUrlOpener( self.uriOpener )
				if pluginName in self._pluginsOptions.keys():
					plugin.setOptions( self._pluginsOptions[pluginName] )
					
				# This will init some plugins like mangle and output
				if pluginType == 'attack' and not self._initialized:
					self._initPlugins()
				return plugin
				
		raise w3afException('Plugin not found')
	
# """"Singleton""""
wCore = w3afCore()

