librex 0.0.1

data/lib/rex/socket/switch_board.rb.ut.rb

@@ -0,0 +1,52 @@
+#!/usr/bin/env ruby
+
+$:.unshift(File.join(File.dirname(__FILE__), '..', '..'))
+
+require 'test/unit'
+require 'rex/socket/switch_board'
+
+class Rex::Socket::SwitchBoard::UnitTest < Test::Unit::TestCase
+
+	Klass = Rex::Socket::SwitchBoard
+
+	def test_add
+		Klass.flush_routes
+		assert_equal(true, Klass.add_route('0.0.0.0', 0, 'foo'))
+		assert_equal(false, Klass.add_route('0.0.0.0', 0, 'foo'))
+		assert_equal(1, Klass.routes.length)
+
+		assert_equal('0.0.0.0', Klass.routes[0].subnet)
+		assert_equal('0.0.0.0', Klass.routes[0].netmask)
+		assert_equal(0, Klass.routes[0].bitmask)
+		assert_equal('foo', Klass.routes[0].comm)
+	end
+
+	def test_remove
+		Klass.flush_routes
+		assert_equal(true, Klass.add_route('0.0.0.0', 0, 'foo'))
+		assert_equal(true, Klass.remove_route('0.0.0.0', 0, 'foo'))
+		assert_equal(false, Klass.remove_route('0.0.0.0', 0, 'foo'))
+		assert_equal(0, Klass.routes.length)
+	end
+
+	def test_best_comm
+		Klass.flush_routes
+		Klass.add_route('0.0.0.0', 0, 'default')
+		Klass.add_route('1.2.3.0', 24, 'spec')
+
+		assert_equal('default', Klass.best_comm('4.5.6.7'))
+		assert_equal('spec', Klass.best_comm('1.2.3.7'))
+	end
+
+	def test_remove_by_comm
+		Klass.flush_routes
+		Klass.add_route('1.2.3.0', 24, 'foo')
+		Klass.add_route('1.2.4.0', 24, 'dog')
+
+		Klass.remove_by_comm('foo')
+
+		assert_equal('dog', Klass.best_comm('1.2.4.7'))
+		assert_nil(Klass.best_comm('1.2.3.7'))
+	end
+
+end

data/lib/rex/socket/tcp.rb

@@ -0,0 +1,76 @@
+require 'rex/socket'
+require 'rex/io/stream'
+
+###
+#
+# This class provides methods for interacting with a TCP client connection.
+#
+###
+module Rex::Socket::Tcp
+
+	include Rex::Socket
+	include Rex::IO::Stream
+	
+	##
+	#
+	# Factory
+	#
+	##
+
+	#
+	# Creates the client using the supplied hash.
+	#
+	def self.create(hash = {})
+		hash['Proto'] = 'tcp' 
+		self.create_param(Rex::Socket::Parameters.from_hash(hash))
+	end
+
+	#
+	# Wrapper around the base socket class' creation method that automatically
+	# sets the parameter's protocol to TCP.
+	#
+	def self.create_param(param)
+		param.proto = 'tcp'
+		Rex::Socket.create_param(param)
+	end
+
+	##
+	#
+	# Stream mixin implementations
+	#
+	##
+
+	#
+	# Calls shutdown on the TCP connection.
+	#
+	def shutdown(how = ::Socket::SHUT_RDWR)
+		begin
+			return (super(how) == 0)
+		rescue ::Exception
+		end
+	end
+
+	#
+	# Returns peer information (host + port) in host:port format.
+	#
+	def peerinfo
+		if (pi = getpeername)
+			return pi[1] + ':' + pi[2].to_s
+		end
+	end
+
+	#
+	# Returns local information (host + port) in host:port format.
+	#
+	def localinfo
+		if (pi = getlocalname)
+			return pi[1] + ':' + pi[2].to_s
+		end
+	end
+
+	# returns socket type
+	def type?
+		return 'tcp'
+	end
+
+end

data/lib/rex/socket/tcp.rb.ut.rb

@@ -0,0 +1,64 @@
+#!/usr/bin/env ruby
+
+$:.unshift(File.join(File.dirname(__FILE__), '..', '..'))
+
+require 'test/unit'
+require 'rex'
+
+class Rex::Socket::Tcp::UnitTest < Test::Unit::TestCase
+
+	def test_tcp
+		port = 65434
+		listener = Rex::Socket.create_tcp_server( 'LocalPort' => port )
+		client = nil
+
+		begin
+			# Connect to the temp server
+			assert_nothing_raised {
+				client = Rex::Socket.create_tcp(
+					'PeerHost' => '127.0.0.1',
+					'PeerPort' => port)
+			}
+
+			assert_kind_of(Rex::Socket::Tcp, client, 'kindof?')
+			assert_equal('127.0.0.1', client.peerhost, 'peerhost')
+			assert_equal(port, client.peerport, 'peerport')
+
+			# Accept the client connection
+			server = listener.accept
+			assert_kind_of(Socket, server, "valid server socket connection")
+
+			# do all of the tests, once for each side
+			{ 'c/s' => [client, server], 's/c' => [server, client] }.each_pair { |mode, sockets|
+				a = sockets[0]
+				b = sockets[1]
+
+				string = "test\n"
+				assert_equal(false, a.has_read_data?(1), "#{mode} : has_read_data?, no data")
+				assert_equal(string.length, b.write(string), "#{mode} : write")
+				assert_equal(true, a.has_read_data?(1), "#{mode} : has_read_data?, with data")
+				assert_equal(string, a.recv(string.length), "#{mode} : recv")
+
+				string = "string\rtest\nwith\x00null"
+				assert_equal(string.length, a << string, "#{mode} : append")
+				tmp = ''; tmp = b.>>
+				assert_equal(string, tmp, "#{mode} : append (reverse)")
+
+				string = "\x00foobar\x00"
+				assert_equal(string.length, a.send(string, 0), "#{mode} : send")
+				assert_equal(string, b.get(), "#{mode} : get")
+			}
+
+			assert_equal(true, client.shutdown(::Socket::SHUT_RD), 'client: shutdown read handle')
+			assert_equal(true, client.shutdown(::Socket::SHUT_WR), 'client: shutdown write handle')
+			assert_nothing_raised {
+				client.close
+				client = nil
+			}
+		ensure
+			client.close if (client)
+			listener.close
+		end
+	end
+
+end

data/lib/rex/socket/tcp_server.rb

@@ -0,0 +1,67 @@
+require 'rex/socket'
+require 'rex/socket/tcp'
+require 'rex/io/stream_server'
+
+###
+#
+# This class provides methods for interacting with a TCP server.  It
+# implements the StreamServer IO interface.
+#
+#
+###
+module  Rex::Socket::TcpServer
+
+	include Rex::Socket
+	include Rex::IO::StreamServer
+
+	##
+	#
+	# Factory
+	#
+	##
+
+	#
+	# Creates the server using the supplied hash.
+	#
+	def self.create(hash = {})
+		hash['Proto'] = 'tcp'
+		hash['Server'] = true
+		self.create_param(Rex::Socket::Parameters.from_hash(hash))
+	end
+
+	#
+	# Wrapper around the base class' creation method that automatically sets
+	# the parameter's protocol to TCP and sets the server flag to true.
+	#
+	def self.create_param(param)
+		param.proto  = 'tcp'
+		param.server = true
+		Rex::Socket.create_param(param)
+	end
+
+	#
+	# Accepts a child connection.
+	#
+	def accept(opts = {})
+		t = super()
+
+		# jRuby compatibility
+		if t.respond_to?('[]')
+			t = t[0]
+		end
+
+		if (t)
+			t.extend(Rex::Socket::Tcp)
+			t.context = self.context
+
+			pn = t.getpeername
+
+			t.peerhost = pn[1]
+			t.peerport = pn[2]
+		end
+
+		t
+	end
+
+end
+

data/lib/rex/socket/tcp_server.rb.ut.rb

@@ -0,0 +1,44 @@
+#!/usr/bin/env ruby
+
+$:.unshift(File.join(File.dirname(__FILE__), '..', '..'))
+
+require 'test/unit'
+require 'rex/socket/tcp_server'
+
+class Rex::Socket::TcpServer::UnitTest < Test::Unit::TestCase
+
+	def test_tcp_server
+		serv_port = 65433
+		serv = Rex::Socket.create_tcp_server(
+			'LocalPort' => serv_port)
+		ccli = nil
+
+		begin
+			assert_kind_of(Rex::Socket::TcpServer, serv, "valid TcpServer")
+			assert_kind_of(Rex::IO::StreamServer, serv, "valid StreamServer")
+
+			# Connect to the server
+			assert_nothing_raised {
+				ccli = Rex::Socket.create_tcp(
+					'PeerHost' => '127.0.0.1',
+					'PeerPort' => serv_port)
+			}
+			assert_kind_of(Rex::Socket::Tcp, ccli, "valid client client Tcp")
+
+			# Accept the client connection
+			scli = serv.accept
+			assert_kind_of(Rex::Socket::Tcp, scli, "valid server client Tcp")
+
+			assert_equal(2, scli.put("Yo"), "scli: put Yo")
+			assert_equal("Yo", ccli.get(), "ccli: get Yo")
+			assert(scli.methods.include?('<<'))
+			assert(scli.methods.include?('>>'))
+			assert(scli.methods.include?('has_read_data?'))
+			
+		ensure
+			ccli.close if (ccli)
+			serv.close
+		end
+	end
+
+end

data/lib/rex/socket/udp.rb

@@ -0,0 +1,157 @@
+require 'rex/socket'
+
+###
+#
+# This class provides methods for interacting with a UDP socket.
+#
+###
+module Rex::Socket::Udp
+
+	include Rex::Socket
+
+	##
+	#
+	# Factory
+	#
+	##
+
+	#
+	# Creates the client using the supplied hash.
+	#
+	def self.create(hash = {})
+		hash['Proto'] = 'udp'
+		# If we have are to bind to a LocalHost we must be a Server to avail of pivoting.
+		# Rex::Socket::Parameters will subsequently turn off the sever flag after the correct
+		# comm has been chosen.
+		if( hash['LocalHost'] )
+			hash['Server'] = true
+		end
+		self.create_param(Rex::Socket::Parameters.from_hash(hash))
+	end
+
+	#
+	# 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'
+		Rex::Socket.create_param(param)
+	end
+
+	##
+	#
+	# UDP connected state methods
+	#
+	##
+
+	#
+	# Write the supplied datagram to the connected UDP socket.
+	#
+	def write(gram)
+		begin
+			return syswrite(gram)
+		rescue  ::Errno::EHOSTUNREACH,::Errno::ENETDOWN,::Errno::ENETUNREACH,::Errno::ENETRESET,::Errno::EHOSTDOWN,::Errno::EACCES,::Errno::EINVAL,::Errno::EADDRNOTAVAIL
+			return nil
+		end
+	end
+
+	alias put write
+
+	#
+	# Read a datagram from the UDP socket.
+	#
+	def read(length = 65535)
+		if length < 0
+			length = 65535
+		end
+		return sysread(length)
+	end
+
+	#
+	# Read a datagram from the UDP socket with a timeout
+	#
+	def timed_read(length = 65535, timeout=def_read_timeout)
+		begin
+			if ((rv = Kernel.select([ fd ], nil, nil, timeout)) and
+			    (rv[0]) and (rv[0][0] == fd)
+			   )
+					return read(length)
+			else
+				return ''
+			end
+		rescue Exception
+			return ''
+		end
+	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)
+
+		# Catch unconnected IPv6 sockets talking to IPv4 addresses
+		peer = Rex::Socket.resolv_nbo(peerhost)
+		if (peer.length == 4 and self.ipv == 6)
+			peerhost = '::ffff:' + Rex::Socket.getaddress(peerhost)
+		end
+
+		begin
+			send(gram, flags, Rex::Socket.to_sockaddr(peerhost, peerport))
+		rescue  ::Errno::EHOSTUNREACH,::Errno::ENETDOWN,::Errno::ENETUNREACH,::Errno::ENETRESET,::Errno::EHOSTDOWN,::Errno::EACCES,::Errno::EINVAL,::Errno::EADDRNOTAVAIL
+			return nil
+		end
+
+	end
+
+	#
+	# Receives a datagram and returns the data and host:port of the requestor
+	# as [ data, host, port ].
+	#
+	def recvfrom(length = 65535, timeout=def_read_timeout)
+
+		begin
+			if ((rv = Kernel.select([ fd ], nil, nil, timeout)) and
+			    (rv[0]) and (rv[0][0] == fd)
+			   )
+					data, saddr    = recvfrom_nonblock(length)
+					af, host, port = Rex::Socket.from_sockaddr(saddr)
+
+					return [ data, host, port ]
+			else
+				return [ '', nil, nil ]
+			end
+		rescue Exception
+			return [ '', nil, nil ]
+		end
+	end
+
+	#
+	# Calls recvfrom and only returns the data
+	#
+	def get(timeout=nil)
+		data, saddr, sport = recvfrom(65535, timeout)
+		return data
+	end
+
+	#
+	# The default number of seconds to wait for a read operation to timeout.
+	#
+	def def_read_timeout
+		10
+	end
+
+	def type?
+		return 'udp'
+	end
+
+end
+