module Msf

module EvasiveTCP
	attr_accessor :_send_size, :_send_delay, :evasive

	def denagle
		begin
			setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
		rescue ::Exception
		end
	end

	def write(buf, opts={})

		return super(buf, opts) if not @evasive

		ret = 0
		idx = 0
		len = @_send_size || buf.length

		while(idx < buf.length)

			if(@_send_delay and idx > 0)
				::IO.select(nil, nil, nil, @_send_delay)
			end

			pkt = buf[idx, len]

			res = super(pkt, opts)
			flush()

			idx += len
			ret += res if res
		end
		ret
	end
end

###
#
# This module provides methods for establish a connection to a remote host and
# communicating with it.
#
###
module Exploit::Remote::Tcp

	#
	# Initializes an instance of an exploit module that exploits a
	# vulnerability in a TCP server.
	#
	def initialize(info = {})
		super

		register_options(
			[
				Opt::RHOST,
				Opt::RPORT
			], Msf::Exploit::Remote::Tcp
		)

		register_advanced_options(
			[
				OptBool.new('SSL',        [ false, 'Negotiate SSL for outgoing connections', false]),
				OptEnum.new('SSLVersion', [ false, 'Specify the version of SSL that should be used', 'SSL3', ['SSL2', 'SSL3', 'TLS1']]),
				Opt::Proxies,
				Opt::CPORT,
				Opt::CHOST,
				OptInt.new('ConnectTimeout', [ true, 'Maximum number of seconds to establish a TCP connection', 10])
			], Msf::Exploit::Remote::Tcp
		)

		register_evasion_options(
			[
				OptInt.new('TCP::max_send_size', [false, 'Maxiumum tcp segment size.  (0 = disable)', 0]),
				OptInt.new('TCP::send_delay', [false, 'Delays inserted before every send.  (0 = disable)', 0])
			], Msf::Exploit::Remote::Tcp
		)
	end

	#
	# Establishes a TCP connection to the specified RHOST/RPORT
	#
	def connect(global = true, opts={})

		dossl = false
		if(opts.has_key?('SSL'))
			dossl = opts['SSL']
		else
			dossl = ssl
			if (datastore.default?('SSL') and rport.to_i == 443)
				dossl = true
			end
		end

		nsock = Rex::Socket::Tcp.create(
			'PeerHost'  =>  opts['RHOST'] || rhost,
			'PeerPort'  => (opts['RPORT'] || rport).to_i,
			'LocalHost' =>  opts['CHOST'] || chost || "0.0.0.0",
			'LocalPort' => (opts['CPORT'] || cport || 0).to_i,
			'SSL'       =>  dossl,
			'SSLVersion'=>  opts['SSLVersion'] || ssl_version,
			'Proxies'   => proxies,
			'Timeout'   => (opts['ConnectTimeout'] || connect_timeout || 10).to_i,
			'Context'   =>
				{
					'Msf'        => framework,
					'MsfExploit' => self,
				})

		# enable evasions on this socket
		set_tcp_evasions(nsock)

		# Set this socket to the global socket as necessary
		self.sock = nsock if (global)

		# Add this socket to the list of sockets created by this exploit
		add_socket(nsock)

		return nsock
	end

	# Enable evasions on a given client
	def set_tcp_evasions(socket)

		if( datastore['TCP::max_send_size'].to_i == 0 and datastore['TCP::send_delay'].to_i == 0)
			return
		end

		return if socket.respond_to?('evasive')

		socket.extend(EvasiveTCP)

		if ( datastore['TCP::max_send_size'].to_i > 0)
			socket._send_size = datastore['TCP::max_send_size']
			socket.denagle
			socket.evasive = true
		end

		if ( datastore['TCP::send_delay'].to_i > 0)
			socket._send_delay = datastore['TCP::send_delay']
			socket.evasive = true
		end
	end

	def handler(nsock = self.sock)
		# If the handler claims the socket, then we don't want it to get closed
		# during cleanup
		if ((rv = super) == Handler::Claimed)
			if (nsock == self.sock)
				self.sock = nil
			end

			# Remove this socket from the list of sockets so that it will not be
			# aborted.
			remove_socket(nsock)
		end

		return rv
	end

	#
	# Closes the TCP connection
	#
	def disconnect(nsock = self.sock)
		begin
			if (nsock)
				nsock.shutdown
				nsock.close
			end
		rescue IOError
		end

		if (nsock == sock)
			self.sock = nil
		end

		# Remove this socket from the list of sockets created by this exploit
		remove_socket(nsock)
	end

	#
	# Performs cleanup, disconnects the socket if necessary
	#
	def cleanup
		super
		disconnect
	end

	##
	#
	# Wrappers for getters
	#
	##

	#
	# Returns the target host
	#
	def rhost
		datastore['RHOST']
	end

	#
	# Returns the remote port
	#
	def rport
		datastore['RPORT']
	end

	#
	# Returns the local host
	#
	def lhost
		datastore['LHOST']
	end

	#
	# Returns the local port
	#
	def lport
		datastore['LPORT']
	end

	#
	# Returns the local host for outgoing connections
	#
	def chost
		datastore['CHOST']
	end

	#
	# Returns the local port for outgoing connections
	#
	def cport
		datastore['CPORT']
	end

	#
	# Returns the boolean indicating SSL
	#
	def ssl
		datastore['SSL']
	end

	#
	# Returns the string indicating SSLVersion
	#
	def ssl_version
		datastore['SSLVersion']
	end

	#
	# Returns the proxy configuration
	#
	def proxies
		datastore['Proxies']
	end

	#
	# Returns the TCP connection timeout
	#
	def connect_timeout
		datastore['ConnectTimeout']
	end

protected

	attr_accessor :sock

end


###
#
# This mixin provides a generic interface for running a TCP server of some
# sort that is designed to exploit clients.  Exploits that include this mixin
# automatically take a passive stance.
#
###
module Exploit::Remote::TcpServer

	def initialize(info = {})
		super(update_info(info,
			'Stance' => Msf::Exploit::Stance::Passive))

		register_options(
			[
				OptBool.new('SSL',        [ false, 'Negotiate SSL for incoming connections', false]),
				OptEnum.new('SSLVersion', [ false, 'Specify the version of SSL that should be used', 'SSL3', ['SSL2', 'SSL3', 'TLS1']]),
				OptAddress.new('SRVHOST', [ true, "The local host to listen on. This must be an address on the local machine or 0.0.0.0", '0.0.0.0' ]),
				OptPort.new('SRVPORT',    [ true, "The local port to listen on.", 8080 ]),
				
			], Msf::Exploit::Remote::TcpServer)

		register_advanced_options(
			[
				OptString.new('ListenerComm', [ false, 'The specific communication channel to use for this service']),
			], Msf::Exploit::Remote::TcpServer)

		register_evasion_options(
			[
				OptInt.new('TCP::max_send_size', [false, 'Maximum tcp segment size.  (0 = disable)', 0]),
				OptInt.new('TCP::send_delay', [false, 'Delays inserted before every send.  (0 = disable)', 0])
			], Msf::Exploit::Remote::Tcp
		)
	end

	#
	# This mixin overrides the exploit method so that it can initiate the
	# service that corresponds with what the client has requested.
	#
	def exploit

		start_service()
		print_status("Server started.")

		# Call the exploit primer
		primer

		# Wait on the service to stop
		self.service.wait
	end

	#
	# Primer method to call after starting service but before handling connections
	#
	def primer
	end

	#
	# Stops the service, if one was created.
	#
	def cleanup
		super
		if(service)
			stop_service()
			print_status("Server stopped.")
		end
	end

	#
	# Called when a client connects.
	#
	def on_client_connect(client)
	end

	#
	# Called when a client has data available for reading.
	#
	def on_client_data(client)
	end

	#
	# Called when a client has disconnected.
	#
	def on_client_close(client)
	end

	#
	# Starts the service.
	#
	def start_service(*args)
		begin

			comm = datastore['ListenerComm']
			if comm == "local"
				comm = ::Rex::Socket::Comm::Local
			else
				comm = nil
			end
			
			self.service = Rex::Socket::TcpServer.create(
				'LocalHost' => srvhost,
				'LocalPort' => srvport,
				'SSL'       => ssl,
				'Comm'      => comm,
				'Context'   =>
					{
						'Msf'        => framework,
						'MsfExploit' => self,
					})

			self.service.on_client_connect_proc = Proc.new { |client|
				on_client_connect(client)
			}
			self.service.on_client_data_proc = Proc.new { |client|
				on_client_data(client)
			}
			self.service.on_client_close_proc = Proc.new { |client|
				on_client_close(client)
			}

			# Start the listening service
			self.service.start

		rescue ::Errno::EACCES => e
			if (srvport.to_i < 1024)
				print_line(" ")
				print_error("Could not start the TCP server: #{e}.")
				print_error(
					"This module is configured to use a privileged TCP port (#{srvport}). " +
					"On Unix systems, only the root user account is allowed to bind to privileged ports." +
					"Please run the framework as root to use this module."
				)
				print_error(
					"On Microsoft Windows systems, this error is returned when a process attempts to "+
					"listen on a host/port combination that is already in use. For example, Windows XP "+
					"will return this error if a process attempts to bind() over the system SMB/NetBIOS services."
				)
				print_line(" ")
			end
			raise e
		end
	end

	#
	# Stops the service.
	#
	def stop_service
		if (service)
			begin
				self.service.deref if self.service.kind_of?(Rex::Service)
				if self.service.kind_of?(Rex::Socket)
					self.service.close
					self.service.stop 
				end
				self.service = nil
			rescue ::Exception
			end
		end
	end

	#
	# Returns the local host that is being listened on.
	#
	def srvhost
		datastore['SRVHOST']
	end

	#
	# Returns the local port that is being listened on.
	#
	def srvport
		datastore['SRVPORT']
	end

	#
	# Returns the SSL option
	#
	def ssl
		datastore['SSL']
	end


	#
	# Re-generates the payload, substituting the current RHOST and RPORT with
	# the supplied client host and port from the socket.
	#
	def regenerate_payload(cli, arch = nil, platform = nil, target = nil)

		ohost = datastore['RHOST']
		oport = datastore['RPORT']
		p = nil

		begin
			# Update the datastore with the supplied client peerhost/peerport
			datastore['RHOST'] = cli.peerhost
			datastore['RPORT'] = cli.peerport

			if ((p = super(arch, platform, target)) == nil)
				print_error("Failed to generate payload")
				return nil
			end

			# Allow the payload to start a new handler
			add_handler({
				'RHOST' => datastore['RHOST'],
				'RPORT' => datastore['RPORT']
			})

		ensure
			datastore['RHOST'] = ohost
			datastore['RPORT'] = oport
		end

		p
	end

protected

	attr_accessor :service # :nodoc:

end

end