require 'msf/core'

module Msf

###
#
# This module exposes an interface that is used when wanting to receive
# notifications about events pertaining to exploitation.
#
###
module ExploitEvent

	#
	# This method is called when an exploit succeeds.
	#
	def on_exploit_success(exploit, session)
	end

end

###
#
# The exploit class acts as the base class for all exploit modules.  It
# provides a common interface for interacting with exploits at the most basic
# level.
#
###
class Exploit < Msf::Module

	##
	#
	# Default compatibility settings for exploit modules.
	#
	##
	module CompatDefaults
		#
		# Default compatibility specifications for payloads
		#
		Payload = 
			{ 
				# Support reverse, bind, find, and noconn connection types
				# for all exploits unless expressly disabled.
				'ConnectionType' => 'reverse bind find noconn none tunnel',
			}
	end

	##
	#
	# The various check codes that can be returned from the ``check'' routine.
	#
	##
	module CheckCode

		#
		# The target is safe and is therefore not exploitable.
		#
		Safe        = [ 0, "The target is not exploitable." ]

		#
		# The target is running the service in requestion but may not be
		# exploitable.
		#
		Detected    = [ 1, "The target service is running, but could not be validated." ]

		#
		# The target appears to be vulnerable.
		#
		Appears     = [ 2, "The target appears to be vulnerable." ]

		#
		# The target is vulnerable.
		#
		Vulnerable  = [ 3, "The target is vulnerable." ]

		#
		# The exploit does not support the check method.
		#
		Unsupported = [ 4, "This exploit does not support check." ]
	end

	#
	# The various basic types of exploits
	#
	module Type

		#
		# Indicates that the exploit is a remote exploit.
		#
		Remote  = "remote"

		#
		# Indicates that the exploit is a local exploit.
		#
		Local   = "local"

		#
		# Indicates that the exploit can work anywhere it damn pleases.
		#
		Omni    = "omnipresent"
	end

	#
	# The types of stances an exploit can take, such as passive or aggressive.
	# Stances indicate whether or not the exploit triggers the exploit without
	# waiting for one or more conditions to be met (aggressive) or whether it
	# must wait for certain conditions to be satisfied before the exploit can
	# be initiated (passive)
	#
	module Stance

		#
		# Used to indicate that an exploit takes an aggressive stance.  This
		# means that the exploit proactively triggers a vulnerability.
		#
		Aggressive = "aggresive"

		#
		# Used to indicate that an exploit takes a passive stance.  This means
		# that the exploit waits for interaction from a client or other entity
		# before being able to trigger the vulnerability.
		#
		Passive    = "passive"
	end

	###
	#
	# The local exploit class is a specialization of the exploit module class that
	# is geared toward exploits that are performed locally.  Locally, in this
	# case, is defined as an exploit that is realized by means other than network
	# communication.
	#
	###
	class Local < Exploit

		#
		# Returns the fact that this exploit is a local exploit.
		#
		def exploit_type
			Exploit::Type::Local
		end
	end

	###
	#
	# The remote exploit class is a specialization of the exploit module class
	# that is geared toward exploits that are performed against targets other than
	# the local machine.  This typically implies exploiting other machines via a
	# network connection, though it is not limited to this scope.
	#
	###
	class Remote < Exploit

		#
		# Initializes the socket array.
		#
		def initialize(info)
			super

			self.sockets = Array.new
		end

		#
		# Returns the fact that this exploit is a remote exploit.
		#
		def exploit_type
			Exploit::Type::Remote
		end

		#
		# Adds a socket to the list of sockets opened by this exploit.
		#
		def add_socket(sock)
			self.sockets << sock
		end

		#
		# Removes a socket from the list of sockets.
		#
		def remove_socket(sock)
			self.sockets.delete(sock)
		end

		#
		# This method is called once a new session has been created on behalf of
		# this exploit instance and all socket connections created by this
		# exploit should be closed.
		#
		def abort_sockets
			sockets.delete_if { |sock|
				sock.abortive_close = true

				begin
					disconnect(sock)
				rescue
				end

				true
			}
		end

	protected

		#
		# The list of sockets established by this exploit.
		#
		attr_accessor :sockets

	end

	#
	# All exploit mixins should be added to the list below
	#
	
	# Behavior
	require 'msf/core/exploit/brute'
	require 'msf/core/exploit/brutetargets'	
	# Payload
	require 'msf/core/exploit/egghunter'
	require 'msf/core/exploit/seh'
	
	# Protocol
	require 'msf/core/exploit/tcp'
	require 'msf/core/exploit/udp'
	require 'msf/core/exploit/smb'
	require 'msf/core/exploit/ftp'
	require 'msf/core/exploit/http'
	require 'msf/core/exploit/dcerpc'
	require 'msf/core/exploit/mssql'
	require 'msf/core/exploit/arkeia'
	require 'msf/core/exploit/ndmp'
	require 'msf/core/exploit/imap'
						
	#
	# Creates an instance of the exploit module.  Mad skillz.
	#
	def initialize(info = {})

		# Ghetto compat mirroring for payload compatibilities.  This mirrors
		#
		# Payload => Compat => xyz
		#
		# to
		#
		# Compat => Payload => xyz
		if (info['Payload'] and info['Payload']['Compat'])
			info['Compat'] = Hash.new if (info['Compat'] == nil)
			info['Compat']['Payload'] = Hash.new if (info['Compat']['Payload'] == nil)
			info['Compat']['Payload'].update(info['Payload']['Compat'])
		end

		# Call the parent constructor after making any necessary modifications
		# to the information hash.
		super(info)

		self.targets = Rex::Transformer.transform(info['Targets'], Array,
				[ Target ], 'Targets')
		self.default_target = info['DefaultTarget']
		self.payload_info = info['Payload'] || {}
		self.session_count = 0
	end

	##
	#
	# Core exploit interface
	#
	# These are the methods that exploits will override to perform various
	# tasks, such as checking a target to see if it's vulnerable, automatically
	# selecting a target, or performing an exploit.
	#
	##

	#
	# Checks to see if the target is vulnerable, returning unsupported if it's
	# not supported.
	#
	# This method is designed to be overriden by exploit modules.
	#
	def check
		CheckCode::Unsupported
	end

	#
	# Kicks off the actual exploit.  Prior to this call, the framework will
	# have validated the data store using the options associated with this
	# exploit module.  It will also pre-generate the desired payload, though
	# exploits can re-generate the payload if necessary. 
	#
	# This method is designed to be overriden by exploit modules.
	#
	def exploit
	end

	#
	# Prepares the module for exploitation, initializes any state, and starts
	# the payload handler.
	#
	def setup
		# Reset the session counts to zero.
		reset_session_counts

		if (payload_instance)
			# Set up the payload handlers
			payload_instance.setup_handler
	
			# Start the payload handler
			payload_instance.start_handler
		end
	end

	#
	# Performs any cleanup that may be necessary, such as disconnecting
	# connections and any other such fun things.  If a payload is active then
	# its handler cleanup routines are called as well.
	#
	def cleanup
		if (payload_instance)
			payload_instance.cleanup_handler
		end
	end

	#
	# Generates the encoded version of the supplied payload using the payload
	# requirements specific to this exploit.  The encoded instance is returned
	# to the caller.  This method is exposed in the manner that it is such
	# that passive exploits and re-generate an encoded payload on the fly
	# rather than having to use the pre-generated one.
	#
	# The return value is an EncodedPayload instance.
	#
	def generate_payload(pinst = nil)
		# Set the encoded payload to the result of the encoding process
		self.payload = generate_single_payload(pinst)

		# Save the payload instance
		self.payload_instance = (pinst) ? pinst : self.payload_instance

		return self.payload
	end

	#
	# This method generates a non-cached payload which is typically useful for
	# passive exploits that will have more than one client.
	#
	def generate_single_payload(pinst = nil)
		if (target == nil)
			raise MissingTargetError, "No target has been specified.",
				caller
		end

		# If a payload instance was supplied, use it, otherwise
		# use the active payload instance
		real_payload = (pinst) ? pinst : self.payload_instance

		if (real_payload == nil)
			raise MissingPayloadError, "No payload has been selected.",
				caller
		end

		# Duplicate the exploit payload requirements
		reqs = self.payload_info.dup

		# Pass save register requirements to the NOP generator
		reqs['SaveRegisters']  = nop_save_registers
		reqs['Prepend']        = payload_prepend
		reqs['PrependEncoder'] = payload_prepend_encoder
		reqs['BadChars']       = payload_badchars
		reqs['Append']         = payload_append
		reqs['MaxNops']        = payload_max_nops
		reqs['MinNops']        = payload_min_nops
		reqs['EncoderType']    = payload_encoder_type
		reqs['EncoderOptions'] = payload_encoder_options

		return EncodedPayload.create(real_payload, reqs)
	end

	#
	# Re-generates an encoded payload, typically called after something in the
	# datastore has changed.
	#
	def regenerate_payload
		generate_single_payload
	end

	##
	#
	# Feature detection
	#
	# These methods check to see if there is a derived implementation of
	# various methods as a way of inferring whether or not a given exploit
	# supports the feature.
	#
	##

	#
	# Returns true if the exploit module supports the check method.
	#
	def supports_check?
		derived_implementor?(Msf::Exploit, 'check')
	end

	#
	# Returns true if the exploit module supports the exploit method.
	#
	def supports_exploit?
		derived_implementor?(Msf::Exploit, 'exploit')
	end

	#
	# Returns a hash of the capabilities this exploit module has support for,
	# such as whether or not it supports check and exploit.
	#
	def capabilities
		{
			'check'       => supports_check?,
			'exploit'     => supports_exploit?
		}
	end

	##
	#
	# Getters/Setters
	#
	# Querying and set interfaces for some of the exploit's attributes.
	#
	##

	#
	# Returns MODULE_EXPLOIT to indicate that this is an exploit module.
	#
	def self.type
		MODULE_EXPLOIT
	end

	#
	# Returns MODULE_EXPLOIT to indicate that this is an exploit module.
	#
	def type
		MODULE_EXPLOIT
	end

	#
	# If we don't know the exploit type, then I guess it's omnipresent!
	#
	def exploit_type
		Type::Omni
	end

	#
	# Generally, all exploits take an aggressive stance.
	#
	def stance
		module_info['Stance'] || Stance::Aggressive
	end

	#
	# Returns if the exploit has a passive stance.
	#
	def passive?
		(stance == Stance::Passive)
	end

	#
	# Returns the active target for this exploit.  If not target has been
	# defined, nil is returned.  If no target was defined but there is a
	# default target, that one will be automatically used.
	#
	def target
		target_idx = datastore['TARGET']

		# Use the default target if one was not supplied.
		if (target_idx == nil and default_target and default_target >= 0)
			target_idx = default_target
		end

		return (target_idx) ? targets[target_idx.to_i] : nil
	end

	#
	# Returns a list of compatible payloads based on platform, architecture,
	# and size requirements.
	#
	def compatible_payloads
		payloads = []
		
		
		c_platform = (target and target.platform) ? target.platform : platform
		c_arch     = (target and target.arch)     ? target.arch     : (arch == []) ? nil : arch

		framework.payloads.each_module(
			'Platform' => c_platform,
			'Arch'     => c_arch ) { |name, mod|

			# Skip over payloads that are too big
			if ((payload_space) and 
			    (framework.payloads.sizes[name] > payload_space))
				dlog("#{refname}: Skipping payload #{name} for being too large", 'core',
					LEV_1)
				next
			end

			# Are we compatible in terms of conventions and connections and
			# what not?
			next if (compatible?(framework.payloads.instance(name)) == false)

			# This one be compatible!
			payloads << [ name, mod ]
		}
	
		return payloads;
	end

	#
	# This method returns the number of bytes that should be adjusted to the
	# stack pointer prior to executing any code.  The number of bytes to adjust
	# is indicated to the routine through the payload 'StackAdjustment'
	# attribute or through a target's payload 'StackAdjustment' attribute.
	#
	def stack_adjustment
		if (target and target.payload_stack_adjustment)
			adj = target.payload_stack_adjustment
		else
			adj = payload_info['StackAdjustment']
		end

		# Get the architecture for the current target or use the one specific to
		# this exploit
		arch = (target and target.arch) ? target.arch : self.arch

		# Default to x86 if we can't find a list of architectures
		if (arch and arch.empty? == false)
			arch = arch.join(", ")
		else
			arch = 'x86'
		end

		(adj != nil) ? Rex::Arch::adjust_stack_pointer(arch, adj) || '' : ''
	end

	#
	# Return any text that should be prepended to the payload.  The payload
	# module is passed so that the exploit can take a guess at architecture
	# and platform if it's a multi exploit.  This automatically takes into
	# account any require stack adjustments.
	#
	def payload_prepend(payload_module = self.payload_instance)
		if (target and target.payload_prepend)
			p = target.payload_prepend
		else 
			p = payload_info['Prepend'] || ''
		end

		stack_adjustment + p
	end

	#
	# Return any text that should be appended to the payload.  The payload
	# module is passed so that the exploit can take a guess at architecture
	# and platform if it's a multi exploit.
	#
	def payload_append(payload_module = self.payload_instance)
		if (target and target.payload_append)
			target.payload_append
		else
			payload_info['Append'] || ''
		end
	end

	#
	# Return any text that should be prepended to the encoder of the payload.
	# The payload module is passed so that the exploit can take a guess 
	# at architecture and platform if it's a multi exploit.
	#
	def payload_prepend_encoder(payload_module = self.payload_instance)
		if (target and target.payload_prepend_encoder)
			p = target.payload_prepend_encoder
		else
			p = payload_info['PrependEncoder'] || ''
		end

		p
	end

	#
	# Maximum number of nops to use as a hint to the framework.
	# Nil signifies that the framework should decide.
	#
	def payload_max_nops
		if (target and target.payload_max_nops)
			target.payload_max_nops
		else
			payload_info['MaxNops'] || nil
		end
	end

	#
	# Minimum number of nops to use as a hint to the framework.
	# Nil snigifies that the framework should decide.
	#
	def payload_min_nops
		if (target and target.payload_min_nops)
			target.payload_min_nops
		else
			payload_info['MinNops'] || nil
		end
	end

	#
	# Returns the maximum amount of room the exploit has for a payload.
	#
	def payload_space
		if (target and target.payload_space)
			target.payload_space
		elsif (payload_info['Space'])
			payload_info['Space'].to_i
		else
			nil
		end
	end

	#
	# Returns the bad characters that cannot be in any payload used by this
	# exploit.
	#
	def payload_badchars
		if (target and target.payload_badchars)
			target.payload_badchars
		else
			payload_info['BadChars']
		end
	end

	#
	# Returns the payload encoder type that is associated with either the
	# current target of the exploit in general.
	#
	def payload_encoder_type
		if (target and target.payload_encoder_type)
			target.payload_encoder_type
		else
			payload_info['EncoderType']
		end
	end

	#
	# Returns the payload encoder option hash that is used to initialize the
	# datastore of the encoder that is selected when generating an encoded
	# payload.
	#
	def payload_encoder_options
		if (target and target.payload_encoder_options)
			target.payload_encoder_options
		else
			payload_info['EncoderOptions']
		end
	end

	##
	#
	# NOP requirements
	# 
	# Hints to the nop generator on how it should perform if it's used.
	#
	##

	#
	# Returns the list of registers that the NOP generator should save, 
	# if any.  It will use the current target's save registers in precedence
	# over those defined globally for the exploit module.
	#
	# If there are no save registers, nil is returned.
	#
	def nop_save_registers
		if (target and target.save_registers)
			return target.save_registers 
		else
			return module_info['SaveRegisters']
		end
	end

	#
	# Returns the first compatible NOP generator for this exploit's payload
	# instance.
	#
	def nop_generator
		return nil if (!payload_instance)

		payload_instance.compatible_nops.each { |nopname, nopmod|
			return nopmod.new
		}
	end

	#
	# Generates a nop sled of a supplied length and returns it to the caller.
	#
	def make_nops(count)
		nop_sled = nil

		# If there is no payload instance then we can't succeed.
		return nil if (!payload_instance)

		payload_instance.compatible_nops.each { |nopname, nopmod|
			# Create an instance of the nop module
			nop = nopmod.new

			# The list of save registers
			save_regs = nop_save_registers || []

			if (save_regs.empty? == true)
				save_regs = nil
			end

			begin
				nop.copy_ui(self)

				nop_sled = nop.generate_sled(count,
					'BadChars'      => payload_badchars || '',
					'SaveRegisters' => save_regs)
			rescue
				wlog("#{self.refname}: Nop generator #{nop.refname} failed to generate sled for exploit: #{$!}",
					'core', LEV_0)
			end

			break
		}

		nop_sled
	end

	##
	#
	# Handler interaction
	#
	##
	
	#
	# Passes the connection to the associated payload handler to see if the
	# exploit succeeded and a connection has been established.  The return
	# value can be one of the Handler::constants.
	#
	def handler(*args)
		return self.payload_instance.handler(*args) if (self.payload_instance)
	end

	##
	#
	# Session tracking
	#
	##
	
	#
	# This is called by the payload when a new session is created
	#
	def on_new_session(session)
		self.session_count += 1
	end
	
	#
	# A boolean for whether a session has been created yet
	#
	def session_created?
		(self.session_count > 0) ? true : false
	end

	#
	# Reset the session counter to zero (which occurs during set up of the
	# exploit prior to calling exploit).
	#
	def reset_session_counts
		self.session_count = 0
	end
	
	##
	#
	# Attributes
	#
	##
	
	#
	# The list of targets.
	#
	attr_reader   :targets
	#
	# The default target.
	#
	attr_reader   :default_target
	#
	# The payload requirement hash.
	#
	attr_reader   :payload_info
	#
	# The active payload instance.
	#
	attr_accessor :payload_instance
	#
	# The encoded payload instance.  An instance of an 
	# EncodedPayload object.
	#
	attr_accessor :payload
	
	#
	# The number of active sessions created by this instance
	#
	attr_reader  :session_count

protected

	#
	# Writable copy of the list of targets.
	#
	attr_writer :targets
	#
	# Writable copy of the default target.
	#
	attr_writer :default_target
	#
	# Writable copy of the payload requirement hash.
	#
	attr_writer :payload_info
	#
	#  Number of sessions created by this exploit instance.
	#
	attr_writer :session_count

	#
	# Overrides the base class method and serves to initialize default
	# compatibilities for exploits
	#
	def init_compat
		super

		#
		# Merge in payload compatible defaults 
		#
		p = module_info['Compat']['Payload']

		CompatDefaults::Payload.each_pair { |k,v|
			(p[k]) ? p[k] << " #{v}" : p[k] = v
		}

		#
		# Set the default save registers if none have been explicitly 
		# specified.
		#
		if (module_info['SaveRegisters'] == nil)
			module_info['SaveRegisters'] = [ 'esp', 'ebp' ]
		end
	end

end

end