dstruct 0.0.1

data/lib/rex/struct2.rb

@@ -0,0 +1,5 @@
+# -*- coding: binary -*-
+
+# 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 @@
+# -*- coding: binary -*-
+
+# 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 @@
+# -*- coding: binary -*-
+
+# 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 @@
+# -*- coding: binary -*-
+
+# 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 @@
+# -*- coding: binary -*-
+
+# 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 @@
+# -*- coding: binary -*-
+
+# 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 @@
+# -*- coding: binary -*-
+
+# 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

data/lib/rex/struct2/s_string.rb

@@ -0,0 +1,72 @@
+# -*- coding: binary -*-
+
+# Rex::Struct2
+module Rex
+module Struct2
+
+class SString
+
+  require 'rex/struct2/element'
+  require 'rex/struct2/constant'
+  include Rex::Struct2::Element
+
+  attr_reader  :size, :default, :pad
+  attr_writer  :default, :pad
+
+  def initialize(size=nil, default=nil, pad=nil)
+    self.size = size
+    @default  = default
+    @pad      = pad
+    reset()
+  end
+
+  def size=(newsize)
+    if !newsize
+      self.restraint = nil
+    else
+      res = Rex::Struct2::Constant.new(newsize)
+      self.restraint = Rex::Struct2::Restraint.new(res, res, false)
+    end
+  end
+
+  def reset
+    self.value = @default
+  end
+
+  def to_s
+    string = self.value
+
+    return if !string
+
+    # pad if short
+    if restraint && restraint.min && self.pad && restraint.min > string.length
+      string += self.pad * (restraint.min - string.length)
+    end
+    # truncate if long
+    if restraint && restraint.max
+      string = string.slice(0, restraint.max)
+    end
+
+    return string
+  end
+
+  def from_s(bytes)
+    # we don't have enough bytes to satisfy our minimum
+    if restraint && restraint.min && bytes.length < restraint.min
+      return
+    end
+
+    if restraint && restraint.max
+      self.value = bytes.slice(0, restraint.max)
+    else
+      self.value = bytes.dup
+    end
+
+
+    return(self.slength)
+  end
+end
+
+# end Rex::Struct2
+end
+end