kaitai-struct-visualizer 0.3

data/lib/kaitai/struct/visualizer/tree.rb

@@ -0,0 +1,253 @@
+require 'benchmark'
+
+require 'kaitai/struct/visualizer/version'
+require 'kaitai/struct/visualizer/node'
+require 'kaitai/struct/visualizer/hex_viewer'
+
+module Kaitai::Struct::Visualizer
+class Tree
+  def initialize(ui, st)
+    @ui = ui
+    @st = st
+    @root = Node.new(self, st, 0)
+    @root.id = '[root]'
+    @max_scr_ln = @ui.rows - 3
+
+    @hv_shift_x = @ui.cols - HexViewer.line_width - 1
+
+    @cur_io = nil
+    @hv = HexViewer.new(ui, nil, @hv_shift_x, self)
+    @hv_hidden = false
+
+    @cur_line = 0
+    @cur_shift = 0
+    @do_exit = false
+  end
+
+  def run
+    c = nil
+    loop {
+      t = redraw
+
+      if @cur_node.nil? and not @cur_line.nil?
+        # gone beyond the end of the tree
+        @cur_line = @root.height - 1
+        clamp_cursor
+        redraw
+      end
+
+      raise '@cur_line is undetermined' if @cur_line.nil?
+      raise '@cur_node is undetermined' if @cur_node.nil?
+
+      thv = Benchmark.realtime {
+        unless @hv_hidden
+          hv_update_io
+
+          unless @cur_node.pos1.nil?
+            if (@hv.addr < @cur_node.pos1) or (@hv.addr >= @cur_node.pos2)
+              @hv.addr = @cur_node.pos1
+              @hv.ensure_visible
+            end
+          end
+
+          @hv.redraw
+          regs = highlight_regions(4)
+          @hv.highlight(regs)
+        end
+      }
+
+      @ui.goto(0, @max_scr_ln + 1)
+      printf "all: %d, tree: %d, tree_draw: %d, hexview: %d, ln: %d, ", (t + thv) * 1e6, t * 1e6, @draw_time * 1e6, thv * 1e6, @ln
+      puts "highlight = #{@cur_node.pos1}..#{@cur_node.pos2}"
+      #puts "keypress: #{c.inspect}"
+
+      begin
+        process_keypress
+      rescue EOFError => e
+        @ui.message_box_exception(e)
+      rescue Kaitai::Struct::Stream::UnexpectedDataError => e
+        @ui.message_box_exception(e)
+      end
+
+      return if @do_exit
+
+      clamp_cursor
+    }
+  end
+
+  def process_keypress
+    c = @ui.read_char_mapped
+    case c
+    when :up_arrow
+      @cur_line -= 1
+      @cur_node = nil
+    when :down_arrow
+      @cur_line += 1
+      @cur_node = nil
+    when :left_arrow
+      if @cur_node.open?
+        @cur_node.close
+      else
+        par = @cur_node.parent
+        if par
+          @cur_line = nil
+          @cur_node = par
+        end
+      end
+    when :right_arrow
+      if @cur_node.openable?
+        if @cur_node.open?
+          @cur_line += 1
+          @cur_node = nil
+        else
+          @cur_node.open
+        end
+      end
+    when :home
+      @cur_line = @cur_shift = 0
+      @cur_node = nil
+    when :end
+      @cur_line = @root.height - 1
+      @cur_node = nil
+    when :pg_up
+      @cur_line -= 20
+      @cur_node = nil
+    when :pg_dn
+      @cur_line += 20
+      @cur_node = nil
+    when :enter
+      if @cur_node.hex?
+        @ui.clear
+        hv = HexViewer.new(@ui, @cur_node.value)
+        hv.redraw
+        hv.run
+        @ui.clear
+        redraw
+      else
+        @cur_node.toggle
+      end
+    when :tab
+      @hv.run
+    when 'H'
+      @hv_hidden = !@hv_hidden
+      @ui.clear
+      redraw
+    when 'q'
+      @do_exit = true
+    end
+  end
+
+  def clamp_cursor
+    if @cur_line
+      @cur_line = 0 if @cur_line < 0
+
+      if @cur_line - @cur_shift < 0
+        @cur_shift = @cur_line
+      end
+      if @cur_line - @cur_shift > @max_scr_ln
+        @cur_shift = @cur_line - @max_scr_ln
+      end
+    end
+  end
+
+  def redraw
+    @draw_time = 0
+    Benchmark.realtime {
+      @ui.clear
+      @ln = 0
+      draw_rec(@root)
+    }
+  end
+
+  def draw_rec(n)
+    scr_ln = @ln - @cur_shift
+    return if @cur_node and scr_ln > @max_scr_ln
+
+    if @ln == @cur_line
+      # Seeking cur_node by cur_line
+      @cur_node = n
+      @ui.bg_color = :gray
+      @ui.fg_color = :black
+    elsif @cur_node == n
+      # Seeking cur_line by cur_node
+      @cur_line = @ln
+      @ui.bg_color = :gray
+      @ui.fg_color = :black
+    end
+
+    @draw_time += Benchmark.realtime {
+#      n.draw(@ui) if scr_ln >= 0
+      n.draw(@ui) if scr_ln >= 0 and scr_ln <= @max_scr_ln
+    }
+
+    @ui.reset_colors if @ln == @cur_line
+    @ln += 1
+    if n.open?
+      n.children.each { |ch|
+        draw_rec(ch)
+        break if scr_ln > @max_scr_ln
+      }
+    end
+  end
+
+  def do_exit
+    @do_exit = true
+  end
+
+  def hv_update_io
+    io = @cur_node.io
+    if io != @cur_io
+      @cur_io = io
+      io.seek(0)
+      buf = io.read_bytes_full
+      @hv.buf = buf
+
+#      @hv.redraw
+    end
+  end
+
+  def highlight_regions(max_levels)
+    node = @cur_node
+    r = []
+    max_levels.times { |i|
+      return r if node.nil?
+      r << [node.pos1, node.pos2]
+      node = node.parent
+    }
+    r
+  end
+
+  def tree_width
+    if @hv_hidden
+      @ui.cols
+    else
+      @hv_shift_x
+    end
+  end
+
+  def self.explore_object(obj, level)
+    root = Node.new(obj, level)
+    if obj.is_a?(Fixnum) or obj.is_a?(String)
+      # do nothing else
+    elsif obj.is_a?(Array)
+      root = Node.new(obj, level)
+      obj.each_with_index { |el, i|
+        n = explore_object(el, level + 1)
+        n.id = i
+        root.add(n)
+      }
+    else
+      root = Node.new(obj, level)
+      obj.instance_variables.each { |k|
+        k = k.to_s
+        next if k =~ /^@_/
+        el = obj.instance_eval(k)
+        n = explore_object(el, level + 1)
+        n.id = k
+        root.add(n)
+      }
+    end
+    root
+  end
+end
+end

data/lib/kaitai/struct/visualizer/version.rb

@@ -0,0 +1,7 @@
+module Kaitai
+  module Struct
+    module Visualizer
+      VERSION = '0.3'
+    end
+  end
+end

data/lib/kaitai/struct/visualizer/visualizer.rb

@@ -0,0 +1,47 @@
+require 'tmpdir'
+
+require 'kaitai/struct/visualizer/version'
+require 'kaitai/tui'
+require 'kaitai/struct/visualizer/tree'
+
+# TODO: should be inside compiled files
+require 'zlib'
+require 'stringio'
+
+module Kaitai::Struct::Visualizer
+class Visualizer
+  def initialize(bin_fn, formats_fn)
+    @bin_fn = bin_fn
+    @formats_fn = formats_fn
+    @primary_format = @formats_fn.shift
+
+    main_class_name = compile_format(@primary_format)
+
+    @formats_fn.each { |fn|
+      compile_format(fn)
+    }
+
+    main_class = Kernel::const_get(main_class_name)
+    @data = main_class.from_file(@bin_fn)
+
+    load_exc = nil
+    begin
+      @data._read
+    rescue EOFError => e
+      load_exc = e
+    rescue Kaitai::Struct::Stream::UnexpectedDataError => e
+      load_exc = e
+    end
+
+    @ui = Kaitai::TUI.new
+    @tree = Tree.new(@ui, @data)
+
+    @tree.redraw
+    @ui.message_box_exception(load_exc) if load_exc
+  end
+
+  def run
+    @tree.run
+  end
+end
+end

data/lib/kaitai/struct/visualizer/visualizer_main.rb

@@ -0,0 +1,24 @@
+require 'kaitai/struct/visualizer/version'
+require 'kaitai/struct/visualizer/visualizer'
+
+module Kaitai::Struct::Visualizer
+
+class ExternalCompilerVisualizer < Visualizer
+  def compile_format(fn)
+    main_class_name = nil
+    Dir.mktmpdir { |code_dir|
+      system("ksc -- --debug -t ruby '#{fn}' -d '#{code_dir}'")
+      exit $?.exitstatus if $?.exitstatus != 0
+
+      compiled_path = Dir.glob("#{code_dir}/*.rb")[0]
+
+      require compiled_path
+
+      main_class_name = File.readlines(compiled_path).grep(/^class /)[0].strip.gsub(/^class /, '').gsub(/ <.*$/, '')
+    }
+
+    return main_class_name
+  end
+end
+
+end

data/lib/kaitai/struct/visualizer/visualizer_ruby.rb

@@ -0,0 +1,18 @@
+require 'ks_ruby_compiler'
+
+class RubyCompilerVisualizer extends Visualizer
+  def compile
+    Dir.mktmpdir { |code_dir|
+      compiled_path = "#{code_dir}/compiled.rb"
+      @compiler = CompileToRuby.new(@format_fn, compiled_path)
+      @compiler.compile
+
+      require compiled_path
+
+      main_class_name = @compiler.type2class(@compiler.desc['meta']['id'])
+      #puts "Main class: #{main_class_name}"
+      main_class = Kernel::const_get(main_class_name)
+      @data = main_class.from_file(@bin_fn)
+    }
+  end
+end

data/lib/kaitai/tui.rb

@@ -0,0 +1,169 @@
+# coding: utf-8
+require 'io/console'
+
+module Kaitai
+
+class TUI
+  attr_reader :cols
+  attr_reader :rows
+
+  def initialize
+    @cols = `tput cols`.to_i
+    @rows = `tput lines`.to_i
+
+    @seq_clear = `tput clear`
+    @seq_sgr0 = `tput sgr0`
+    @seq_fgcolor = []
+    @seq_bgcolor = []
+  end
+
+  def clear
+    print @seq_clear
+  end
+
+  ##
+  # Put the cursor up to screen position (x, y). First line is 0,
+  # first column is 0.
+  def goto(x, y)
+    #print `tput cup #{y} #{x}`
+    printf "\e[%d;%dH", y + 1, x + 1
+  end
+
+  COLORS = {
+    :black => 0,
+    :gray => 7,
+    :gray0 => 232,
+    :gray1 => 233,
+    :gray2 => 234,
+    :gray3 => 235,
+    :gray4 => 236,
+    :gray5 => 237,
+    :gray6 => 238,
+    :gray7 => 239,
+    :gray8 => 240,
+    :gray9 => 241,
+    :gray10 => 242,
+    :gray11 => 243,
+    :gray12 => 244,
+    :gray13 => 245,
+    :gray14 => 246,
+    :gray15 => 247,
+    :gray16 => 248,
+    :gray17 => 249,
+    :gray18 => 250,
+    :gray19 => 251,
+    :gray20 => 252,
+    :gray21 => 253,
+    :gray22 => 254,
+    :gray23 => 255,
+  }
+
+  def fg_color=(col)
+    #print @seq_fgcolor[col] ||= `tput setaf #{col}`
+    code = COLORS[col]
+    raise "Invalid color: #{col}" unless code
+    print "\e[38;5;#{code}m"
+  end
+
+  def bg_color=(col)
+    #print @seq_bgcolor[col] ||= `tput setab #{col}`
+    code = COLORS[col]
+    raise "Invalid color: #{col}" unless code
+    print "\e[48;5;#{code}m"
+  end
+
+  def reset_colors
+    print @seq_sgr0
+  end
+
+  # Reads keypresses from the user including 2 and 3 escape character sequences.
+  def read_char
+    $stdin.echo = false
+    $stdin.raw!
+
+    input = $stdin.getc.chr
+    if input == "\e" then
+      input << $stdin.read_nonblock(3) rescue nil
+      input << $stdin.read_nonblock(2) rescue nil
+    end
+  ensure
+    $stdin.echo = true
+    $stdin.cooked!
+
+    return input
+  end
+
+  def read_char_mapped
+    c = read_char
+    c2 = KEY_MAP[c]
+    c2 ? c2 : c
+  end
+
+  KEY_MAP = {
+    "\t" => :tab,
+    "\r" => :enter,
+    "\e[A" => :up_arrow,
+    "\e[B" => :down_arrow,
+    "\e[C" => :right_arrow,
+    "\e[D" => :left_arrow,
+    "\e[5~" => :pg_up,
+    "\e[6~" => :pg_dn,
+    "\e[H" => :home,
+    "\e[F" => :end,
+  }
+
+  def message_box_exception(e)
+    message_box("Error while parsing", e.message)
+  end
+
+  SINGLE_CHARSET = '┌┐└┘─│'
+  HEAVY_CHARSET  = '┏┓┗┛━┃'
+  DOUBLE_CHARSET = '╔╗╚╝═║'
+
+  CHAR_TL = 0
+  CHAR_TR = 1
+  CHAR_BL = 2
+  CHAR_BR = 3
+  CHAR_H  = 4
+  CHAR_V  = 5
+
+  def message_box(header, msg)
+    top_y = @rows / 2 - 5
+    draw_rectangle(10, top_y, @cols - 20, 10)
+    goto(@cols / 2 - (header.length / 2) - 1, top_y)
+    print ' ', header, ' '
+    goto(11, top_y + 1)
+    puts msg
+    draw_button(@cols / 2 - 10, top_y + 8, 10, 'OK')
+    loop {
+      c = read_char_mapped
+      return if c == :enter
+    }
+  end
+
+  def draw_rectangle(x, y, w, h, charset = DOUBLE_CHARSET)
+    goto(x, y)
+    print charset[CHAR_TL]
+    print charset[CHAR_H] * (w - 2)
+    print charset[CHAR_TR]
+
+    ((y + 1)..(y + h - 1)).each { |i|
+      goto(x, i)
+      print charset[CHAR_V]
+      print ' ' * (w - 2)
+      print charset[CHAR_V]
+    }
+
+    goto(x, y + h)
+    print charset[CHAR_BL]
+    print charset[CHAR_H] * (w - 2)
+    print charset[CHAR_BR]
+  end
+
+  def draw_button(x, y, w, caption)
+    goto(x, y)
+    puts "[ #{caption} ]"
+  end
+end
+
+end