#!/usr/bin/ruby

#
# this is a ruby library to manipulate serialized java objects
# should handle almost all objects
#
# intended use: retrieve an object, use this lib to patch it, then replace the original (aka man in the middle)
# (full object creation from scratch is too painful)
# you can also quickly inspect the object structure using 'puts jobj'
#
# (c) 2007 Yoann Guillot
# distributes under the terms of the wtfpl (sam.zoy.org/wtfpl/)
#


class JavaObj
	Type = { 0x70 => :null, 0x71 => :reference, 0x72 => :classdesc, 0x73 => :object,
		0x74 => :string, 0x75 => :array, 0x76 => :class, 0x77 => :blockdata,
		0x78 => :endblockdata, 0x79 => :reset, 0x7a => :blockdatalong, 0x7b => :exception,
		0x7c => :longstring, 0x7d => :proxyclassdesc, 0x7e => :enum }

	Typecode = { ?B => :byte, ?C => :char, ?D => :double, ?F => :float, ?I => :int, ?J => :long,
		?S => :short, ?Z => :bool, ?[ => :array, ?L => :object }

	attr_accessor :bin, :ptr, :handle2obj, :objects

	def initialize
	end

	def [](i) objects.first[i] end
	def []=(i, v) objects.first[i] = v end
	def to_s ; objects.to_s end

	def read(len)
		@ptr += len
		@bin[@ptr-len, len]
	end
	def readbyte
		read(1)[0]
	end
	def readshort
		read(2).unpack('n').first
	end
	def readint
		read(4).unpack('N').first
	end
	def readlong
		(readint << 32) | readint
	end

	def exception(msg)
		RuntimeError.new "#{msg}, found #{Type[@bin[@ptr-1]]} #{@bin[@ptr-1]} #{'0x%x' % @bin[@ptr-1]} at #{@ptr-1} #{'0x%x' % (@ptr-1)}"
	end

	# structure parser
	
	def readcontent
		b = @bin[@ptr]
		case Type[b]
		when :blockdata, :blockdatalong; readblock
		when nil
		else readobject
		end
	end

	def readobject
puts "reading #{Type[@bin[@ptr]]} at #{@ptr} #{'0x%x' % @ptr}" if $DEBUG
		case Type[@bin[@ptr]]
		when :object; Object.new.read(self)
		when :class; readclass
		when :array; Array.new.read(self)
		when :string, :longstring; String.new.read(self)
		when :enum; Enum.new.read(self)
		when :classdesc, :proxyclassdesc; readclassdesc
		when :null; readbyte ; nil
		when :reference; readprevobj
		when :exception; readexception
		when :reset; :reset
		else raise self, 'object expected'
		end
	end

	def readprevobj
		raise self, 'prevobj expected' if Type[readbyte] != :reference

		ref = readint
		if not o = @handle2obj[ref]
			# raise
			puts "invalid prev handle #{'%x' % ref} at #{@ptr-4} #{'%x' % (@ptr-4)}" 
			return "prevobj #{'%x' % ref}"
		end
		o
	end

	def readclass
		case Type[@bin[@ptr]]
		when :class; Class.new.read(self)
		when :null; nil
		when :reference; readprevobj
		else readbyte ; raise self, 'classlike expected'
		end
	end

	class Class
		attr_accessor :desc

		def read(jr)
			raise jr, 'class expected' if Type[jr.readbyte] != :class

			@desc = jr.readclassdesc
			jr.newhandle self

			self
		end

		def write(jr)
			jr.writetype :class
			jr.writeobj @desc
			jr.newhandle self
		end
	end

	def readclassdesc
		case Type[@bin[@ptr]]
		when :classdesc; ClassDesc.new.read(self)
		when :proxyclassdesc; ProxyClassDesc.new.read(self)
		when :null; readbyte ; nil
		when :reference; readprevobj
		else readbyte ; raise self, 'classdesclike expected'
		end
	end

	class ClassDesc
		attr_accessor :classname, :serial, :flags, :fields, :annotation, :superclassdesc

		def each_parent(&b)
			if superclassdesc
				@superclassdesc.each_parent(&b)
			end
			yield self
		end
			
		def parent_fields
			sf = []
			sf = @superclassdesc.parent_fields if superclassdesc
			sf + @fields
		end

		def parent_flags
			(@superclassdesc ? @superclassdesc.parent_flags : []) | @flags
		end

		def read(jr)
			raise jr, 'classdesc expected' if Type[jr.readbyte] != :classdesc
			@classname = jr.readutf
			@serial = jr.readlong
			jr.newhandle self

			f = jr.readbyte
			@flags = []
			@flags << :write_method if f & 1 == 1
			@flags << :serializable if f & 2 == 2
			@flags << :externalizable if f & 4 == 4
			@flags << :block_data if f & 8 == 8
			@flags << :enum if f & 16 == 16

			@fields = []
			fieldcount = jr.readshort
			fieldcount.times { @fields << jr.readfielddesc }
			a = jr.readblock
			@annotation = a if a
			sc = jr.readclassdesc
			@superclassdesc = sc if sc
puts 'found classdesc ' + inspect if $DEBUG
			self
		end

		def write(jr)
			jr.writetype :classdesc
			jr.writeutf @classname
			jr.writelong @serial
			jr.newhandle self

			f = 0
			f |= 1 if @flags.include? :write_method
			f |= 2 if @flags.include? :serializable
			f |= 4 if @flags.include? :externalizable
			f |= 8 if @flags.include? :block_data
			f |= 16 if @flags.include? :enum
			jr.writebyte f

			jr.writeshort @fields.length
			@fields.each { |f| jr.writeobj(f) }
			jr.writeblock annotation
			jr.writeobj superclassdesc
		end
	end

	class ProxyClassDesc
		attr_accessor :proxies, :annotation, :superclassdesc

		def read(jr)
			raise jr, 'proxyclassdesc expected' if Type[jr.readbyte] != :proxyclassdesc

			jr.newhandle self
			@proxies = []
			jr.readint.times { @proxies << jr.readutf }
			a = jr.readblock
			@annotation = a if a
			sc = jr.readclassdesc
			@superclassdesc = sc if sc

			self
		end

		def write(jr)
			jr.writetype :proxyclassdesc
			jr.newhandle self
			jr.writeint @proxies.length
			@proxies.each { |p| jr.writeutf p }
			jr.writeblock annotation
			jr.writeobj superclassdesc
		end
	end

	class Array
		attr_accessor :classdesc, :values

		def to_s(done=[])
			done += [self]
			ret = "["
			@values.each { |v|
				if done.include? v
					str = '...'
				else
					case v
					when Object, Array; str = v.to_s(done).to_a.map { |s| "  #{s}" }.join.chomp + ',' + "\n"
					else str = v.inspect + ', '
					end
				end

				ret << "#{str}"
			}
			ret << "]"
			ret
		end

		def read(jr)
			raise jr, 'array expected' if Type[jr.readbyte] != :array

			@classdesc = jr.readclassdesc
			jr.newhandle self
			@values = []
			jr.readint.times { @values << jr.readobject }

			self
		end

		def write(jr)
			jr.writetype :array
			jr.writeobj @classdesc
			jr.newhandle self
			jr.writeint(@values.length)
			@values.each { |v| jr.writeobj v }
		end

		def each(&b) @values.each(&b) end
		def [](*a) @values[*a] end
		def []=(*a) @values[*a[0..-2]]=a[-1] end
	end

	class Object
		attr_accessor :classdesc, :values, :annotations

		def to_s(done=[])
			case @classdesc.classname
			when 'java.sql.Timestamp'
				ts = @annotations.first.first.data.reverse.unpack('q').first
				return case ts
				when -0x38831c799000; '1/1/0 0:0:0'
				when  0x1d8fd5268400; '1/1/3000 0:0:0'
				else Time.at(ts/1000).strftime('%d/%m/%Y %H:%M:%S') rescue 'date:%x' % ts
				end
			when 'java.lang.Long'; return self['value'].to_s
			when 'java.lang.Double'; return self['value'].to_s
			end
			done += [self]
			ret = "#{@classdesc.classname}\n"
			ret << "annotation #{@annotations.inspect}\n" if annotations
			@classdesc.parent_fields.zip(@values).each { |f, v|
				if done.include? v
					str = '...'
				else
					case v
					when Object, Array; str = "\n" + v.to_s(done).to_a.map { |s| "  #{s}" }.join.chomp
					else str = v.inspect
					end
				end

				ret << " (#{f.type}) #{f.name} = #{str}\n"
			}
			ret
		end

		def read(jr)
			raise jr, 'new object expected' if Type[jr.readbyte] != :object

			@classdesc = jr.readclassdesc
			jr.newhandle self
			@values = []
			# @annotations = [] if @classdesc.parent_flags.include? :write_method
			raise 'externalizable data unhandled' if @classdesc.flags.include? :externalizable

			@classdesc.each_parent { |c|
				c.fields.each { |f|
#			@classdesc.parent_fields.each { |f|
puts "fld #{f.type} #{f.name} at #{'0x%x' % jr.ptr} #{jr.bin[jr.ptr, 16].inspect}" if $DEBUG
				case f.type
				when :object; v = jr.readobject
				when :array;  v = jr.readobject
				when :int;    v = jr.readint
				when :bool;   v = jr.readbyte
				when :char;   v = jr.readshort.chr
				when :short;  v = jr.readshort
				when :long;   v = jr.readlong
				when :double; v = jr.read(8).unpack('G').first
				when :float;  v = jr.read(4).unpack('g').first
				else raise jr, "unhandled field #{f.type} #{jr.bin[jr.ptr, 16].inspect}, #{@classdesc.parent_fields.map { |p| [p.name, p.type] }.inspect}" 
				end
puts "fld #{f.type} #{f.name} found #{v.inspect}" if $DEBUG
				@values << v
				}
				(@annotations ||= []) << jr.readblock if c.flags.include? :write_method
			}

			self
		end

		def write(jr)
			jr.writetype :object

			jr.writeobj @classdesc
			jr.newhandle self
			val = @values.dup
			ann = @annotations.dup if ann
			@classdesc.each_parent { |c|
			#@classdesc.parent_fields.zip(@values).each { |f, v|
				c.fields.each { |f|
				v = val.shift
				case f.type
				when :object; jr.writeobj v
				when :array;  jr.writeobj v
				when :int;    jr.writeint v
				when :bool;   jr.writebyte v
				when :char;   jr.writeshort v[0]
				when :short;  jr.writeshort v
				when :long;   jr.writelong v
				when :double; jr.bin << [v].pack('G')
				when :float;  jr.bin << [v].pack('g')
				else raise "unhandled field #{f.type}"
				end
				}
				jr.writeblock(ann.shift) if ann and c.flags.include? :write_method
			}
		end

		def [](idx)
			pf = @classdesc.parent_fields
			idx = pf.index(pf.find { |f| f.name.downcase == idx.downcase })
			raise "unknown index #{idx} not in #{pf.map { |f| f.name }.inspect}" if not idx
			@values[idx]
		end

		def []=(idx, val)
			pf = @classdesc.parent_fields
			idx = pf.index(pf.find { |f| f.name.downcase == idx.downcase })
			raise "unknown index #{idx} not in #{pf.map { |f| f.name }.inspect}" if not idx
			@values[idx] = val
		end
	end

	class String
		attr_accessor :data

		def initialize(data=nil)
			@data = data if data
		end

		def read(jr)
			b = jr.readbyte

			case Type[b]
			when :string; len = jr.readshort
			when :longstring; len = jr.readlong
			else raise jr, 'string expected'
			end

			@data = jr.read(len)

			#jr.newhandle(self)
			#self

			jr.newhandle(@data)
			@data
		end

		def write(jr)
			if @data.length > 0xffff
				jr.writetype :longstring
				jr.writelong @data.length
			else
				jr.writetype :string
				jr.writeshort @data.length
			end
			jr.bin << @data
		end
	end

	class Enum
		attr_accessor :classdesc, :str

		def read(jr)
			raise jr, 'enum expected' if Type[jr.readbyte] != :enum

			@classdesc = jr.readclassdesc
			jr.newhandle self
			@str = jr.readobject
		end

		def write(jr)
			jr.writetype :enum
			jr.writeobj @classdesc
			jr.newhandle self
			jr.writeobj @str
		end
	end

	def readfielddesc
		FieldDesc.new.read(self)
	end

	class FieldDesc
		attr_accessor :type, :name, :objclass, :objclass

		def read(jr)
			opos = jr.ptr
			@type = Typecode[jr.readbyte]
			raise jr, 'invalid typecode' if not @type
			@name = jr.readutf
			case @type
			when :array, :object; @objclass = jr.readobject
			end
puts "new field descriptor at #{opos} #{'0x%x' % opos} #{inspect}" if $DEBUG

			self
		end

		def write(jr)
			jr.writebyte Typecode.index(@type)
			jr.writeutf @name
			case @type
			when :array, :object; jr.writeobj @objclass
			end
		end
	end

	def readblock
		b = []
		loop do
			b << 
			case Type[readbyte]
			when :blockdata; Block.new.read(self, false)
			when :blockdatalong; Block.new.read(self, true)
			when :endblockdata; break
			else @ptr -= 1 ; readobject
			end
		end
		b.empty? ? nil : b
	end

	def writeblock(b)
		b.to_a.each { |e|
			writeobj e
		}
		writetype :endblockdata
	end

	class Block
		attr_accessor :data

		def read(jr, long=false)
			len = (long ? jr.readint : jr.readbyte)
puts "block data start at #{'0x%x' % jr.ptr}, #{len}" if $DEBUG
			@data = jr.read(len)

			self
		end

		def write(jr)
			if data
				if @data.length > 255
					jr.writetype :blockdata
					jr.writeint(@data.length)
				else
					jr.writetype :blockdatalong
					jr.writebyte(@data.length)
				end
				jr.bin << @data
			end
		end
	end

	def readutf
		len = readshort
		read(len)
	end

	def newhandle(obj)
		@handlectr ||= 0x7e0000-1
		@handlectr += 1
		@handle2obj[@handlectr] = obj
		@handlectr
	end

	# reads a binary stream
	def load(bin)
		@ptr = 0
		@bin = bin
		@handle2obj = {}

		raise 'java load: no data' if bin.empty?
		raise 'invalid javaobjstream signature' if readshort != 0xaced
		raise 'invalid javaobjstream version' if readshort != 5
		
		@objects = []
		while o = readcontent
			@objects << o
		end

puts "unread bytes from stream: #{read(1024)}" if $DEBUG and @ptr < @bin.length

		self
	end

	def self.load(bin)
		new.load(bin)
	end

	# writes a stream
	
	def writebyte(b)
		@bin << b
	end
	def writeshort(v)
		@bin << [v].pack('n')
	end
	def writeint(i)
		@bin << [i].pack('N')
	end
	def writelong(l)
		writeint((l >> 32) & 0xffff_ffff)
		writeint(l & 0xffff_ffff)
	end
	def writetype(t)
		writebyte Type.index(t)
	end
	def writeutf(str)
		writeshort(str.length)
		@bin << str
	end

	def writeobj(obj)
		#if h = @handle2obj.index(obj)	# optimizes dup strings => not binary identical
		if h = @handle2obj.find { |k, v| v.object_id == obj.object_id } and h = h[0]
			writetype(:reference)
			writeint(h)
		elsif not obj
			writetype(:null)
		elsif obj.kind_of? ::String
			String.new(obj).write(self)
			newhandle(obj)
		else
			obj.write self
		end
	end

	def save
		@bin = ''
		@handlectr = nil
		@handle2obj = {}
		writeshort(0xaced)
		writeshort(5)
		@objects.each { |o|
			writeobj(o)
		}
		@bin
	end
end

if $0 == __FILE__
	begin
	require 'pp'

	if ARGV.empty?
		require 'clipboard'
		bin = Clipboard.getdata
		File.open('raw', 'wb') { |fd| fd.write bin }
	else
		bin = File.open(ARGV.shift, 'rb') { |fd| fd.read }
	end

	# bin.hexdump
	jo = JavaObj.new
	#$DEBUG = true
	jo.load(bin)

	if jo.objects.first.kind_of? JavaObj::Object
		puts jo.objects
	else
		pp jo.objects
	end

	if $DEBUG
	# check that obj.serialize == original string
	oldbin = jo.bin.dup
	jo.save

	if jo.bin != oldbin
		File.open('raw.resave', 'wb') { |fd| fd.write jo.bin }
		diffidx = 0
		diffidx += 1 while oldbin[diffidx] == jo.bin[diffidx]
		diffend = -1
		diffend -= 1 while oldbin[diffend] == jo.bin[diffend]
		puts "difference from #{diffidx} #{'0x%x' % diffidx} to #{diffend} #{'0x%x' % diffend}"
		diffidx -= 5 ; diffidx = 0 if diffidx < 0
		diffend += 5 ; diffend = -1 if diffend > -1
		p oldbin[diffidx..diffend], jo.bin[diffidx..diffend]
	else
		puts "resave identical" if $DEBUG
	end
	end

	rescue
		puts $!, $!.backtrace
	end

	if RUBY_PLATFORM =~ /mswin32/i
		puts 'press enter'
		gets
	end
end