packetfu 1.1.2 → 1.1.3

data/lib/packetfu/protos/udp.rb

@@ -0,0 +1,240 @@
+module PacketFu
+
+	# UDPHeader is a complete UDP struct, used in UDPPacket. Many Internet-critical protocols
+	# rely on UDP, such as DNS and World of Warcraft.
+	#
+	# For more on UDP packets, see http://www.networksorcery.com/enp/protocol/udp.htm
+	#
+	# ==== Header Definition
+	#  Int16   :udp_src
+	#  Int16   :udp_dst
+	#  Int16   :udp_len  Default: calculated
+	#  Int16   :udp_sum  Default: 0. Often calculated. 
+	#  String  :body
+	class UDPHeader < Struct.new(:udp_src, :udp_dst, :udp_len, :udp_sum, :body)
+
+		include StructFu
+
+		def initialize(args={})
+			super(
+				Int16.new(args[:udp_src]),
+				Int16.new(args[:udp_dst]),
+				Int16.new(args[:udp_len] || udp_calc_len),
+				Int16.new(args[:udp_sum]),
+				StructFu::String.new.read(args[:body])
+			)
+		end
+
+		# Returns the object in string form.
+		def to_s
+			self.to_a.map {|x| x.to_s}.join
+		end
+
+		# Reads a string to populate the object.
+		def read(str)
+			force_binary(str)
+			return self if str.nil?
+			self[:udp_src].read(str[0,2])
+			self[:udp_dst].read(str[2,2])
+			self[:udp_len].read(str[4,2])
+			self[:udp_sum].read(str[6,2])
+			self[:body].read(str[8,str.size])
+			self
+		end
+
+		# Setter for the UDP source port.
+		def udp_src=(i); typecast i; end
+		# Getter for the UDP source port.
+		def udp_src; self[:udp_src].to_i; end
+		# Setter for the UDP destination port.
+		def udp_dst=(i); typecast i; end
+		# Getter for the UDP destination port.
+		def udp_dst; self[:udp_dst].to_i; end
+		# Setter for the length field. Usually should be recalc()'ed instead.
+		def udp_len=(i); typecast i; end
+		# Getter for the length field.
+		def udp_len; self[:udp_len].to_i; end
+		# Setter for the checksum. Usually should be recalc()'ed instad.
+		def udp_sum=(i); typecast i; end
+		# Getter for the checksum.
+		def udp_sum; self[:udp_sum].to_i; end
+
+		# Returns the true length of the UDP packet.
+		def udp_calc_len
+			body.to_s.size + 8
+		end
+
+		# Recalculates calculated fields for UDP.
+		def udp_recalc(args=:all)
+			arg = arg.intern if arg.respond_to? :intern
+			case args
+			when :udp_len
+				self.udp_len = udp_calc_len
+			when :all
+				self.udp_recalc(:udp_len)
+			else
+				raise ArgumentError, "No such field `#{arg}'"
+			end
+		end
+
+		# Equivalent to udp_src.to_i
+		def udp_sport
+			self.udp_src
+		end
+
+		# Equivalent to udp_src=
+		def udp_sport=(arg)
+			self.udp_src=(arg)
+		end
+
+		# Equivalent to udp_dst
+		def udp_dport
+			self.udp_dst
+		end
+		
+		# Equivalent to udp_dst=
+		def udp_dport=(arg)
+			self.udp_dst=(arg)
+		end
+
+		# Readability aliases
+
+		def udp_sum_readable
+			"0x%04x" % udp_sum
+		end
+
+	end
+
+	# UDPPacket is used to construct UDP Packets. They contain an EthHeader, an IPHeader, and a UDPHeader.
+	#
+	# == Example
+	#
+	#   udp_pkt = PacketFu::UDPPacket.new
+	#   udp_pkt.udp_src=rand(0xffff-1024) + 1024
+	#   udp_pkt.udp_dst=53
+	# 
+	#   udp_pkt.ip_saddr="1.2.3.4"
+	#   udp_pkt.ip_daddr="10.20.30.40"
+	#
+	#   udp_pkt.recalc
+	#   udp_pkt.to_f('/tmp/udp.pcap')
+	#
+	# == Parameters
+	#
+	#  :eth
+	#    A pre-generated EthHeader object.
+	#  :ip
+	#    A pre-generated IPHeader object.
+	#  :flavor
+	#    TODO: Sets the "flavor" of the UDP packet. UDP packets don't tend have a lot of
+	#    flavor, but their underlying ip headers do.
+	#  :config
+	#   A hash of return address details, often the output of Utils.whoami?
+	class UDPPacket < Packet
+
+		attr_accessor :eth_header, :ip_header, :udp_header
+
+		def self.can_parse?(str)
+			return false unless str.size >= 28
+			return false unless EthPacket.can_parse? str
+			return false unless IPPacket.can_parse? str
+			return false unless str[23,1] == "\x11"
+			return true
+		end
+
+		def read(str=nil, args={})
+			raise "Cannot parse `#{str}'" unless self.class.can_parse?(str)
+			@eth_header.read(str)
+			@ip_header.read(str[14,str.size])
+			@eth_header.body = @ip_header
+			if args[:strip]
+				udp_len = str[16,2].unpack("n")[0] - 20
+				@udp_header.read(str[14+(@ip_header.ip_hlen),udp_len])
+			else
+				@udp_header.read(str[14+(@ip_header.ip_hlen),str.size])
+			end
+			@ip_header.body = @udp_header
+			super(args)
+			self
+		end
+
+		def initialize(args={})
+			@eth_header = EthHeader.new(args).read(args[:eth])
+			@ip_header = IPHeader.new(args).read(args[:ip])
+			@ip_header.ip_proto=0x11
+			@udp_header = UDPHeader.new(args).read(args[:icmp])
+			@ip_header.body = @udp_header
+			@eth_header.body = @ip_header
+			@headers = [@eth_header, @ip_header, @udp_header]
+			super
+			udp_calc_sum
+		end
+
+		# udp_calc_sum() computes the UDP checksum, and is called upon intialization. 
+		# It usually should be called just prior to dropping packets to a file or on the wire. 
+		def udp_calc_sum
+			# This is /not/ delegated down to @udp_header since we need info
+			# from the IP header, too.
+			checksum = (ip_src.to_i >> 16)
+			checksum += (ip_src.to_i & 0xffff)
+			checksum += (ip_dst.to_i >> 16)
+			checksum += (ip_dst.to_i & 0xffff)
+			checksum += 0x11
+			checksum += udp_len.to_i
+			checksum += udp_src.to_i
+			checksum += udp_dst.to_i
+			checksum += udp_len.to_i
+			if udp_len.to_i >= 8
+				# For IP trailers. This isn't very reliable. :/
+				real_udp_payload = payload.to_s[0,(udp_len.to_i-8)] 
+			else
+				# I'm not going to mess with this right now.
+				real_udp_payload = payload 
+			end
+			chk_payload = (real_udp_payload.size % 2 == 0 ? real_udp_payload : real_udp_payload + "\x00")
+			chk_payload.unpack("n*").each {|x| checksum = checksum+x}
+			checksum = checksum % 0xffff
+			checksum = 0xffff - checksum
+			checksum == 0 ? 0xffff : checksum
+			@udp_header.udp_sum = checksum
+		end
+
+		# udp_recalc() recalculates various fields of the UDP packet. Valid arguments are:
+		#
+		#   :all
+		#     Recomputes all calculated fields.
+		#   :udp_sum
+		#     Recomputes the UDP checksum.
+		#   :udp_len
+		#     Recomputes the UDP length.
+		def udp_recalc(args=:all)
+			case args
+			when :udp_len
+				@udp_header.udp_recalc
+			when :udp_sum
+				udp_calc_sum
+			when :all
+				@udp_header.udp_recalc
+				udp_calc_sum
+			else
+				raise ArgumentError, "No such field `#{arg}'"
+			end
+		end
+
+		# Peek provides summary data on packet contents.
+		def peek_format
+			peek_data = ["U  "]
+			peek_data << "%-5d" % self.to_s.size
+			peek_data << "%-21s" % "#{self.ip_saddr}:#{self.udp_sport}"
+			peek_data << "->"
+			peek_data << "%21s" % "#{self.ip_daddr}:#{self.udp_dport}"
+			peek_data << "%23s" % "I:"
+			peek_data << "%04x" % self.ip_id
+			peek_data.join
+		end
+
+	end
+
+end
+
+# vim: nowrap sw=2 sts=0 ts=2 ff=unix ft=ruby

data/lib/packetfu/structfu.rb

@@ -0,0 +1,294 @@
+# 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
+
+	# Handle deep copies correctly. Marshal in 1.9, re-read myself on 1.8
+	def clone
+		begin
+			Marshal.load(Marshal.dump(self))
+		rescue
+			self.class.new.read(self.to_s)
+		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
+			@packstr = (self.e == :big) ? "n" : "v"
+			[(self.v || self.d)].pack(@packstr)
+	 	end
+
+	end
+  
+	# Int16be is a two byte value in big-endian format. The endianness cannot be altered.
+	class Int16be < Int16
+		undef :endian=
+	end
+
+	# Int16le is a two byte value in little-endian format. The endianness cannot be altered.
+	class Int16le < Int16
+		undef :endian=
+		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
+			@packstr = (self.e == :big) ? "N" : "V"
+			[(self.v || self.d)].pack(@packstr)
+	 	end
+
+	end
+
+	# Int32be is a four byte value in big-endian format. The endianness cannot be altered.
+	class Int32be < Int32
+		undef :endian=
+	end
+
+	# Int32le is a four byte value in little-endian format. The endianness cannot be altered.
+	class Int32le < Int32
+		undef :endian=
+		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)
+			if int < Int
+				super(int.new,string,mode)
+				calc
+			else
+				raise "IntStrings need a StructFu::Int for a length."
+			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)
+		PacketFu.force_binary(str)
+	end
+
+end
+
+# vim: nowrap sw=2 sts=0 ts=2 ff=unix ft=ruby