require 'rex/proto/smb'
require 'rex/proto/dcerpc'
require 'rex/encoder/ndr'

module Msf

###
#
# This mixin provides utility methods for interacting with a SMB/CIFS service on
# a remote machine.  These methods may generally be useful in the context of
# exploitation.  This mixin extends the Tcp exploit mixin. Only one SMB
# service can be accessed at a time using this class.
#
###

module Exploit::Remote::SMB

	include Exploit::Remote::Tcp
	SIMPLE = Rex::Proto::SMB::SimpleClient
	XCEPT  = Rex::Proto::SMB::Exceptions
	CONST  = Rex::Proto::SMB::Constants
	
	# Alias over the Rex DCERPC protocol modules
	DCERPCPacket   = Rex::Proto::DCERPC::Packet
	DCERPCClient   = Rex::Proto::DCERPC::Client
	DCERPCResponse = Rex::Proto::DCERPC::Response
	DCERPCUUID     = Rex::Proto::DCERPC::UUID
	NDR            = Rex::Encoder::NDR

	def initialize(info = {})
		super
	
		register_evasion_options(
		[
			OptBool.new('SMB::pipe_evasion',     [ true, 'Enable segmented read/writes for SMB Pipes', 'False']),
			OptInt.new('SMB::pipe_write_min_size', [ true, 'Minimum buffer size for pipe writes',  1]),
			OptInt.new('SMB::pipe_write_max_size', [ true, 'Maximum buffer size for pipe writes', 1024]),
			OptInt.new('SMB::pipe_read_min_size',  [ true, 'Minimum buffer size for pipe reads',  1]),
			OptInt.new('SMB::pipe_read_max_size',  [ true, 'Maximum buffer size for pipe reads', 1024]),
			OptInt.new('SMB::pad_data_level',  [ true, 'Place extra padding between headers and data (level 0-3)', 0]),
			OptInt.new('SMB::pad_file_level',  [ true, 'Obscure path names used in open/create (level 0-3)', 0]),
			OptInt.new('SMB::obscure_trans_pipe_level',  [ true, 'Obscure PIPE string in TransNamedPipe (level 0-3)', 0]),
			
		], Msf::Exploit::Remote::SMB)

		register_advanced_options(
		[
			OptBool.new('SMBDirect', [ true, 'The target port is a raw SMB service (not NetBIOS)', 'True' ]),
			OptString.new('SMBUser', [ false, 'The username to authenticate as', '']),
			OptString.new('SMBPass', [ false, 'The password for the specified username', '']),
			OptString.new('SMBDomain',  [ false, 'The Windows domain to use for authentication', 'WORKGROUP']),
			OptString.new('SMBName', [ true, 'The NetBIOS hostname (required for port 139 connections)', '*SMBSERVER'])
		], Msf::Exploit::Remote::SMB)

		register_options(
			[
				Opt::RHOST,
				OptInt.new('RPORT', [ true, 'Set the SMB service port', 445])	
			], Msf::Exploit::Remote::SMB)
	end

	def connect()
		
		disconnect()
		
		super
	
		self.simple = SIMPLE.new(self.sock, datastore['SMBDirect'])

		# setup pipe evasion foo
		if datastore['SMB::pipe_evasion']
			# XXX - insert code to change the instance of the read/write functions to do segmentation
		end
		
		if (datastore['SMB::pad_data_level'])
			self.simple.client.evasion_opts['pad_data'] = datastore['SMB::pad_data_level']
		end

		if (datastore['SMB::pad_file_level'])
			self.simple.client.evasion_opts['pad_file'] = datastore['SMB::pad_file_level']
		end

		if (datastore['SMB::obscure_trans_pipe_level'])
			self.simple.client.evasion_opts['obscure_trans_pipe'] = datastore['SMB::obscure_trans_pipe_level']
		end
	end

	# Convert a standard ASCII string to 16-bit Unicode
	def unicode(str)
		Rex::Text.to_unicode(str)
	end
	
	# This method establishes a SMB session over the default socket
	def smb_login
		simple.login(
			datastore['SMBName'], 
			datastore['SMBUser'],
			datastore['SMBPass'],
			datastore['SMBDomain']
		)
		
		simple.connect("\\\\#{datastore['RHOST']}\\IPC$")
	end

	# This method returns the native operating system of the peer	
	def smb_peer_os
		self.simple.client.peer_native_os
	end
	
	# This method returns the native lanman version of the peer
	def smb_peer_lm
		self.simple.client.peer_native_lm
	end
	
	# This method opens a handle to an IPC pipe
	def smb_create(pipe)
		self.simple.create_pipe(pipe)
	end
	
	def smb_hostname
		datastore['SMBName'] || '*SMBSERVER'
	end
	
	attr_accessor :simple

end

###
#
# This mixin provides a minimal SMB server
#
###

module Exploit::Remote::SMBServer
	include Exploit::Remote::TcpServer
	CONST = ::Rex::Proto::SMB::Constants
	CRYPT = ::Rex::Proto::SMB::Crypt
	UTILS = ::Rex::Proto::SMB::Utils
	XCEPT = ::Rex::Proto::SMB::Exceptions
	EVADE = ::Rex::Proto::SMB::Evasions

	def initialize(info = {})
		super

		register_options(
			[
				OptPort.new('SRVPORT',    [ true, "The local port to listen on.", 139 ])
			], self.class)
	end

	def setup
		super
		@state = {}
	end

	def on_client_connect(client)
		# print_status("New SMB connection from #{client.peerhost}:#{client.peerport}")
		smb_conn(client)
	end
	
	def on_client_data(client)
		# print_status("New data from #{client.peerhost}:#{client.peerport}")
		smb_recv(client)
		true
	end
	
	def on_client_close(client)
		smb_stop(client)
	end
	
	def smb_conn(c)
		@state[c] = {:name => "#{c.peerhost}:#{c.peerport}", :ip => c.peerhost, :port => c.peerport}
	end
	
	def smb_stop(c)
		@state.delete(c)
	end
	
	def smb_recv(c)
		smb = @state[c]
		smb[:data] ||= ''
		smb[:data] << c.get_once
		
		while(smb[:data].length > 0)
		
			return if smb[:data].length < 4

			plen = smb[:data][2,2].unpack('n')[0]

			return if smb[:data].length < plen+4

			buff = smb[:data].slice!(0, plen+4)

			pkt_nbs = CONST::NBRAW_PKT.make_struct
			pkt_nbs.from_s(buff)

			# print_status("NetBIOS request from #{smb[:name]} #{pkt_nbs.v['Type']} #{pkt_nbs.v['Flags']} #{buff.inspect}")

			# Check for a NetBIOS name request
			if (pkt_nbs.v['Type'] == 0x81)
				# Accept any name they happen to send
				
				host_dst = UTILS.nbname_decode(pkt_nbs.v['Payload'][1,32]).gsub(/[\x00\x20]+$/, '')
				host_src = UTILS.nbname_decode(pkt_nbs.v['Payload'][35,32]).gsub(/[\x00\x20]+$/, '')
				
				smb[:nbdst] = host_dst
				smb[:nbsrc] = host_src
				
				# print_status("NetBIOS session request from #{smb[:name]} (asking for #{host_dst} from #{host_src})")
				c.write("\x82\x00\x00\x00")
				next
			end
			
			
			# 
			# TODO: Support AndX parameters
			# 
			

			# Cast this to a generic SMB structure
			pkt = CONST::SMB_BASE_PKT.make_struct
			pkt.from_s(buff)

			# Only response to requests, ignore server replies
			if (pkt['Payload']['SMB'].v['Flags1'] & 128 != 0)
				print_status("Ignoring server response from #{smb[:name]}")		
				next
			end
			
			cmd = pkt['Payload']['SMB'].v['Command']
			begin
				smb_cmd_dispatch(cmd, c, buff)
			rescue ::Interrupt
				raise $!
			rescue ::Exception => e
				print_status("Error processing request from #{smb[:name]} (#{cmd}): #{e.class} #{e.to_s} #{e.backtrace.to_s}")
				next				
			end		
		end	
	end
	
	def smb_cmd_dispatch(cmd, c, buff)
		smb = @state[c]		
		print_status("Received command #{cmd} from #{smb[:name]}")	
	end
	
	def smb_set_defaults(c, pkt)
		smb = @state[c]
		pkt['Payload']['SMB'].v['ProcessID'] = smb[:process_id].to_i
		pkt['Payload']['SMB'].v['UserID'] = smb[:user_id].to_i
		pkt['Payload']['SMB'].v['TreeID'] = smb[:tree_id].to_i
		pkt['Payload']['SMB'].v['MultiplexID'] = smb[:multiplex_id].to_i
	end
	
end


end