# -*- coding: binary -*-
##
# $Id$
# Version: $Revision$
##

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

module Rex
module OLE

require 'rex/ole/direntry'

#
# This class serves as the root directory entry in addition to
# an abstraction around the concept of a directory as a whole.
#
class Directory < DirEntry

	# XXX: num_entries is not maintained once a stream/storage is added!
	attr_accessor :num_entries

	def initialize(stg)
		super

		@num_entries = 1
	end


	# woop, recursive each
	def yield_entries(de, &block)
		block.call(de)
		de.each { |el|
			yield_entries(el, &block)
		}
	end
	def each_entry(&block)
		yield_entries(self, &block)
	end


	def set_ministream_params(start, size)
		@_sectStart = start
		@_ulSize = size
	end

	def link_item(parent, child)
		# set sid, advance count
		child.sid = @num_entries
		@num_entries += 1

		# link item to siblings and/or parent
		if (parent._sidChild == DIR_NOSTREAM)
			parent._sidChild = child.sid
			dlog("Linking #{child.name} as THE child of #{parent.name} as sid #{child.sid}", 'rex', LEV_3)
		else
			sib = nil
			parent.each { |el|
				if (el._sidLeftSib == DIR_NOSTREAM)
					sib = el
					el._sidLeftSib = child.sid
					dlog("Linking #{child.name} as the LEFT sibling of #{sib.name} as sid #{child.sid}", 'rex', LEV_3)
					break
				end
				if (el._sidRightSib == DIR_NOSTREAM)
					sib = el
					el._sidRightSib = child.sid
					dlog("Linking #{child.name} as the RIGHT sibling of #{sib.name} as sid #{child.sid}", 'rex', LEV_3)
					break
				end
			}
			if (not sib)
				raise RuntimeError, 'Unable to find a sibling to link to in the directory'
			end
		end
		parent << child
	end


	#
	# low-level functions
	#
	def from_s(sid, buf)
		super

		if (@_sidRightSib != DIR_NOSTREAM)
			raise RuntimeError, 'Root Entry is invalid! (has right sibling)'
		end
		if (@_sidLeftSib != DIR_NOSTREAM)
			raise RuntimeError, 'Root Entry is invalid! (has left sibling)'
		end
	end

	def read
		@children = []
		visited = []
		entries = []
		root_node = nil
		sect = @stg.header._sectDirStart
		while (sect != SECT_END)

			if (visited.include?(sect))
				raise RuntimeError, 'Sector chain loop detected (0x%08x)' % sect
			end
			visited << sect

			sbuf = @stg.read_sector(sect, @stg.header.sector_size)
			while (sbuf.length >= DIRENTRY_SZ)
				debuf = sbuf.slice!(0, DIRENTRY_SZ)

				type = Util.get8(debuf, 0x42)
				case type
				when STGTY_ROOT
					if (entries.length != 0)
						raise RuntimeError, 'Root Entry found, but not first encountered!'
					end
					if (root_node)
						raise RuntimeError, 'Multiple root directory sectors detected (0x%08x)' % sect
					end
					de = self
					root_node = de

				when STGTY_STORAGE
					de = SubStorage.new @stg

				when STGTY_STREAM
					de = Stream.new @stg

				when STGTY_INVALID
					# skip invalid entries
					next

				else
					raise RuntimeError, 'Unsupported directory entry type (0x%02x)' % type
				end

				# read content
				de.from_s(entries.length, debuf)
				entries << de
			end
			sect = @stg.next_sector(sect)
		end

		@num_entries = entries.length

		# sort out the tree structure, starting with the root
		if (@_sidChild != DIR_NOSTREAM)
			populate_children(entries, root_node, @_sidChild)
		end
	end


	# recursively add entries to their proper parents :)
	def populate_children(entries, parent, sid)
		node = entries[sid]
		dlog("populate_children(entries, \"#{parent.name}\", #{sid}) - node: #{node.name}", 'rex', LEV_3)
		parent << node
		if (node.type == STGTY_STORAGE) and (node._sidChild != DIR_NOSTREAM)
			populate_children(entries, node, node._sidChild)
		end
		if (node._sidLeftSib != DIR_NOSTREAM)
			populate_children(entries, parent, node._sidLeftSib)
		end
		if (node._sidRightSib != DIR_NOSTREAM)
			populate_children(entries, parent, node._sidRightSib)
		end
	end

	# NOTE: this may not be necessary if we were to use each_entry
	def flatten_tree(entries, parent)
		entries << parent
		parent.each { |el|
			flatten_tree(entries, el)
		}
	end


	def write
		# flatten the directory again
		entries = []
		flatten_tree(entries, self)
		dlog("flattened tree has #{entries.length} entries...", 'rex', LEV_3)

		# count directory sectors
		ds_count = entries.length / 4
		if ((entries.length % 4) > 0)
			# one more sector to hold the rest
			ds_count += 1
		end

		# put the root entry first
		sbuf = self.pack

		# add the rest
		prev_sect = nil
		dir_start = nil
		entries.each { |de|
			# we already got the root entry, no more!
			next if (de.type == STGTY_ROOT)

			dir = de.pack
			dlog("writing dir entry #{de.name}", 'rex', LEV_3)
			sbuf << dir

			if (sbuf.length == @stg.header.sector_size)
				# we have a full sector, add it!
				sect = @stg.write_sector(sbuf, nil, prev_sect)
				prev_sect = sect
				dir_start ||= sect
				# reset..
				sbuf = ""
			end
		}

		# still a partial sector left?
		if (sbuf.length > 0)
			# add it! (NOTE: it will get padded with nul bytes if its not sector sized)
			sect = @stg.write_sector(sbuf, nil, prev_sect)
			prev_sect = sect
			dir_start ||= sect
		end

		@stg.header._sectDirStart = dir_start
	end

end

end
end