##
# $Id$
# Version: $Revision$
##

##
# Rex::OLE - an OLE implementation
# written in 2010 by Joshua J. Drake <jduck [at] metasploit.com>
##

module Rex
module OLE

class DIFAT

	def initialize stg
		@stg = stg
		@entries = []
	end

	#
	# convenience access to entries
	#
	def []=(idx,expr)
		@entries[idx] = expr
	end

	def [](idx)
		@entries[idx]
	end

	def +(expr)
		@entries += expr
		self
	end

	def <<(expr)
		@entries << expr
	end

	def length
		@entries.length
	end

	def slice!(start,stop)
		@entries.slice!(start,stop)
	end

	def reset
		@entries = []
	end

	def each
		@entries.each { |el|
			yield el
		}
	end

	#
	# woop
	#
	def to_s
		ret = "{ "
		@entries.each { |el|
			ret << ", " if (ret.length > 2)
			case el
			when SECT_END
				ret << "END"
			when SECT_DIF
				ret << "DIF"
			when SECT_FAT
				ret << "FAT"
			when SECT_FREE
				ret << "FREE"
			else
				ret << "0x%x" % el
			end
		}
		ret << " }"
		ret
	end

	#
	# low-level functions
	#
	def read
		@entries = []

		# start with the header part
		@entries += @stg.header._sectFat

		# double indirect fat
		sect = @stg.header._sectDifStart
		while (sect != SECT_END)
			if (@entries.include?(sect))
				raise RuntimeError, 'Sector chain loop detected (0x%08x)' % sect
			end

			@entries << sect
			buf = @stg.read_sector(sect, @stg.header.sector_size)

			# the last sect ptr in the block becomes the next entry
			sect = Util.get32(buf, ((@stg.header.idx_per_sect)-1) * 4)
		end

		# don't need these free ones, but it doesn't hurt to keep them.
		#@difat.delete(SECT_FREE)
	end

	def write
		len = @entries.length
		first109 = @entries.dup

		rest = nil
		if (len > 109)
			rest = first109.slice!(109,len)
		end

		@stg.header._sectFat = []
		@stg.header._sectFat += first109
		if (len < 109)
			need = 109 - len
			need.times {
				@stg.header._sectFat << SECT_FREE
			}
		end

		if (rest and rest.length > 0)
			raise RuntimeError, 'TODO: support writing DIF properly!'
			# may require adding more fat sectors :-/
			#@stg.header._csectDif = rest.length
			#@stg.header._sectDifStart = idx
		end

		@stg.header._csectFat = len
	end

end

end
end