metasploit-framework/lib/packetfu/structfu.rb

275 lines
7.0 KiB
Ruby

# StructFu, a nifty way to leverage Ruby's built in Struct class
# to create meaningful binary data.
module StructFu
# Normally, self.size and self.length will refer to the Struct
# size as an array. It's a hassle to redefine, so this introduces some
# shorthand to get at the size of the resultant string.
def sz
self.to_s.size
end
alias len sz
# Typecast is used mostly by packet header classes, such as IPHeader,
# TCPHeader, and the like. It takes an argument, and casts it to the
# expected type for that element.
def typecast(i)
c = caller[0].match(/.*`([^']+)='/)[1]
self[c.intern].read i
end
# Used like typecast(), but specifically for casting Strings to StructFu::Strings.
def body=(i)
if i.kind_of? ::String
typecast(i)
elsif i.kind_of? StructFu
self[:body] = i
elsif i.nil?
self[:body] = StructFu::String.new.read("")
else
raise ArgumentError, "Can't cram a #{i.class} into a StructFu :body"
end
end
# Ints all have a value, an endianness, and a default value.
# Note that the signedness of Int values are implicit as
# far as the subclasses are concerned; to_i and to_f will
# return Integer/Float versions of the input value, instead
# of attempting to unpack the pack value. (This can be a useful
# hint to other functions).
#
# ==== Header Definition
#
# Fixnum :value
# Symbol :endian
# Fixnum :width
# Fixnum :default
class Int < Struct.new(:value, :endian, :width, :default)
alias :v= :value=
alias :v :value
alias :e= :endian=
alias :e :endian
alias :w= :width=
alias :w :width
alias :d= :default=
alias :d :default
# This is a parent class definition and should not be used directly.
def to_s
raise StandardError, "StructFu::Int#to_s accessed, must be redefined."
end
# Returns the Int as an Integer.
def to_i
(self.v || self.d).to_i
end
# Returns the Int as a Float.
def to_f
(self.v || self.d).to_f
end
def initialize(value=nil, endian=nil, width=nil, default=nil)
super(value,endian,width,default=0)
end
# Reads either an Integer or a packed string, and populates the value accordingly.
def read(i)
self.v = i.kind_of?(Integer) ? i.to_i : i.to_s.unpack(@packstr).first
self
end
end
# Int8 is a one byte value.
class Int8 < Int
def initialize(v=nil)
super(v,nil,w=1)
@packstr = "C"
end
# Returns a one byte value as a packed string.
def to_s
[(self.v || self.d)].pack("C")
end
end
# Int16 is a two byte value.
class Int16 < Int
def initialize(v=nil, e=:big)
super(v,e,w=2)
@packstr = (self.e == :big) ? "n" : "v"
end
# Returns a two byte value as a packed string.
def to_s
[(self.v || self.d)].pack(@packstr)
end
end
# Int16be is a two byte value in big-endian format.
class Int16be < Int16
end
# Int16le is a two byte value in little-endian format.
class Int16le < Int16
def initialize(v=nil, e=:little)
super(v,e)
@packstr = (self.e == :big) ? "n" : "v"
end
end
# Int32 is a four byte value.
class Int32 < Int
def initialize(v=nil, e=:big)
super(v,e,w=4)
@packstr = (self.e == :big) ? "N" : "V"
end
# Returns a four byte value as a packed string.
def to_s
[(self.v || self.d)].pack(@packstr)
end
end
# Int32be is a four byte value in big-endian format.
class Int32be < Int32
end
# Int32le is a four byte value in little-endian format.
class Int32le < Int32
def initialize(v=nil, e=:little)
super(v,e)
end
end
# Strings are just like regular strings, except it comes with a read() function
# so that it behaves like other StructFu elements.
class String < ::String
def read(str)
str = str.to_s
self.replace str
self
end
end
# Provides a primitive for creating strings, preceeded by
# an Int type of length. By default, a string of length zero with
# a one-byte length is presumed.
#
# Note that IntStrings aren't used for much, but it seemed like a good idea at the time.
class IntString < Struct.new(:int, :string, :mode)
def initialize(string='',int=Int8,mode=nil)
unless int.respond_to?(:ancestors) && int.ancestors.include?(StructFu::Int)
raise StandardError, "Invalid length (#{int.inspect}) associated with this String."
else
super(int.new,string,mode)
calc
end
end
# Calculates the size of a string, and sets it as the value.
def calc
int.v = string.to_s.size
self.to_s
end
# Returns the object as a string, depending on the mode set upon object creation.
def to_s
if mode == :parse
"#{int}" + [string].pack("a#{len}")
elsif mode == :fix
self.int.v = string.size
"#{int}#{string}"
else
"#{int}#{string}"
end
end
# By redefining #string=, we can ensure the correct value
# is calculated upon assignment. If you'd prefer to have
# an incorrect value, use the syntax, obj[:string]="value"
# instead. Note, by using the alternate form, you must
# #calc before you can trust the int's value. Think of the
# = assignment as "set to equal," while the []= assignment
# as "boxing in" the value. Maybe.
def string=(s)
self[:string] = s
calc
end
# Shorthand for querying a length. Note that the usual "length"
# and "size" refer to the number of elements of this struct.
def len
self[:int].value
end
# Override the size, if you must.
def len=(i)
self[:int].value=i
end
# Read takes a string, assumes an int width as previously
# defined upon initialization, but makes no guarantees
# the int value isn't lying. You're on your own to test
# for that (or use parse() with a :mode set).
def read(s)
unless s[0,int.width].size == int.width
raise StandardError, "String is too short for type #{int.class}"
else
int.read(s[0,int.width])
self[:string] = s[int.width,s.size]
end
self.to_s
end
# parse() is like read(), except that it interprets the string, either
# based on the declared length, or the actual length. Which strategy
# is used is dependant on which :mode is set (with self.mode).
#
# :parse : Read the length, and then read in that many bytes of the string.
# The string may be truncated or padded out with nulls, as dictated by the value.
# :fix : Skip the length, read the rest of the string, then set the length to what it ought to be.
# else : If neither of these modes are set, just perfom a normal read(). This is the default.
def parse(s)
unless s[0,int.width].size == int.width
raise StandardError, "String is too short for type #{int.class}"
else
case mode
when :parse
int.read(s[0,int.width])
self[:string] = s[int.width,int.value]
if string.size < int.value
self[:string] += ("\x00" * (int.value - self[:string].size))
end
when :fix
self.string = s[int.width,s.size]
else
return read(s)
end
end
self.to_s
end
end
end
class Struct
# Monkeypatch for Struct to include some string safety -- anything that uses
# Struct is going to presume binary strings anyway.
def force_binary(str)
str.force_encoding "binary" if str.respond_to? :force_encoding
end
end
# vim: nowrap sw=2 sts=0 ts=2 ff=unix ft=ruby