librex 0.0.1

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

@@ -0,0 +1,44 @@
+#!/usr/bin/env 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 = 55432
+		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

data/lib/rex/struct2.rb

@@ -0,0 +1,5 @@
+#!/usr/bin/env ruby
+
+# just a shim to load all of the Struct2 libraries
+
+require 'rex/struct2/c_struct_template'

data/lib/rex/struct2/c_struct.rb

@@ -0,0 +1,181 @@
+#!/usr/bin/env ruby
+
+# Rex::Struct2
+module Rex
+module Struct2
+
+require 'rex/struct2/s_struct'
+
+class CStruct_Values
+
+	def initialize(obj)
+		@obj = obj
+	end
+
+	def [](*args)
+		o = @obj[*args]
+		return if !o
+		return o.value
+	end
+
+	def []=(*args)
+		o = @obj[*args[0 .. -2]]
+		return if !o
+		o.value = args[-1]
+	end
+
+	# this one is for HD, the whiniest girl around...
+	# allow for like v.field = whatever
+	def method_missing(sym, *args)
+		if sym.to_s[-1] == "="[0]
+			return self[sym.to_s[0 .. -2]] = args[0]
+		else
+			return self[sym.to_s]
+		end
+	end
+end
+
+class CStruct < SStruct
+
+	require 'rex/struct2/element'
+	require 'rex/struct2/generic'
+	require 'rex/struct2/s_string'
+	require 'rex/struct2/c_struct_template'
+	require 'rex/struct2/restraint'
+
+	include Rex::Struct2::Element
+
+	attr_reader  :v
+
+	@@dt_table = {
+	  'int8'      => proc { |*a| Rex::Struct2::Generic.new('C',  true, *a) },
+	  'uint8'     => proc { |*a| Rex::Struct2::Generic.new('C', false, *a) },
+	  'int16v'    => proc { |*a| Rex::Struct2::Generic.new('v',  true, *a) },
+	  'uint16v'   => proc { |*a| Rex::Struct2::Generic.new('v', false, *a) },
+	  'int32v'    => proc { |*a| Rex::Struct2::Generic.new('V',  true, *a) },
+	  'uint32v'   => proc { |*a| Rex::Struct2::Generic.new('V', false, *a) },
+	  'int64v'    => proc { |*a| Rex::Struct2::Generic.new('q',  true, *a) },
+	  'uint64v'   => proc { |*a| Rex::Struct2::Generic.new('Q', false, *a) },
+	  'int16n'    => proc { |*a| Rex::Struct2::Generic.new('n',  true, *a) },
+	  'uint16n'   => proc { |*a| Rex::Struct2::Generic.new('n', false, *a) },
+	  'int32n'    => proc { |*a| Rex::Struct2::Generic.new('N',  true, *a) },
+	  'uint32n'   => proc { |*a| Rex::Struct2::Generic.new('N', false, *a) },
+	  'string'    => proc { |*a| Rex::Struct2::SString.new(*a) },
+	  'sstruct'   => proc { |*a| Rex::Struct2::SStruct.new(*a) },
+	  'object'    => proc { |o| o },
+	  'template'  => proc { |o| o.make_struct },
+	}
+
+	# CStruct.typedef(name, factory, ... )
+	def CStruct.typedef(*args)
+		while args.length >= 2
+			name    = args.shift
+			factory = args.shift
+			@@dt_table[name] = factory
+		end
+	end
+
+	def initialize(*dts)
+		super()
+		@name_table = [ ]
+		@v = Rex::Struct2::CStruct_Values.new(self)
+
+		return self.add_from_dt(*dts)
+	end
+
+	def add_from_dt(*dts)
+		dts.each { | dt |
+			return if !dt.kind_of?(Array) || dt.length < 2
+
+			type = dt[0]
+			name = dt[1]
+
+			factory = @@dt_table[type]
+
+			return if !factory
+
+			# call with the arguments passed in
+			obj = factory.call(*(dt[2 .. -1]))
+
+			self.add_object(name, obj)
+		}
+
+		return dts.length
+	end
+
+	def add_object(*objs)
+		while objs.length >= 2
+			@name_table << objs.shift
+			self        << objs.shift
+		end
+	end
+	# apply_restraint( name, restraint, name2, restraint2 ... )
+	def apply_restraint(*ress)
+		while ress.length >= 2
+			name = ress.shift
+			res  = ress.shift
+			self[name].restraint = res
+
+			# update the restrainted object, so it will update the value
+			# of the restrainter, with the initial size.  If you don't
+			# want this behavior, um, you'll have to be careful with what
+			# you supply as default values...
+			self[name].update_restraint
+		end
+		return self
+	end
+
+	# create_restraints( [ name, stuff_to_restraint_constructor ] ... )
+	def create_restraints(*ress)
+		ress.each { |r|
+			# make a copy before we modify...
+			r = r.dup
+			# resolve names into objects
+			r[1] = self[r[1]] if r[1]
+			r[2] = self[r[2]] if r[2]
+
+			# build and apply the restraint
+			self.apply_restraint(r[0], Rex::Struct2::Restraint.new(*r[1 .. -1]))
+		}
+
+		return self
+	end
+
+	# ya ya, I know, these are weird.  I'm not sure why I even bothered
+	# to inherit from array...
+	def [](index, *other)
+		if index.kind_of?(String)
+			i = @name_table.index(index)
+			return if !i
+			return super(i)
+		else
+			return super(index, *other)
+		end
+	end
+
+	def []=(index, *other)
+		if index.kind_of?(String)
+			i = @name_table.index(index)
+			return if !i
+			return super(i, *other)
+		else
+			return super(index, *other)
+		end
+	end
+	
+	# Produce a list of field names
+	def keys
+		@name_table
+	end
+	
+	# Iterate through all fields and values
+	def each_pair(&block)
+		@name_table.each do |k|
+			block.call(k, self.v[k])
+		end
+	end
+end
+
+# end Rex::Struct2
+end
+end

data/lib/rex/struct2/c_struct_template.rb

@@ -0,0 +1,39 @@
+#!/usr/bin/env ruby
+
+# Rex::Struct2
+module Rex
+module Struct2
+
+class CStructTemplate
+
+	require 'rex/struct2/c_struct'
+
+	attr_reader  :template, :template_create_restraints, :template_apply_restraint
+	attr_writer  :template, :template_create_restraints, :template_apply_restraint
+
+	def initialize(*tem)
+		self.template = tem
+		self.template_create_restraints = [ ]
+		self.template_apply_restraint = [ ]
+	end
+
+	def create_restraints(*ress)
+		self.template_create_restraints = ress
+		return self
+	end
+
+	def apply_restraint(*ress)
+		self.template_apply_restraint = ress
+		return self
+	end
+
+	def make_struct
+		Rex::Struct2::CStruct.new(*self.template).
+		  create_restraints(*self.template_create_restraints).
+		  apply_restraint(*self.template_apply_restraint)
+	end
+end
+
+# end Rex::Struct2
+end
+end

data/lib/rex/struct2/constant.rb

@@ -0,0 +1,26 @@
+#!/usr/bin/env ruby
+
+# Rex::Struct2
+module Rex
+module Struct2
+
+# this is a "constant" element.  It's not actually constant, you can set it
+# via the constructor and value.  It doesn't do from_s/to_s, etc.
+
+# what use is it? Well it's useful for doing constant restraints (like fix
+# sized arrays), and probably not a ton more.
+
+class Constant
+	
+	require 'rex/struct2/element'
+	include Rex::Struct2::Element
+
+	def initialize(value)
+		self.value = value
+	end
+
+end
+
+# end Rex::Struct2
+end
+end

data/lib/rex/struct2/element.rb

@@ -0,0 +1,44 @@
+#!/usr/bin/env ruby
+
+# Rex::Struct2
+module Rex
+module Struct2
+
+module Element
+
+	# elements should have to_s, but we don't define it here because
+	# it will just overlap with inheritence and cause issues
+
+	attr_reader  :value, :restraint, :container
+	attr_writer  :restraint, :container
+
+	# update the restraints on any value change
+	def value=(newval)
+		@value = newval
+		self.update_restraint
+	end
+
+	# avoid conflicting with normal namespace length()
+	def slength
+		to_s().length()
+	end
+
+	def update_restraint
+		if self.restraint
+			# Sort of a hack, but remove the restraint before we update, so we aren't using
+			# the old restraint during calculating the restraint update value
+			old_restraint, self.restraint = self.restraint, nil
+			old_restraint.update(self.slength)
+			self.restraint = old_restraint
+		end
+
+		if self.container
+			self.container.update_restraint
+		end
+	end
+
+end
+
+# end Rex::Struct2
+end
+end

data/lib/rex/struct2/generic.rb

@@ -0,0 +1,73 @@
+#!/usr/bin/env ruby
+
+# Rex::Struct2
+module Rex
+module Struct2
+
+class Generic
+	
+	require 'rex/struct2/element'
+	include Rex::Struct2::Element
+
+	attr_reader  :default
+	attr_writer  :default
+
+	attr_accessor :mask, :check_mask
+
+	def initialize(packspec, signed=false, default=nil)
+		@packspec = packspec
+		@default  = default
+
+		bytelen = [ -1 ].pack(@packspec).length
+		self.mask = (1 << (8 * bytelen)) - 1
+
+		if signed
+			self.check_mask = 1 << (8 * bytelen - 1)
+		else
+			self.check_mask = 0
+		end
+
+		reset()
+	end
+
+	def reset
+		self.value = @default
+	end
+
+	def to_s
+		# I realize this will bomb out if this isn't an integer, for
+		# example if it is nil.  That should only happen for a user
+		# error so that's what I want it to do...
+		string = [ @value ].pack(@packspec)
+		
+		if restraint && restraint.max
+			return string.slice(0, restraint.max)
+		else
+			return string
+		end
+		# what to do for restraint.min?!?
+	end
+
+	def from_s(bytes)
+		value = bytes.unpack(@packspec)[0]
+		# return nil on unpack error
+		return if !value
+		len = slength()
+		# error on any restraint issues
+		return if restraint && restraint.max && len > restraint.max
+		return if restraint && restraint.min && len < restraint.min
+		# else set our value and return length used for this element
+
+		if (value & check_mask) != 0
+			value = -((~value & mask) + 1)
+		end
+
+		self.value = value
+		return(len)
+	end
+
+end
+
+# end Rex::Struct2
+end
+end

data/lib/rex/struct2/restraint.rb

@@ -0,0 +1,54 @@
+#!/usr/bin/env ruby
+
+# Rex::Struct2
+module Rex
+module Struct2
+
+class Restraint
+
+	attr_reader  :max_object, :min_object, :should_update,
+	  :max_transform, :min_transform, :max_inv_transform, :min_inv_transform
+	attr_writer  :max_object, :min_object, :should_update,
+	  :max_transform, :min_transform, :max_inv_transform, :min_inv_transform
+
+
+	def initialize(
+	  max_object=nil, min_object=nil, should_update=false,
+	  max_transform=nil, min_transform=nil,
+	  max_inv_transform=nil, min_inv_transform=nil
+	)
+		@max_object    = max_object
+		@min_object    = min_object
+		@should_update = should_update
+
+		def_trans = proc {|i| i}
+
+		@max_transform       = max_transform == nil ? def_trans : max_transform
+		@min_transform       = min_transform == nil ? def_trans : min_transform
+		@max_inv_transform   = max_inv_transform == nil ? def_trans : max_inv_transform
+		@min_inv_transform   = min_inv_transform == nil ? def_trans : min_inv_transform
+	end
+
+	def min
+		return if !min_object
+		return min_object.value
+	end
+
+	def max
+		return if !max_object
+		return max_object.value
+	end
+
+	# update values if request (ie string set field to its length)
+	def update(value)
+		return if !@should_update
+
+		max_object.value = max_inv_transform.call(value) if max_object
+		min_object.value = min_inv_transform.call(value) if min_object
+	end
+
+end
+
+# end Rex::Struct2
+end
+end