#    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 is a linux/x86 debugger with a console interface
#

require 'metasm'

module Ansi
	CursHome = "\e[H".freeze
	ClearLineAfter  = "\e[0K"
	ClearLineBefore = "\e[1K"
	ClearLine = "\e[2K"
	ClearScreen = "\e[2J"
	def self.set_cursor_pos(y=1,x=1) "\e[#{y};#{x}H" end
	Reset = "\e[m"
	Colors = [:black, :red, :green, :yellow, :blue, :magenta, :cyan, :white, :aoeu, :reset]
	def self.color(*args)
		fg = true
		"\e[" << args.map { |a|
			case a
			when :bold; 2
			when :negative; 7
			when :normal; 22
			when :positive; 27
			else
				if col = Colors.index(a)
					add = (fg ? 30 : 40)
					fg = false
					col+add
				end
			end
		}.compact.join(';') << 'm'
	end
	def self.hline(len) "\e(0" << 'q'*len << "\e(B" end

	TIOCGWINSZ = 0x5413
	TCGETS = 0x5401
	TCSETS = 0x5402
	CANON = 2
	ECHO  = 8
	def self.get_terminal_size
		s = ''.ljust(8)
		$stdin.ioctl(TIOCGWINSZ, s) >= 0 ? s.unpack('SS') : [80, 25]
	end
	def self.set_term_canon(bool)
		ttys = ''.ljust(256)
		$stdin.ioctl(TCGETS, ttys)
		tty = ttys.unpack('C*')
		if bool
			tty[12] &= ~(ECHO|CANON)
		else
			tty[12] |= ECHO|CANON
		end
		$stdin.ioctl(TCSETS, tty.pack('C*'))
	end

	ESC_SEQ = {'A' => :up, 'B' => :down, 'C' => :right, 'D' => :left,
		'1~' => :home, '2~' => :inser, '3~' => :suppr, '4~' => :end,
		'5~' => :pgup, '6~' => :pgdown,
		'P' => :f1, 'Q' => :f2, 'R' => :f3, 'S' => :f4,
		'15~' => :f5, '17~' => :f6, '18~' => :f7, '19~' => :f8,
		'20~' => :f9, '21~' => :f10, '23~' => :f11, '24~' => :f12,
		'[A' => :f1, '[B' => :f2, '[C' => :f3, '[D' => :f4, '[E' => :f5,
		'H' => :home, 'F' => :end,
	}
	def self.getkey
		c = $stdin.getc
		raise 'nonblocking $stdin?' if not c
		return c if c != ?\e
		c = $stdin.getc
		if c != ?[ and c != ?O
			$stdin.ungetc c
			return ?\e
		end
		seq = ''
		loop do
			c = $stdin.getc
			seq << c
			case c; when ?a..?z, ?A..?Z, ?~; break end
		end
		ESC_SEQ[seq] || seq
	end
end

class LinDebug
	attr_accessor :win_reg_height, :win_data_height, :win_code_height, :win_prpt_height
	def init_screen
		Ansi.set_term_canon(true)
		@win_reg_height = 2
		@win_data_height = 20
		@win_code_height = 20
		resize
	end

	def fini_screen
		Ansi.set_term_canon(false)
		$stdout.write Ansi.color(:normal, :reset)
		$stdout.flush
	end

	def win_data_start; @win_reg_height end
	def win_code_start; win_data_start+win_data_height end
	def win_prpt_start; win_code_start+win_code_height end

	Color = {:changed => Ansi.color(:cyan, :bold), :border => Ansi.color(:green),
		:normal => Ansi.color(:white, :black, :normal), :hilight => Ansi.color(:blue, :white, :normal),
		:status => Ansi.color(:black, :cyan)}

	# yields but keep from reentring (return defretval in this case)
	def once(name, defretval=nil)
		@once ||= {}
		if not @once[name]
			@once[name] = true
			begin
				defretval = yield
			ensure
				@once[name] = false
			end
		end
		defretval
	end

	attr_accessor :dataptr, :codeptr, :dbg, :promptlog, :command
	def initialize(dbg)
		@dbg = dbg
		@dbg.set_log_proc { |l| add_log l }
		@datafmt = 'db'

		@prompthistlen = 20
		@prompthistory = []
		@promptloglen = 200
		@promptlog = []
		@promptbuf = ''
		@promptpos = 0
		@log_off = 0
		@console_width = 80
		@oldregs = {}

		@running = false
		@focus = :prompt
		@command = {}
		load_commands
		trap('WINCH') { resize }
	end

	def init_rs
		@codeptr = @dataptr = @dbg.pc	# avoid initial faults
	end

	def main_loop
		begin
			begin
				init_screen
				init_rs
				main_loop_inner
			rescue Errno::ESRCH
				log "target does not exist anymore"
			ensure
				fini_screen
				$stdout.print Ansi.set_cursor_pos(@console_height, 1)
			end
		rescue
			puts $!, $!.backtrace
		end
		puts @promptlog.last
	end

	# optimize string to display to stdout
	# currently only skips unchanged lines
	# could also match end of lines (spaces), but would need to check for color codes etc
	def optimize_screen(buf)
	end

	# display the text buffer screenlines to the screen, leaves the cursor at (cursx, cursy), converts cursor pos from 0-base to 1-base
	# screenlines is a big text buffer with 1 line per tobeshown screen line (e.g. no Ansi cursor pos)
	# screenlines must be screen wide
	def display_screen(screenlines, cursx, cursy)

		@oldscreenbuf ||= []
		lines = screenlines.lines
		oldlines = @oldscreenbuf
		@oldscreenbuf = lines
		screenlines = lines.zip(oldlines).map { |l, ol| l == ol ? "\n" : l }.join

		while screenlines[-1] == ?\n
			screenlines.chop!
		end
		starty = 1
		while screenlines[0] == ?\n
			screenlines = screenlines[1..-1]
			starty += 1
		end

		$stdout.write Ansi.set_cursor_pos(starty, 1) + screenlines + Ansi.set_cursor_pos(cursy+1, cursx+1)

		$stdout.flush
	end

	def update
		return if not @running
		display_screen updateregs + updatedata + updatecode + updateprompt, @promptpos+1, @console_height-2
	end

	def updateregs
		once(:updateregs, "\n\n") { _updateregs }
	end

	def _updateregs
		pvrsz = 0
		words = @dbg.register_list.map { |r|
			rs = r.to_s.rjust(pvrsz)
			pvrsz = rs.length
			rv = @dbg[r]
			["#{rs}=%0#{@dbg.register_size[r]/4}X " % rv,
				(@oldregs[r] != rv)]
		} + @dbg.flag_list.map { |fl|
			fv = @dbg.get_flag(fl)
			[fv ? fl.to_s.upcase : fl.to_s.downcase,
				(@oldregs[fl] != fv)]
		}

		text = ' '
		linelen = 1	# line length w/o ansi colors

		owr = @win_reg_height
		@win_reg_height = 1
		words.each { |w, changed|
			if linelen + w.length >= @console_width - 1
				text << (' '*([@console_width-linelen, 0].max)) << "\n "
				linelen = 1
				@win_reg_height += 1
			end

			text << Color[:changed] if changed
			text << w
			text << Color[:normal] if changed
			text << ' '

			linelen += w.length+1
		}
		resize if owr != @win_reg_height
		text << (' '*([@console_width-linelen, 0].max)) << "\n"
	end

	def updatecode
		once(:updatecode, "...\n"*@win_code_height) { _updatecode }
	end

	def _updatecode
		if @codeptr
			addr = @codeptr
		elsif @oldregs[@dbg.register_pc] and @oldregs[@dbg.register_pc] < @dbg.pc and @oldregs[@dbg.register_pc] + 8 >= @dbg.pc
			addr = @oldregs[@dbg.register_pc]
		else
			addr = @dbg.pc
		end
		@codeptr = addr

		addrsz = @dbg.register_size[@dbg.register_pc]
		addrfmt = "%0#{addrsz/4}X"
		if not @dbg.addr2module(addr) and @dbg.shortname !~ /remote/
			base = addr & ((1 << addrsz) - 0x1000)
			@noelfsig ||= {}	# cache elfmagic notfound
			if not @noelfsig[base] and base < ((1 << addrsz) - 0x1_0000)
				self.statusline = " scanning for elf header at #{addrfmt % base}"
				128.times {
					@statusline = " scanning for elf header at #{addrfmt % base}"
					if not @noelfsig[base] and @dbg[base, Metasm::ELF::MAGIC.length] == Metasm::ELF::MAGIC
						@dbg.loadsyms(base, base.to_s(16))
						break
					else
						@noelfsig[base] = true	# XXX an elf may be mmaped here later..
					end
					base -= 0x1000
					break if base < 0
				}
				self.statusline = nil
			end
		end

		text = ''
		text << Color[:border]
		title = @dbg.addrname(addr)
		pre  = [@console_width-100, 6].max
		post = @console_width - (pre + title.length + 2)
		text << Ansi.hline(pre) << ' ' << title << ' ' << Ansi.hline(post) << Color[:normal] << "\n"

		seg = ''
		seg = ('%04X' % @dbg['cs']) << ':' if @dbg.cpu.shortname =~ /ia32|x64/

		cnt = @win_code_height
		while (cnt -= 1) > 0
			if @dbg.symbols[addr]
				text << ('    ' << @dbg.symbols[addr] << ?:) << Ansi::ClearLineAfter << "\n"
				break if (cnt -= 1) <= 0
			end
			text << Color[:hilight] if addr == @dbg.pc
			text << seg
			if @dbg.shortname =~ /remote/ and @dbg.realmode
				text << (addrfmt % (addr - 16*@dbg['cs']))
			else
				text << (addrfmt % addr)
			end
			di = @dbg.di_at(addr)
			di = nil if di and addr < @dbg.pc and addr+di.bin_length > @dbg.pc
			len = (di ? di.bin_length : 1)
			text << '  '
			text << @dbg.memory[addr, [len, 10].min].to_s.unpack('C*').map { |c| '%02X' % c }.join.ljust(22)
			if di
				text <<
				if addr == @dbg.pc
					"*#{di.instruction}".ljust([@console_width-(addrsz/4+seg.length+24), 0].max)
				else
					" #{di.instruction}" << Ansi::ClearLineAfter
				end
			else
				text << ' <unk>' << Ansi::ClearLineAfter
			end
			text << Color[:normal] if addr == @dbg.pc
			addr += len
			text << "\n"
		end
		text
	end

	def updatedata
		once(:updatedata, "...\n"*@win_data_height) { _updatedata }
	end

	def _updatedata
		addrsz = @dbg.register_size[@dbg.register_pc]
		addrfmt = "%0#{addrsz/4}X"

		@dataptr &= ((1 << addrsz) - 1)
		addr = @dataptr

		text = ''
		text << Color[:border]
		title = @dbg.addrname(addr)
		pre  = [@console_width-100, 6].max
		post = [@console_width - (pre + title.length + 2), 0].max
		text << Ansi.hline(pre) << ' ' << title << ' ' << Ansi.hline(post) << Color[:normal] << "\n"

		seg = ''
		seg = ('%04X' % @dbg['ds']) << ':' if @dbg.cpu.shortname =~ /^ia32/

		cnt = @win_data_height
		while (cnt -= 1) > 0
			raw = @dbg.memory[addr, 16].to_s
			text << seg << (addrfmt % addr) << '  '
			case @datafmt
			when 'db'; text << raw[0,8].unpack('C*').map { |c| '%02x ' % c }.join << ' ' <<
				   raw[8,8].to_s.unpack('C*').map { |c| '%02x ' % c }.join
			when 'dw'; text << raw.unpack('S*').map { |c| '%04x ' % c }.join
			when 'dd'; text << raw.unpack('L*').map { |c| '%08x ' % c }.join
			end
			text << ' ' << raw.unpack('C*').map { |c| (0x20..0x7e).include?(c) ? c : 0x2e }.pack('C*')
			text << Ansi::ClearLineAfter << "\n"
			addr += 16
		end
		text
	end

	def updateprompt
		once(:updateprompt, "\n"*@win_prpt_height) { _updateprompt }
	end

	def _updateprompt
		text = ''
		text << Color[:border] << Ansi.hline(@console_width) << Color[:normal] << "\n"

		@log_off = @promptlog.length - 2 if @log_off >= @promptlog.length
		@log_off = 0 if @log_off < 0
		len = @win_prpt_height - 2
		len.times { |i|
			i += @promptlog.length - @log_off - len
			text << ((@promptlog[i] if i >= 0) || '')
			text << Ansi::ClearLineAfter << "\n"
		}
		text << ':' << @promptbuf << Ansi::ClearLineAfter << "\n"
		text << Color[:status] << statusline.chomp.ljust(@console_width) << Color[:normal]
	end

	def statusline
		@statusline ||= '    Enter a command (help for help)'
	end
	def statusline=(s)
		@statusline = s
		update
	end

	def resize
		@console_height, @console_width = Ansi.get_terminal_size
		@win_data_height = 1 if @win_data_height < 1
		@win_code_height = 1 if @win_code_height < 1
		if @win_data_height + @win_code_height + @win_reg_height + 3 > @console_height
			@win_data_height = @console_height/2 - 4
			@win_code_height = @console_height/2 - 4
		end
		@win_prpt_height = @console_height-(@win_data_height+@win_code_height+@win_reg_height) - 1
		@oldscreenbuf = []
		update
	end

	def log(*strs)
	    strs.each { |str|
		raise str.inspect if not str.kind_of? ::String
		str = str.chomp
		if str.length > @console_width
			# word wrap
			str.scan(/.{0,#@console_width}/) { |str_| log str_ }
			return
		end
		@promptlog << str
		@promptlog.shift if @promptlog.length > @promptloglen
	    }
	end

	def add_log(l)
		log l
		puts l if not @running
		update rescue puts l
	end

	def exec_prompt
		@log_off = 0
		log ':'+@promptbuf
		return if @promptbuf == ''
		str = @promptbuf
		@prompthistory << @promptbuf
		@prompthistory.shift if @prompthistory.length > @prompthistlen
		@promptbuf = ''
		@promptpos = @promptbuf.length

		cmd, str = str.split(/\s+/, 2)
		if @command.has_key? cmd
			@command[cmd].call(str.to_s)
		else
			if cmd and (poss = @command.keys.find_all { |c| c[0, cmd.length] == cmd }).length == 1
				@command[poss.first].call(str.to_s)
			else
				log 'unknown command'
			end
		end
	end

	def preupdate
		@dbg.register_list.each { |r| @oldregs[r] = @dbg[r] }
		@dbg.flag_list.each { |fl| @oldregs[fl] = @dbg.get_flag(fl) }
	end

	def updatecodeptr
		@codeptr ||= @dbg.pc
		if @codeptr > @dbg.pc or @codeptr < @dbg.pc - 6*@win_code_height
			@codeptr = @dbg.pc
		elsif @codeptr != @dbg.pc
			addr = @codeptr
			addrs = []
			while addr <= @dbg.pc
				addrs << addr
				addrs << addr if @dbg.symbols[addr]
				o = ((di = @dbg.di_at(addr)) ? di.bin_length : 0)
				addr += ((o == 0) ? 1 : o)
			end
			if addrs.length > @win_code_height-4
				@codeptr = addrs[-(@win_code_height-4)]
			end
		end
		updatedataptr
	end

	def updatedataptr
	end

	def singlestep
		self.statusline = ' target singlestepping...'
		preupdate
		@dbg.singlestep_wait
		updatecodeptr
		@statusline = nil
	end
	def stepover
		self.statusline = ' target running...'
		preupdate
		@dbg.stepover_wait
		updatecodeptr
		@statusline = nil
	end
	def cont(*a)
		self.statusline = ' target running...'
		preupdate
		@dbg.continue_wait(*a)
		updatecodeptr
		@statusline = nil
	end
	def stepout
		self.statusline = ' target running...'
		preupdate
		@dbg.stepout_wait
		updatecodeptr
		@statusline = nil
	end
	def syscall
		self.statusline = ' target running to next syscall...'
		preupdate
		@dbg.syscall_wait
		updatecodeptr
		@statusline = nil
	end

	def main_loop_inner
		@prompthistory = ['']
		@histptr = nil
		@running = true
		update
		while @running
			if not IO.select [$stdin], nil, nil, 0
				begin
					update
				rescue Errno::ESRCH
					break
				end
			end
			break if handle_keypress(Ansi.getkey)
		end
	end

	def handle_keypress(k)
			case k
			when ?\4; log 'exiting'; return true	 # eof
			when ?\e; @focus = :prompt
			when :f5;  cont
			when :f6;  syscall
			when :f10; stepover
			when :f11; singlestep
			when :f12; stepout
			when :up
				case @focus
				when :prompt
					if not @histptr
						@prompthistory << @promptbuf
						@histptr = 2
					else
						@histptr += 1
						@histptr = 1 if @histptr > @prompthistory.length
					end
					@promptbuf = @prompthistory[-@histptr].dup
					@promptpos = @promptbuf.length
				when :data
					@dataptr -= 16
				when :code
					@codeptr ||= @dbg.pc
					@codeptr -= (1..10).find { |off|
						di = @dbg.di_at(@codeptr-off)
						di.bin_length == off if di
					} || 10
				end
			when :down
				case @focus
				when :prompt
					if not @histptr
						@prompthistory << @promptbuf
						@histptr = @prompthistory.length
					else
						@histptr -= 1
						@histptr = @prompthistory.length if @histptr < 1
					end
					@promptbuf = @prompthistory[-@histptr].dup
					@promptpos = @promptbuf.length
				when :data
					@dataptr += 16
				when :code
					@codeptr ||= @dbg.pc
					di = @dbg.di_at(@codeptr)
					@codeptr += (di ? (di.bin_length || 1) : 1)
				end
			when :left;  @promptpos -= 1 if @promptpos > 0
			when :right; @promptpos += 1 if @promptpos < @promptbuf.length
			when :home;  @promptpos = 0
			when :end;   @promptpos = @promptbuf.length
			when :backspace, ?\x7f; @promptbuf[@promptpos-=1, 1] = '' if @promptpos > 0
			when :suppr; @promptbuf[@promptpos, 1] = '' if @promptpos < @promptbuf.length
			when :pgup
				case @focus
				when :prompt; @log_off += @win_prpt_height-3
				when :data; @dataptr -= 16*(@win_data_height-1)
				when :code
					@codeptr ||= @dbg.pc
					(@win_code_height-1).times {
						@codeptr -= (1..10).find { |off|
							di = @dbg.di_at(@codeptr-off)
							di.bin_length == off if di
						} || 10
					}
				end
			when :pgdown
				case @focus
				when :prompt; @log_off -= @win_prpt_height-3
				when :data; @dataptr += 16*(@win_data_height-1)
				when :code
					@codeptr ||= @dbg.pc
					(@win_code_height-1).times { @codeptr += ((o = @dbg.di_at(@codeptr)) ? [o.bin_length, 1].max : 1) }
				end
			when ?\t
				if not @promptbuf[0, @promptpos].include? ' '
					poss = @command.keys.find_all { |c| c[0, @promptpos] == @promptbuf[0, @promptpos] }
					if poss.length > 1
						log poss.sort.join(' ')
					elsif poss.length == 1
						@promptbuf[0, @promptpos] = poss.first + ' '
						@promptpos = poss.first.length+1
					end
				end
			when ?\n
				@histptr = nil
				begin
					exec_prompt
				rescue Exception
					log "error: #$!", *$!.backtrace
				end
			when ?\ ..?~
				@promptbuf[@promptpos, 0] = k.chr
				@promptpos += 1
			else log "unknown key pressed #{k.inspect}"
			end
			nil
	end

	def load_commands
		@command['kill'] = lambda { |str|
			@dbg.kill
			@running = false
			log 'killed'
		}
		@command['quit'] = @command['detach'] = @command['exit'] = lambda { |str|
			@dbg.detach
			@running = false
		}
		@command['closeui'] = lambda { |str|
			@running = false
		}
		@command['bpx'] = lambda { |str|
			@dbg.bpx @dbg.resolve(str)
		}
		@command['bphw'] = @command['hwbp'] = lambda { |str|
			type, str = str.split(/\s+/, 2)
			@dbg.hwbp @dbg.resolve(str.to_s), type
		}
		@command['bt'] = lambda { |str| @dbg.stacktrace { |a,t| add_log "#{'%x' % a} #{t}" } }
		@command['d'] =  lambda { |str| @dataptr = @dbg.resolve(str) if str.length > 0 }
		@command['db'] = lambda { |str| @datafmt = 'db' ; @dataptr = @dbg.resolve(str) if str.length > 0 }
		@command['dw'] = lambda { |str| @datafmt = 'dw' ; @dataptr = @dbg.resolve(str) if str.length > 0 }
		@command['dd'] = lambda { |str| @datafmt = 'dd' ; @dataptr = @dbg.resolve(str) if str.length > 0 }
		@command['r'] =  lambda { |str|
			r, str = str.split(/\s+/, 2)
			if r == 'fl'
				@dbg.toggle_flag(str.to_sym)
			elsif not @dbg[r]
				log "bad reg #{r}"
			elsif str and str.length > 0
				@dbg[r] = @dbg.resolve(str)
			else
				log "#{r} = #{@dbg[r]}"
			end
		}
		@command['g'] = lambda { |str|
			@dbg.go @dbg.resolve(str)
		}
		@command['u'] = lambda { |str| @codeptr = @dbg.resolve(str) }
		@command['ruby'] = lambda { |str| instance_eval str }
		@command['wd'] = lambda { |str|
			@focus = :data
			if str.length > 0
				@win_data_height = @dbg.resolve(str)
				resize
			end
		}
		@command['wc'] = lambda { |str|
			@focus = :code
			if str.length > 0
				@win_code_height = @dbg.resolve(str)
				resize
			end
		}
		@command['wp'] = lambda { |str| @focus = :prompt }
		@command['?'] = lambda { |str|
			val = @dbg.resolve(str)
			log "#{val} 0x#{val.to_s(16)} #{[val].pack('L').inspect}"
		}
		@command['syscall'] = lambda { |str|
			@dbg.syscall_wait(str)
		}
	end
end


if $0 == __FILE__
	require 'optparse'
	opts = { :sc_cpu => 'Ia32' }
	OptionParser.new { |opt|
		opt.on('-m map', '--map filemap') { |f| opts[:filemap] = f }
		opt.on('--cpu cpu') { |c| opts[:sc_cpu] = c }
	}.parse!(ARGV)

	case ARGV.first
	when /^(tcp:|udp:)?..+:/, /^ser:/
		opts[:sc_cpu] = eval(opts[:sc_cpu]) if opts[:sc_cpu] =~ /[.(\s:]/
		opts[:sc_cpu] = opts[:sc_cpu].new if opts[:sc_cpu].kind_of?(::Class)
		dbg = Metasm::GdbRemoteDebugger.new(ARGV.first, opts[:sc_cpu])
	else
		dbg = Metasm::LinDebugger.new(ARGV.join(' '))
	end
	dbg.load_map(opts[:filemap]) if opts[:filemap]
	LinDebug.new(dbg).main_loop
end
