#!/usr/bin/env ruby
#    This file is part of Metasm, the Ruby assembly manipulation suite
#    Copyright (C) 2006-2009 Yoann GUILLOT
#
#    Licence is LGPL, see LICENCE in the top-level directory


# This sample implements a trivial binary diffing algorithm between two programs
# the programs have first to be disassembled, and then the diff algorith will
# (try to) identify identical functions in both dasm graphs
# Currently there is NO fuzzy matching whatsoever, so the function graphs have to
# be exactly the same in both programs to be recognized.
# You can still force a comparaison between two functions, but the results will be bad.
#
# This file can be run as a standalone application (eg 'ruby bindiff file1 file2')
# or as a disassembler plugin (see dasm-plugin/bindiff)

require 'metasm'

module ::Metasm
class BinDiffWidget < Gui::DrawableWidget
	attr_accessor :status

	COLORS = { :same => '8f8', :similar => 'cfc', :badarg => 'fcc', :badop => 'faa', :default => '888' }

	def initialize_widget(d1=nil, d2=nil)
		self.dasm1 = d1 if d1
		self.dasm2 = d2 if d2
		@status = nil
	end

	def dasm1; @dasm1 end
	def dasm1=(d)
		@dasm1 = d
		@func1 = nil
		@funcstat1 = nil
		@dasmcol1 = {}
		@dasm1.gui.bg_color_callback = lambda { |a1| COLORS[@dasmcol1[a1] || :default] }
		@match_func = nil
	end

	def dasm2; @dasm2 end
	def dasm2=(d)
		@dasm2 = d
		@func2 = nil
		@funcstat1 = nil
		@dasmcol2 = {}
		@dasm2.gui.bg_color_callback = lambda { |a2| COLORS[@dasmcol2[a2] || :default] }
		@match_func = nil
	end

	def curaddr1; @dasm1.gui.curaddr end
	def curaddr2; @dasm2.gui.curaddr end
	def curfunc1; @dasm1.find_function_start(curaddr1) end
	def curfunc2; @dasm2.find_function_start(curaddr2) end
	def func1; @func1 ||= set_status('funcs 1') { create_funcs(@dasm1) } end
	def func2; @func2 ||= set_status('funcs 2') { create_funcs(@dasm2) } end
	def funcstat1; @funcstat1 ||= set_status('func stats 1') { create_funcs_stats(func1, @dasm1) } end
	def funcstat2; @funcstat2 ||= set_status('func stats 2') { create_funcs_stats(func2, @dasm2) } end

	def paint
		draw_string_color(:black, @font_width, 3*@font_height, @status || 'idle')
	end

	def gui_update
		@dasm1.gui.gui_update rescue nil
		@dasm2.gui.gui_update rescue nil
		redraw
	end

	def set_status(st=nil)
		ost = @status
		@status = st
		redraw
		if block_given?
			ret = protect { yield }
			set_status ost
			ret
		end
	end

	def keypress(key)
		case key
		when ?A
			keypress(?D)
			keypress(?f)
			keypress(?i)
		when ?D
			disassemble_all
		when ?c
			disassemble
		when ?C
			disassemble(:disassemble_fast)
		when ?f
			funcstat1
			funcstat2
		when ?g
			inputbox('address to go', :text => Expression[@dasm1.gui.curaddr]) { |v|
				@dasm1.gui.focus_addr_autocomplete(v)
				@dasm2.gui.focus_addr_autocomplete(v)
			}
		when ?M
			show_match_funcs
		when ?m
			match_one_func(curfunc1, curfunc2)

		when ?r
			puts 'reload'
			load __FILE__
			gui_update

		when ?Q
			Gui.main_quit
		end
	end

	def keypress_ctrl(key)
		case key
		when ?C
			disassemble(:disassemble_fast_deep)
		when ?r
			inputbox('code to eval') { |c| messagebox eval(c).inspect[0, 512], 'eval' }
		end
	end

	def disassemble_all
		@func1 = @func2 = @funcstat1 = @funcstat2 = nil
		@dasm1.load_plugin 'dasm_all'
		@dasm2.load_plugin 'dasm_all'
		set_status('dasm_all 1') { @dasm1.dasm_all_section '.text' }
		set_status('dasm_all 2') { @dasm2.dasm_all_section '.text' }
		gui_update
	end

	def disassemble(method=:disassemble)
		@func1 = @func2 = @funcstat1 = @funcstat2 = nil
		set_status('dasm 1') {
			@dasm1.send(method, curaddr1)
			@dasm1.gui.focus_addr(curaddr1, :graph)
		}
		set_status('dasm 2') {
			@dasm2.send(method, curaddr2)
			@dasm2.gui.focus_addr(curaddr2, :graph)
		}
		gui_update
	end


	def show_match_funcs
		match_funcs

		gui_update
		Gui.main_iter
		list = [['addr 1', 'addr 2', 'score']]
		f1 = func1.keys
		f2 = func2.keys
		match_funcs.each { |a1, (a2, s)|
			list << [(@dasm1.get_label_at(a1) || Expression[a1]), (@dasm2.get_label_at(a2) || Expression[a2]), '%.4f' % s]
			f1.delete a1
			f2.delete a2
		}
		f1.each { |a1| list << [(@dasm1.get_label_at(a1) || Expression[a1]), '?', 'nomatch'] }
		f2.each { |a2| list << ['?', (@dasm2.get_label_at(a2) || Expression[a2]), 'nomatch'] }
		listwindow("matches", list) { |i| @dasm1.gui.focus_addr i[0], nil, true ; @dasm2.gui.focus_addr i[1], nil, true }
	end

	# func addr => { funcblock => list of funcblock to }
	def create_funcs(dasm)
		f = {}
		dasm.entrypoints.to_a.each { |ep| dasm.function[ep] ||= DecodedFunction.new }
		dasm.function.each_key { |a|
			next if not dasm.di_at(a)
			f[a] = create_func(dasm, a)
			Gui.main_iter
		}
		f
	end

	def create_func(dasm, a)
		h = {}
		todo = [a]
		while a = todo.pop
			next if h[a]
			h[a] = []
			dasm.decoded[a].block.each_to_samefunc(dasm) { |ta|
				next if not dasm.di_at(ta)
				todo << ta
				h[a] << ta
			}
		end
		h
	end

	def create_funcs_stats(f, dasm)
		fs = {}
		f.each { |a, g|
			fs[a] = create_func_stats(dasm, a, g)
			Gui.main_iter
		}
		fs
	end

	def create_func_stats(dasm, a, g)
		s = {}
		s[:blocks] = g.length

		s[:edges] = 0	# nr of edges
		s[:leaves] = 0	# nr of nodes with no successor
		s[:ext_calls] = 0	# nr of jumps out_of_func
		s[:loops] = 0	# nr of jump back

		todo = [a]
		done = []
		while aa = todo.pop
			next if done.include? aa
			done << aa
			todo.concat g[aa]

			s[:edges] += g[aa].length
			s[:leaves] += 1 if g[aa].empty?
			dasm.decoded[aa].block.each_to_otherfunc(dasm) { s[:ext_calls] += 1 }
		end

		# loop detection
		# find the longest distance to the root w/o loops
		g = g.dup
		while eliminate_one_loop(a, g)
			s[:loops] += 1
		end

		s
	end

	def eliminate_one_loop(a, g)
		stack = []
		index = {}
		reach_index = {}
		done = false

		curindex = 0
		
		trajan = lambda { |e|
			index[e] = curindex
			reach_index[e] = curindex
			curindex += 1
			stack << e
			g[e].each { |ne|
				if not index[ne]
					trajan[ne]
					break if done
					reach_index[e] = [reach_index[e], reach_index[ne]].min
				elsif stack.include? ne
					reach_index[e] = [reach_index[e], reach_index[ne]].min
				end
			}
			break if done
			if index[e] == reach_index[e]
				if (e == stack.last and not g[e].include? e)
					stack.pop
					next
				end
				# e is the entry in the loop, cut the loop here
				tail = reach_index.keys.find { |ee| reach_index[ee] == index[e] and g[ee].include? e }
				g[tail] -= [e]	# patch g, but don't modify the original g value (ie -= instead of delete)
				done = true	# one loop found & removed, try again
			end
		}

		trajan[a]
		done
	end

	def rematch_funcs
		@match_funcs = nil
		match_funcs
	end

	def match_funcs
		@match_funcs ||= {}

		layout_match = {}

		set_status('match func layout') {
		funcstat1.each { |a, s|
			next if @match_funcs[a]
			layout_match[a] = []
			funcstat2.each { |aa, ss|
				layout_match[a] << aa if s == ss
			}
			Gui.main_iter
		}
		}

		set_status('match funcs') {
		# refine the layout matching with actual function matching
		already_matched = []
		layout_match.each { |f1, list|
puts "matching #{Expression[f1]}" if $VERBOSE
begin
			f2 = (list - already_matched).sort_by { |f| match_func(f1, f, false, false) }.first
			if f2
				already_matched << f2
				score = match_func(f1, f2)
				@match_funcs[f1] = [f2, score]
			end
rescue Interrupt
	puts 'abort this one'
	sleep 0.2	# allow a 2nd ^c do escalate
end
			Gui.main_iter
		}
		}

		puts "matched #{@match_funcs.length} - unmatched #{func1.length - @match_funcs.length}"
		@match_funcs
	end

	def match_one_func(a1, a2)
		s = match_func(a1, a2)
		puts "match score: #{s}"
		@match_funcs ||= {}
		@match_funcs[a1] = [a2, s]
		gui_update
	end

	# return how much match a func in d1 and a func in d2
	def match_func(a1, a2, do_colorize=true, verb=true)
		f1 = func1[a1]
		f2 = func2[a2]
		raise "dasm1 has no function at #{Expression[a1]}" if not f1
		raise "dasm2 has no function at #{Expression[a2]}" if not f2
		todo1 = [a1]
		todo2 = [a2]
		done1 = []
		done2 = []
		score = 0.0	# average of the (local best) match_block scores
		score += 0.01 if @dasm1.get_label_at(a1) != @dasm2.get_label_at(a2)	# for thunks
		score_div = [f1.length, f2.length].max.to_f
		# XXX this is stupid and only good for perfect matches (and even then it may fail)
		# TODO handle block split etc (eg instr-level diff VS block-level)
		while a1 = todo1.shift
			next if done1.include? a1
			t = todo2.map { |a| [a, match_block(@dasm1.decoded[a1].block, @dasm2.decoded[a].block)] }
			a2 = t.sort_by { |a, s| s }.first
			if not a2
				break
			end
			score += a2[1] / score_div
			a2 = a2[0]
			done1 << a1
			done2 << a2
			todo1.concat f1[a1]
			todo2.concat f2[a2]
			todo2 -= done2
			colorize_blocks(a1, a2) if do_colorize
		end

		score += (f1.length - f2.length).abs * 3 / score_div	# block count difference -> +3 per block

		score
	end

	def match_block(b1, b2)
		# 0 = perfect match (same opcodes, same args)
		# 1 = same opcodes, same arg type
		# 2 = same opcodes, diff argtypes
		# 3 = some opcode difference
		# 4 = full block difference
		score = 0
		score_div = [b1.list.length, b2.list.length].max.to_f
		common_start = 0
		common_end = 0

		# basic diff-style: compare start while it's good, then end, then whats left
		# should handle most simples cases well
		len = [b1.list.length, b2.list.length].min
		while common_start < len and (s = match_instr(b1.list[common_start], b2.list[common_start])) <= 1
			score += s / score_div
			common_start += 1
		end

		while common_start+common_end < len and (s = match_instr(b1.list[-1-common_end], b2.list[-1-common_end])) <= 1
			score += s / score_div
			common_end += 1
		end

		# TODO improve the middle part matching (allow insertions/suppressions/swapping)
		b1.list[common_start..-1-common_end].zip(b2.list[common_start..-1-common_end]).each { |di1, di2|
			score += match_instr(di1, di2) / score_div
		}

		yield(common_start, common_end) if block_given?	# used by colorize_blocks

		score += (b1.list.length - b2.list.length).abs * 3 / score_div	# instr count difference -> +3 per instr

		score
	end

	def colorize_blocks(a1, a2)
		b1 = @dasm1.decoded[a1].block
		b2 = @dasm2.decoded[a2].block

		common_start = common_end = 0
		match_block(b1, b2) { |a, b| common_start = a ; common_end = b }

		b1.list[0..-1-common_end].zip(b2.list[0..-1-common_end]).each { |di1, di2|
			next if not di1 or not di2
			@dasmcol1[di1.address] = @dasmcol2[di2.address] = [:same, :similar, :badarg, :badop][match_instr(di1, di2)]
		}
		b1.list[-common_end..-1].zip(b2.list[-common_end..-1]).each { |di1, di2|
			next if not di1 or not di2
			@dasmcol1[di1.address] = @dasmcol2[di2.address] = [:same, :similar, :badarg, :badop][match_instr(di1, di2)]
		}
	end

	def match_instr(di1, di2)
		if not di1 or not di2 or di1.opcode.name != di2.opcode.name
			3
		elsif di1.instruction.args.map { |a| a.class } != di2.instruction.args.map { |a| a.class }
			2
		elsif di1.instruction.to_s.gsub(/loc_\w+/, 'loc_') != di2.instruction.to_s.gsub(/loc_\w+/, 'loc_')	# local labels	 TODO compare blocks targeted
			1
		else
			0
		end
	end

	# show in window 1 the match of the function found in win 2
	def sync1
		c2 = curfunc2
		if a1 = match_funcs.find_key { |k| match_funcs[k][0] == c2 }
			@dasm1.gui.focus_addr(a1)
		end
	end

	def sync2
		if a2 = match_funcs[curfunc1]
			@dasm2.gui.focus_addr(a2[0])
		end
	end
end

class BinDiffWindow < Gui::Window
	def initialize_window(d1=nil, d2=nil)
		self.widget = BinDiffWidget.new(d1, d2)
	end

	def build_menu
		menu = new_menu
		addsubmenu(menu, 'load file 1') { openfile('file 1') { |f| loadfile1(f) } }
		addsubmenu(menu, 'load file 2') { openfile('file 2') { |f| loadfile2(f) } }
		addsubmenu(menu)
		addsubmenu(menu, '_disassemble from there', '^C') { widget.disassemble(:disassemble_fast_deep) }
		addsubmenu(menu, 'co_mpare current functions', 'm') { widget.match_one_func(widget.curfunc1, widget.curfunc2) }
		addsubmenu(menu, 'compare all funct_ions', 'M') { widget.show_match_funcs }
		addsubmenu(menu, '_goto', 'g') { widget.keypress ?g }
		addsubmenu(menu)
		addsubmenu(menu, 'sync win 2', '2') { widget.sync2 }
		addsubmenu(menu, 'sync win 1', '1') { widget.sync1 }
		addsubmenu(menu)
		addsubmenu(menu, '_quit', 'Q') { Gui.main_quit }

		addsubmenu(@menu, '_File', menu)
	end

	def loadfile1(f)
		exe = AutoExe.orshellcode { Ia32.new }.decode_file(f)
		d = exe.init_disassembler
		Gui::DasmWindow.new("bindiff - 1 - #{f}").display(d)
		widget.dasm1 = d
	end

	def loadfile2(f)
		exe = AutoExe.orshellcode { Ia32.new }.decode_file(f)
		d = exe.init_disassembler
		Gui::DasmWindow.new("bindiff - 2 - #{f}").display(d)
		widget.dasm2 = d
	end
end
end

if $0 == __FILE__ and not defined? $bindiff_loaded
# allow reloading the file for easier diff algorithm test
$bindiff_loaded = true

require 'optparse'

$VERBOSE = true

# parse arguments
opts = {}
OptionParser.new { |opt|
	opt.banner = 'Usage: bindiff.rb [options] <executable> [<entrypoints>]'
	opt.on('-P <plugin>', '--plugin <plugin>', 'load a metasm disassembler plugin') { |h| (opts[:plugin] ||= []) << h }
	opt.on('-e <code>', '--eval <code>', 'eval a ruby code') { |h| (opts[:hookstr] ||= []) << h }
	opt.on('--map1 <mapfile>', 'load a map file (addr <-> name association)') { |f| opts[:map1] = f }
	opt.on('--map2 <mapfile>', 'load a map file (addr <-> name association)') { |f| opts[:map2] = f }
	opt.on('-c <header>', '--c-header <header>', 'read C function prototypes (for external library functions)') { |h| opts[:cheader] = h }
	opt.on('-a', '--autoload', 'loads all relevant files with same filename (.h, .map..)') { opts[:autoload] = true }
	opt.on('-v', '--verbose') { $VERBOSE = true }	# default
	opt.on('-q', '--no-verbose') { $VERBOSE = false }
	opt.on('-d', '--debug') { $DEBUG = $VERBOSE = true }
	opt.on('-A', 'match everything on start') { opts[:doit] = true }
}.parse!(ARGV)

if exename1 = ARGV.shift
	w1 = Metasm::Gui::DasmWindow.new("#{exename1} - bindiff1 - metasm disassembler")
	exe1 = w1.loadfile(exename1)
	if opts[:autoload]
		basename1 = exename1.sub(/\.\w\w?\w?$/, '')
		opts[:map1] ||= basename1 + '.map' if File.exist?(basename1 + '.map')
		opts[:cheader] ||= basename1 + '.h' if File.exist?(basename1 + '.h')
	end
end

if exename2 = ARGV.shift
	w2 = Metasm::Gui::DasmWindow.new("#{exename2} - bindiff2 - metasm disassembler")
	exe2 = w2.loadfile(exename2)
	if opts[:autoload]
		basename2 = exename2.sub(/\.\w\w?\w?$/, '')
		opts[:map2] ||= basename2 + '.map' if File.exist?(basename2 + '.map')
		opts[:cheader] ||= basename2 + '.h' if File.exist?(basename2 + '.h')
	end
end

if exe1
	dasm1 = exe1.init_disassembler
	dasm1.load_map opts[:map1] if opts[:map1]
	dasm1.parse_c_file opts[:cheader] if opts[:cheader]
end

if exe2
	dasm2 = exe2.init_disassembler
	dasm2.load_map opts[:map2] if opts[:map2]
	dasm2.parse_c_file opts[:cheader] if opts[:cheader]
end

ep = ARGV.dup

w1.dasm_widget.focus_addr ep.first if w1 and not ep.empty?
w2.dasm_widget.focus_addr ep.first if w2 and not ep.empty?

opts[:plugin].to_a.each { |p| dasm1.load_plugin(p) if dasm1 ; dasm2.load_plugin(p) if dasm2 }
opts[:hookstr].to_a.each { |f| eval f }

ep.each { |e| dasm1.disassemble_fast_deep(e) if dasm1 ; dasm2.disassemble_fast_deep(e) if dasm2 }

bd = Metasm::BinDiffWindow.new(dasm1, dasm2)

bd.widget.keypress ?A if opts[:doit]

Metasm::Gui.main

end