udp

git-svn-id: file:///home/svn/incoming/trunk@2554 4d416f70-5f16-0410-b530-b9f4589650da

lib/rex.rb.ts.rb

@@ -6,6 +6,7 @@ require 'Rex/Socket.rb.ut'
 require 'Rex/Socket/Tcp.rb.ut'
 require 'Rex/Socket/SslTcp.rb.ut'
 require 'Rex/Socket/TcpServer.rb.ut'
+require 'Rex/Socket/Udp.rb.ut'
 require 'Rex/Socket/Parameters.rb.ut'
 require 'Rex/Socket/Comm/Local.rb.ut'
 
@@ -18,6 +19,7 @@ class Rex::TestSuite
 		suite << Rex::Socket::Tcp::UnitTest.suite
 		suite << Rex::Socket::SslTcp::UnitTest.suite
 		suite << Rex::Socket::TcpServer::UnitTest.suite
+		suite << Rex::Socket::Udp::UnitTest.suite
 		suite << Rex::Socket::Comm::Local::UnitTest.suite
 
 		return suite;

lib/rex/io/stream.rb

@@ -136,7 +136,7 @@ module Stream
 	#
 	def get(timeout = def_read_timeout, ltimeout = def_read_loop_timeout, opts = {})
 		# For those people who are used to being able to use a negative timeout!
-		if (timeout < 0)
+		if (timeout and timeout.to_i < 0)
 			timeout = nil
 		end
 

lib/rex/socket.rb

@@ -24,7 +24,7 @@ class Socket
 	#
 	##
 
-	def self.create(opts)
+	def self.create(opts = {})
 		return create_param(Rex::Socket::Parameters.from_hash(opts))
 	end
 
@@ -32,7 +32,7 @@ class Socket
 		return param.comm.create(param)
 	end
 
-	def self.create_tcp(opts)
+	def self.create_tcp(opts = {})
 		return create_param(Rex::Socket::Parameters.from_hash(opts.merge('Proto' => 'tcp')))
 	end
 	
@@ -40,6 +40,10 @@ class Socket
 		return create_tcp(opts.merge('Server' => true))
 	end
 
+	def self.create_udp(opts = {})
+		return create_param(Rex::Socket::Parameters.from_hash(opts.merge('Proto' => 'udp')))
+	end
+
 	##
 	#
 	# Serialization
@@ -58,6 +62,19 @@ class Socket
 		return data.pack('snCCCCa8')
 	end
 
+	#
+	# Returns the address family, host, and port of the supplied sockaddr as
+	# [ af, host, port ]
+	#
+	def self.from_sockaddr(saddr)
+		up = saddr.unpack('snCCCC')
+
+		af   = up.shift
+		port = up.shift
+
+		return [ af, up.join('.'), port ]
+	end
+
 	##
 	#
 	# Class initialization

lib/rex/socket.rb.ut.rb

@@ -37,4 +37,21 @@ class Rex::Socket::UnitTest < Test::Unit::TestCase
 		assert_equal("\x02\x00\x00\x16\x01\x02\x03\x04" + "\x00" * 8, Rex::Socket.to_sockaddr("1.2.3.4", 22), "1.2.3.4 addr, port 22 sockaddr")
 	end
 
+	def test_from_sockaddr
+		af, host, port = Rex::Socket.from_sockaddr("\x00" * 16)
+		assert_equal(0, af, "zero af")
+		assert_equal('0.0.0.0', host, "zero host")
+		assert_equal(0, port, "zero port")
+
+		af, host, port = Rex::Socket.from_sockaddr("\x02\x00\x00\x16" + "\x00" * 12)
+		assert_equal(2, af, "af = 2")
+		assert_equal('0.0.0.0', host, "zero host")
+		assert_equal(22, port, "port = 22")
+
+		af, host, port = Rex::Socket.from_sockaddr("\x02\x00\x00\x16\x01\x02\x03\x04" + "\x00" * 8)
+		assert_equal(2, af, "af = 2")
+		assert_equal('1.2.3.4', host, "zero host")
+		assert_equal(22, port, "port = 22")
+	end
+
 end

lib/rex/socket/comm/local.rb

@@ -1,5 +1,7 @@
 require 'Rex/Socket'
 require 'Rex/Socket/Tcp'
+require 'Rex/Socket/SslTcp'
+require 'Rex/Socket/Udp'
 
 ###
 #
@@ -62,6 +64,8 @@ class Rex::Socket::Comm::Local
 					end
 
 					return klass.new(sock, param)
+				when 'udp'
+					return Rex::Socket::Udp.new(sock, param)
 			end
 		end
 	end

lib/rex/socket/udp.rb

@@ -0,0 +1,85 @@
+require 'Rex/Socket'
+
+###
+#
+# Udp
+# ---
+#
+# This class provides methods for interacting with a UDP socket.
+#
+###
+class Rex::Socket::Udp < Rex::Socket
+
+	##
+	#
+	# Factory
+	#
+	##
+
+	#
+	# Wrapper around the base socket class' creation method that automatically
+	# sets the parameter's protocol to UDP
+	#
+	def self.create_param(param)
+		param.proto = 'udp'
+
+		super(param)
+	end
+
+	##
+	#
+	# UDP connected state methods
+	#
+	##
+
+	#
+	# Write the supplied datagram to the connected UDP socket
+	#
+	def write(gram)
+		return sock.write(gram)
+	end
+
+	#
+	# Read a datagram from the UDP socket
+	#
+	def read(length = 65535)
+		return sock.read(length)
+	end
+
+	#alias send write
+	#alias recv read
+
+	##
+	#
+	# UDP non-connected state methods
+	#
+	##
+
+	#
+	# Sends a datagram to the supplied host:port with optional flags
+	#
+	def sendto(gram, peerhost, peerport, flags = 0)
+		return sock.send(gram, flags, Rex::Socket.to_sockaddr(peerhost, peerport))
+	end
+
+	#
+	# Receives a datagram and returns the data and host:port of the requestor
+	# as [ data, host, port ]
+	#
+	def recvfrom(length = 65535)
+		data, saddr    = sock.recvfrom(length)
+		af, host, port = Rex::Socket.from_sockaddr(saddr)
+
+		return [ data, host, port ]
+	end
+
+	#
+	# Checks for whether or not any datagrams are pending read
+	#
+	def has_read_data?(timeout = nil)
+		timeout = timeout.to_i if (timeout)
+
+		return (select([ poll_fd ], nil, nil, timeout) != nil)
+	end
+
+end

lib/rex/socket/udp.rb.ut.rb

@@ -0,0 +1,44 @@
+#!/usr/bin/ruby
+
+$:.unshift(File.join('..', '..', File.dirname(__FILE__)))
+
+require 'test/unit'
+require 'Rex/Socket/Udp'
+
+class Rex::Socket::Udp::UnitTest < Test::Unit::TestCase
+
+	def test_udp
+		serv_port = 65432
+		serv = Rex::Socket::Udp.create(
+			'LocalHost' => '127.0.0.1',
+			'LocalPort' => serv_port)
+
+		begin
+		  assert_kind_of(Rex::Socket::Udp, serv, "valid Udp server instance")
+
+		  # Test connected socket
+		  concli = Rex::Socket::Udp.create(
+			  'PeerHost' => '127.0.0.1',
+			  'PeerPort' => serv_port)
+
+		  assert_equal('127.0.0.1', concli.peerhost, "matching peerhost")
+		  assert_equal(serv_port, concli.peerport, "matching peerport")
+		  assert_equal(2, concli.write('yo'), "write succeeded")
+
+		  data, host, port = serv.recvfrom
+
+		  assert_equal('yo', data, "read data match")
+		  assert_equal('127.0.0.1', host, "matching client host")
+
+			# Test non-connected socket
+			concli = Rex::Socket::Udp.create
+
+			assert_equal(3, concli.sendto('bob', '127.0.0.1', serv_port), "sendto")
+		   data, host, port = serv.recvfrom
+			assert_equal('bob', data, "read data match")
+		ensure
+			serv.close
+		end
+	end
+
+end