ScratchABit.py

   1 #!/usr/bin/env python3
   2 # ScratchABit - interactive disassembler
   3 #
   4 # Copyright (c) 2015 Paul Sokolovsky
   5 #
   6 # This program is free software; you can redistribute it and/or modify
   7 # it under the terms of the GNU General Public License as published by
   8 # the Free Software Foundation, either version 3 of the License, or
   9 # (at your option) any later version.
  10 #
  11 # This program is distributed in the hope that it will be useful, but
  12 # WITHOUT ANY WARRANTY; without even the implied warranty of
  13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14 # General Public License for more details.
  15 #
  16 # You should have received a copy of the GNU General Public License
  17 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  18 import sys
  19 import os
  20 import os.path
  21 import time
  22 import re
  23 import string
  24 import binascii
  25 import logging as log
  26 import argparse
  27
  28 from scratchabit import engine
  29 import idaapi
  30
  31 from picotui.widgets import *
  32 from picotui import editorext as editor
  33 from picotui.screen import Screen
  34 from picotui.editorext import Viewer
  35 from picotui.menu import *
  36 from picotui.dialogs import *
  37
  38 from scratchabit import utils
  39 from scratchabit import help
  40 from scratchabit import saveload
  41 from scratchabit import actions
  42 from scratchabit import uiprefs
  43
  44
  45 HEIGHT = 21
  46
  47 MENU_PREFS = 2000
  48 MENU_PLUGIN = 2001
  49 MENU_ADD_TO_FUNC = 2002
  50
  51
  52 class AppClass:
  53     pass
  54
  55 APP = AppClass()
  56
  57
  58 def disasm_one(p):
  59     insnsz = p.ana()
  60     p.out()
  61     print("%08x %s" % (p.cmd.ea, p.cmd.disasm))
  62     p.cmd.ea += p.cmd.size
  63     p.cmd.size = 0
  64
  65
  66 class DisasmViewer(editor.EditorExt):
  67
  68     def __init__(self, *args):
  69         super().__init__(*args)
  70         self.model = None
  71         self.addr_stack = []
  72         self.search_str = ""
  73         self.def_color = C_PAIR(C_CYAN, C_BLUE)
  74
  75     def set_model(self, model):
  76         self.model = model
  77         self.set_lines(model.lines())
  78         # Invalidate top_line. Assuming goto_*() will be called
  79         # after set_model().
  80         self.top_line = sys.maxsize
  81
  82     def show_line(self, l, i):
  83         global show_bytes
  84         res = l
  85         if not isinstance(l, str):
  86             res = "%08x " % l.ea
  87             if show_bytes > 0:
  88                 bin = ""
  89                 if not l.virtual:
  90                     b = self.model.AS.get_bytes(l.ea, l.size)
  91                     bin = str(binascii.hexlify(b[:show_bytes]), "ascii")
  92                     if l.size > show_bytes:
  93                         bin += "+"
  94                 res += idaapi.fillstr(bin, show_bytes * 2 + 1)
  95             res += l.indent + l.render()
  96
  97         COLOR_MAP = {
  98             engine.Label: C_PAIR(C_GREEN, C_BLUE),
  99             engine.AreaWrapper: C_PAIR(C_YELLOW, C_BLUE),
 100             engine.FunctionWrapper: C_PAIR(C_B_YELLOW, C_BLUE),
 101             engine.Xref: C_PAIR(C_MAGENTA, C_BLUE),
 102             engine.Unknown: C_PAIR(C_WHITE, C_BLUE),
 103             engine.Data: C_PAIR(C_MAGENTA, C_BLUE),
 104             engine.String: C_PAIR(C_B_MAGENTA, C_BLUE),
 105             engine.Fill: C_PAIR(C_B_BLUE, C_BLUE),
 106         }
 107         c = COLOR_MAP.get(type(l), self.def_color)
 108         self.attr_color(c)
 109         super().show_line(res, i)
 110         self.attr_reset()
 111
 112
 113     def handle_input(self, key):
 114         try:
 115             return super().handle_input(key)
 116         except Exception as ex:
 117             self.show_exception(ex)
 118             return None
 119
 120
 121     def goto_addr(self, to_addr, col=None, from_addr=None):
 122         if to_addr is None:
 123             self.show_status("No address-like value to go to")
 124             return
 125         subno = -1
 126         if isinstance(to_addr, tuple):
 127             to_addr, subno = to_addr
 128         adj_addr = self.model.AS.adjust_addr_reverse(to_addr)
 129         if adj_addr is None:
 130             self.show_status("Unknown address: 0x%x" % to_addr)
 131             return
 132         to_addr = adj_addr
 133
 134         # If we can position cursor within current screen, do that,
 135         # to avoid jumpy UI
 136         no = self.model.addr2line_no(to_addr, subno)
 137         if no is not None:
 138             if self.line_visible(no):
 139                 self.goto_line(no, col=col)
 140                 if from_addr is not None:
 141                     self.addr_stack.append(from_addr)
 142                 return
 143
 144         # Otherwise, re-render model around needed address, and redraw screen
 145         t = time.time()
 146         model = engine.render_partial_around(to_addr, 0, HEIGHT * 2)
 147         self.show_status("Rendering time: %fs" % (time.time() - t))
 148         if not model:
 149             self.show_status("Unknown address: 0x%x" % to_addr)
 150             return
 151         self.set_model(model)
 152
 153         no = self.model.addr2line_no(to_addr, subno)
 154         if no is not None:
 155             if from_addr is not None:
 156                 self.addr_stack.append(from_addr)
 157             if not self.goto_line(no, col=col):
 158                 # Need to redraw always, because we changed underlying model
 159                 self.redraw()
 160         else:
 161             self.show_status("Unknown address: %x" % to_addr)
 162
 163     def update_model(self, stay_on_real=False):
 164         """Re-render model and update screen in such way that cursor stayed
 165         on the same line (as far as possible).
 166         stay_on_real == False - try to stay on same relative line no. for
 167         the current address.
 168         stay_on_real == True - try to stay on the line which contains real
 169         bytes for the current address (use this if you know that cursor
 170         stayed on such line before the update).
 171         """
 172         addr, subno = self.cur_addr_subno()
 173         t = time.time()
 174         model = engine.render_partial_around(addr, subno, HEIGHT * 2)
 175         self.show_status("Rendering time: %fs" % (time.time() - t))
 176         self.set_model(model)
 177         if stay_on_real:
 178             self.cur_line = model.target_addr_lineno_real
 179         else:
 180             self.cur_line = model.target_addr_lineno
 181         self.top_line = self.cur_line - self.row
 182         #log.debug("update_model: addr=%x, row=%d, cur_line=%d, top_line=%d" % (addr, self.row, self.cur_line, self.top_line))
 183         self.redraw()
 184
 185     def handle_cursor_keys(self, key):
 186         cl = self.cur_line
 187         if super().handle_cursor_keys(key):
 188             if self.cur_line == cl:
 189                 return True
 190             #log.debug("handle_cursor_keys: cur: %d, total: %d", self.cur_line, self.total_lines)
 191             if self.cur_line <= HEIGHT or self.total_lines - self.cur_line <= HEIGHT:
 192                 log.debug("handle_cursor_keys: triggering update")
 193                 self.update_model()
 194
 195             return True
 196         else:
 197             return False
 198
 199     def cur_addr(self):
 200         line = self.get_cur_line()
 201         return line.ea
 202
 203     # Address of the next line. It may be the same address as the
 204     # current line, as several lines may "belong" to the same address,
 205     # (virtual lines like headers, etc.)
 206     def next_line_addr_subno(self):
 207         try:
 208             l = self.content[self.cur_line + 1]
 209             return (l.ea, l.subno)
 210         except:
 211             return None
 212
 213     # Return next address following the current line. May need to skip
 214     # few virtual lines.
 215     def next_addr(self):
 216         addr = self.cur_addr()
 217         n = self.cur_line + 1
 218         try:
 219             while self.content[n].ea == addr:
 220                 n += 1
 221             return self.content[n].ea
 222         except:
 223             return None
 224
 225     def cur_addr_subno(self):
 226         line = self.get_cur_line()
 227         return (line.ea, line.subno)
 228
 229     def cur_operand_no(self, line):
 230         col = self.col - engine.DisasmObj.LEADER_SIZE - len(line.indent)
 231         #self.show_status("Enter pressed: %s, %s" % (col, line))
 232         for i, pos in enumerate(line.arg_pos):
 233             if pos[0] <= col <= pos[1]:
 234                 return i
 235         return -1
 236
 237     def analyze_status(self, cnt):
 238         self.show_status("Analyzing (%d insts so far)" % cnt)
 239
 240     def expect_flags(self, fl, allowed_flags):
 241         if fl not in allowed_flags:
 242             self.show_status("Undefine first (u key)")
 243             return False
 244         return True
 245
 246
 247     def show_exception(self, e):
 248         log.exception("Exception processing user command")
 249         L = 5
 250         T = 2
 251         W = 70
 252         H = 20
 253         self.dialog_box(L, T, W, H)
 254         v = Viewer(L + 1, T + 1, W - 2, H - 2)
 255         import traceback
 256         v.set_lines([
 257             "Exception occured processing the command. Press Esc to continue.",
 258             "Recommended action is saving database, quitting and comparing",
 259             "database files with backup copies for possibility of data loss",
 260             "or corruption. The exception was also logged to scratchabit.log.",
 261             "Please report way to reproduce it to",
 262             "https://github.com/pfalcon/ScratchABit/issues",
 263             "",
 264         ] + traceback.format_exc().splitlines())
 265         v.loop()
 266         self.redraw()
 267
 268
 269     def resolve_expr(self, expr):
 270         if expr:
 271             if expr[0].isdigit():
 272                 return int(expr, 0)
 273             else:
 274                 words = expr.split("+", 1)
 275                 addend = 0
 276                 if len(words) > 1:
 277                     try:
 278                         addend = int(words[1], 0)
 279                     except:
 280                         pass
 281                 to_addr = self.model.AS.resolve_label(words[0])
 282                 if to_addr is None:
 283                     return
 284                 return to_addr + addend
 285
 286
 287     def require_non_func(self, fl):
 288         if fl & 0x7f != self.model.AS.CODE:
 289             self.show_status("Code required")
 290             return False
 291         if fl & self.model.AS.FUNC:
 292             self.show_status("Already a function")
 293             return False
 294         return True
 295
 296
 297     def handle_edit_key(self, key):
 298         line = self.get_cur_line()
 299         if key == editor.KEY_ENTER:
 300             line = self.get_cur_line()
 301             log.info("Enter pressed: %s" % line)
 302             op_no = self.cur_operand_no(line)
 303             self.show_status("Enter pressed: %s, %s" % (self.col, op_no))
 304             to_addr = None
 305             # No longer try to jump only to addresses in args, parse
 306             # textual representation below
 307             if False and isinstance(line, engine.DisasmObj):
 308                 if op_no >= 0:
 309                     o = line[op_no]
 310                     to_addr = o.get_addr()
 311                 if to_addr is None:
 312                     o = line.get_operand_addr()
 313                     if o:
 314                         to_addr = o.get_addr()
 315             if to_addr is None:
 316                 pos = self.col - line.LEADER_SIZE - len(line.indent)
 317                 word = utils.get_word_at_pos(line.cache, pos)
 318                 self.show_status("Enter pressed: %s, %s, %s" % (self.col, op_no, word))
 319                 to_addr = self.resolve_expr(word)
 320                 if to_addr is None:
 321                     self.show_status("Unknown address: %s" % word)
 322                     return
 323             self.goto_addr(to_addr, from_addr=self.cur_addr_subno())
 324         elif key == editor.KEY_ESC:
 325             if self.addr_stack:
 326                 self.show_status("Returning")
 327                 self.goto_addr(self.addr_stack.pop())
 328         elif key == b"q":
 329             res = ACTION_OK
 330             if self.model.AS.changed:
 331                 res = DConfirmation("There're unsaved changes. Quit?").result()
 332             if res == ACTION_OK:
 333                 return editor.KEY_QUIT
 334             self.redraw()
 335         elif key == b"\x1b[5;5~":  # Ctrl+PgUp
 336             self.goto_addr(self.model.AS.min_addr(), from_addr=line.ea)
 337         elif key == b"\x1b[6;5~":  # Ctrl+PgDn
 338             self.goto_addr(self.model.AS.max_addr(), from_addr=line.ea)
 339         elif key == b"c":
 340             addr = self.cur_addr()
 341             self.show_status("Analyzing at %x" % addr)
 342             engine.add_entrypoint(addr, False)
 343             engine.analyze(self.analyze_status)
 344             self.update_model()
 345
 346         elif key == b"F":
 347             addr = self.cur_addr()
 348             fl = self.model.AS.get_flags(addr, 0xff)
 349             if not self.require_non_func(fl):
 350                 return
 351             self.show_status("Retracing as a function...")
 352             self.model.AS.make_label("fun_", addr)
 353             engine.add_entrypoint(addr, True)
 354             engine.analyze(self.analyze_status)
 355             self.update_model()
 356             self.show_status("Retraced as a function")
 357
 358         elif key == MENU_ADD_TO_FUNC:
 359             addr = self.cur_addr()
 360             if actions.add_code_to_func(APP, addr):
 361                 self.update_model()
 362
 363         elif key == b"d":
 364             addr = self.cur_addr()
 365             fl = self.model.AS.get_flags(addr)
 366             if not self.expect_flags(fl, (self.model.AS.DATA, self.model.AS.UNK)):
 367                 return
 368             if fl == self.model.AS.UNK:
 369                 self.model.AS.set_flags(addr, 1, self.model.AS.DATA, self.model.AS.DATA_CONT)
 370             else:
 371                 sz = self.model.AS.get_unit_size(addr)
 372                 self.model.undefine_unit(addr)
 373                 sz *= 2
 374                 if sz > 4: sz = 1
 375                 self.model.AS.set_flags(addr, sz, self.model.AS.DATA, self.model.AS.DATA_CONT)
 376             self.update_model()
 377         elif key == b"a":
 378             addr = self.cur_addr()
 379             fl = self.model.AS.get_flags(addr)
 380             if not self.expect_flags(fl, (self.model.AS.DATA, self.model.AS.UNK)):
 381                 return
 382             sz = 0
 383             label = "s_"
 384             while True:
 385                 b = self.model.AS.get_byte(addr)
 386                 fl = self.model.AS.get_flags(addr)
 387                 if not (0x20 <= b <= 0x7e or b in (0x0a, 0x0d)):
 388                     if b == 0:
 389                         sz += 1
 390                     break
 391                 if fl not in (self.model.AS.UNK, self.model.AS.DATA, self.model.AS.DATA_CONT):
 392                     break
 393                 c = chr(b)
 394                 if c < '0' or c in string.punctuation:
 395                     c = '_'
 396                 label += c
 397                 addr += 1
 398                 sz += 1
 399             if sz > 0:
 400                 self.model.AS.set_flags(self.cur_addr(), sz, self.model.AS.STR, self.model.AS.DATA_CONT)
 401                 self.model.AS.make_unique_label(self.cur_addr(), label)
 402                 self.update_model()
 403         elif key == b"f":
 404             addr = self.cur_addr()
 405             fl = self.model.AS.get_flags(addr)
 406             if not self.expect_flags(fl, (self.model.AS.UNK,)):
 407                 return
 408
 409             off, area = self.model.AS.addr2area(self.cur_addr())
 410             # Don't cross area boundaries with filler
 411             remaining = area[engine.END] - addr + 1
 412             sz = 0
 413             while remaining:
 414                 try:
 415                     fl = self.model.AS.get_flags(addr)
 416                 except engine.InvalidAddrException:
 417                     break
 418                 if fl != self.model.AS.UNK:
 419                     break
 420                 b = self.model.AS.get_byte(addr)
 421                 if b not in (0, 0xff):
 422                     self.show_status("Filler must consist of 0x00 or 0xff")
 423                     return
 424                 sz += 1
 425                 addr += 1
 426                 remaining -= 1
 427             if sz > 0:
 428                 self.model.AS.make_filler(self.cur_addr(), sz)
 429                 self.update_model()
 430
 431         elif key == b"u":
 432             addr = self.cur_addr()
 433             self.model.undefine_unit(addr)
 434             self.update_model()
 435
 436         elif key == b"h":
 437             op_no = self.cur_operand_no(self.get_cur_line())
 438             if op_no >= 0:
 439                 addr = self.cur_addr()
 440                 subtype = self.model.AS.get_arg_prop(addr, op_no, "subtype")
 441                 if subtype != engine.IMM_ADDR:
 442                     next_subtype = {
 443                         engine.IMM_UHEX: engine.IMM_UDEC,
 444                         engine.IMM_UDEC: engine.IMM_UHEX,
 445                     }
 446                     self.model.AS.set_arg_prop(addr, op_no, "subtype", next_subtype[subtype])
 447                     self.redraw()
 448                     self.show_status("Changed arg #%d to %s" % (op_no, next_subtype[subtype]))
 449         elif key == b"o":
 450             addr = self.cur_addr()
 451             line = self.get_cur_line()
 452             o = line.get_operand_addr()
 453             if not o:
 454                 self.show_status("Cannot convert operand to offset")
 455                 return
 456             if o.type != idaapi.o_imm or not self.model.AS.is_valid_addr(o.get_addr()):
 457                 self.show_status("Cannot convert operand to offset: #%s: %s" % (o.n, o.type))
 458                 return
 459
 460             if self.model.AS.get_arg_prop(addr, o.n, "subtype") == engine.IMM_ADDR:
 461                 self.model.AS.unmake_arg_offset(addr, o.n, o.get_addr())
 462             else:
 463                 self.model.AS.make_arg_offset(addr, o.n, o.get_addr())
 464             self.update_model(True)
 465         elif key == b";":
 466             addr = self.cur_addr()
 467             comment = self.model.AS.get_comment(addr) or ""
 468             res = DMultiEntry(60, 5, comment.split("\n"), title="Comment:").result()
 469             if res != ACTION_CANCEL:
 470                 res = "\n".join(res).rstrip("\n")
 471                 self.model.AS.set_comment(addr, res)
 472                 self.update_model()
 473             else:
 474                 self.redraw()
 475         elif key == b"n":
 476             addr = self.cur_addr()
 477             label = self.model.AS.get_label(addr)
 478             def_label = self.model.AS.get_default_label(addr)
 479             s = label or def_label
 480             while True:
 481                 res = DTextEntry(30, s, title="New label:").result()
 482                 if not res:
 483                     break
 484                 if res == def_label:
 485                     res = addr
 486                 else:
 487                     if self.model.AS.label_exists(res):
 488                         s = res
 489                         self.show_status("Duplicate label")
 490                         continue
 491                 self.model.AS.set_label(addr, res)
 492                 if not label:
 493                     # If it's new label, we need to add it to model
 494                     self.update_model()
 495                     return
 496                 break
 497             self.redraw()
 498         elif key == b"g":
 499             d = Dialog(4, 4, title="Go to")
 500             d.add(1, 1, WLabel("Label/addr:"))
 501             entry = WAutoComplete(20, "", self.model.AS.get_label_list())
 502             entry.popup_h = 12
 503             entry.finish_dialog = ACTION_OK
 504             d.add(13, 1, entry)
 505             d.add(1, 2, WLabel("Press Down to auto-complete"))
 506             res = d.loop()
 507             self.redraw()
 508
 509             if res == ACTION_OK:
 510                 value = entry.get_text()
 511                 if '0' <= value[0] <= '9':
 512                     addr = int(value, 0)
 513                 else:
 514                     addr = self.model.AS.resolve_label(value)
 515                 self.goto_addr(addr, from_addr=self.cur_addr())
 516
 517         elif key == editor.KEY_F1:
 518             help.help(self)
 519             self.redraw()
 520         elif key == b"S":
 521             self.show_status("Saving...")
 522             saveload.save_state(project_dir)
 523             self.model.AS.changed = False
 524             self.show_status("Saved.")
 525         elif key == b"\x11":  # ^Q
 526             class IssueList(WListBox):
 527                 def render_line(self, l):
 528                     return "%08x %s" % l
 529             d = Dialog(4, 4, title="Problems list")
 530             lw = IssueList(40, 16, self.model.AS.get_issues())
 531             lw.finish_dialog = ACTION_OK
 532             d.add(1, 1, lw)
 533             res = d.loop()
 534             self.redraw()
 535             if res == ACTION_OK:
 536                 val = lw.get_cur_line()
 537                 if val:
 538                     self.goto_addr(val[0], from_addr=self.cur_addr())
 539
 540         elif key == b"i":
 541             off, area = self.model.AS.addr2area(self.cur_addr())
 542             props = area[engine.PROPS]
 543             percent = 100 * off / (area[engine.END] - area[engine.START] + 1)
 544             func = self.model.AS.lookup_func(self.cur_addr())
 545             func = self.model.AS.get_label(func.start) if func else None
 546             status = "Area: 0x%x %s (%s): %.1f%%, func: %s" % (
 547                 area[engine.START], props.get("name", "noname"), props["access"], percent, func
 548             )
 549             subarea = self.model.AS.lookup_subarea(self.cur_addr())
 550             if subarea:
 551                 status += ", subarea: " + subarea[2]
 552             self.show_status(status)
 553         elif key == b"I":
 554             from scratchabit import memmap
 555             addr = memmap.show(self.model.AS, self.cur_addr())
 556             if addr is not None:
 557                 self.goto_addr(addr, from_addr=self.cur_addr())
 558             self.redraw()
 559         elif key == b"W":
 560             out_fname = "out.lst"
 561             with open(out_fname, "w") as f:
 562                 engine.render_partial(actions.TextSaveModel(f, self), 0, 0, 10000000)
 563             self.show_status("Disassembly listing written: " + out_fname)
 564         elif key == b"\x17":  # Ctrl+W
 565             outfile = actions.write_func_by_addr(APP, self.cur_addr(), feedback_obj=self)
 566             if outfile:
 567                 self.show_status("Wrote file: %s" % outfile)
 568         elif key == b"\x15":  # Ctrl+U
 569             # Next undefined
 570             addr = self.cur_addr()
 571             flags = self.model.AS.get_flags(addr)
 572             if flags == self.model.AS.UNK:
 573                 # If already on undefined, skip the current stride of them,
 574                 # as they indeed go in batches.
 575                 while True:
 576                     flags = self.model.AS.get_flags(addr)
 577                     if flags != self.model.AS.UNK:
 578                         break
 579                     addr = self.model.AS.next_addr(addr)
 580                     if addr is None:
 581                         break
 582
 583             if addr is not None:
 584                 while True:
 585                     flags = self.model.AS.get_flags(addr)
 586                     if flags == self.model.AS.UNK:
 587                         self.goto_addr(addr, from_addr=self.cur_addr())
 588                         break
 589                     addr = self.model.AS.next_addr(addr)
 590                     if addr is None:
 591                         break
 592
 593             if addr is None:
 594                 self.show_status("There're no further undefined strides")
 595
 596         elif key == b"\x06":  # Ctrl+F
 597             # Next non-function
 598             addr = self.cur_addr()
 599             flags = self.model.AS.get_flags(addr, 0xff)
 600             if flags == self.model.AS.CODE:
 601                 # If already on non-func code, skip the current stride of it,
 602                 # as it indeed go in batches.
 603                 while True:
 604                     flags = self.model.AS.get_flags(addr, 0xff)
 605                     self.show_status("fl=%x" % flags)
 606                     if flags not in (self.model.AS.CODE, self.model.AS.CODE_CONT):
 607                         break
 608                     addr = self.model.AS.next_addr(addr)
 609                     if addr is None:
 610                         break
 611
 612             if addr is not None:
 613                 while True:
 614                     flags = self.model.AS.get_flags(addr, 0xff)
 615                     if flags == self.model.AS.CODE:
 616                         self.goto_addr(addr, from_addr=self.cur_addr())
 617                         break
 618                     addr = self.model.AS.next_addr(addr)
 619                     if addr is None:
 620                         break
 621
 622             if addr is None:
 623                 self.show_status("There're no further non-function code strides")
 624
 625         elif key in (b"/", b"?"):  # "/" and Shift+"/"
 626
 627             class FoundException(Exception): pass
 628
 629             class TextSearchModel(engine.Model):
 630                 def __init__(self, substr, ctrl, this_addr, this_subno):
 631                     super().__init__()
 632                     self.search = substr
 633                     self.ctrl = ctrl
 634                     self.this_addr = this_addr
 635                     self.this_subno = this_subno
 636                     self.cnt = 0
 637                 def add_line(self, addr, line):
 638                     super().add_line(addr, line)
 639                     # Skip virtual lines before the line from which we started
 640                     if addr == self.this_addr and line.subno < self.this_subno:
 641                         return
 642                     txt = line.render()
 643                     idx = txt.find(self.search)
 644                     if idx != -1:
 645                         raise FoundException((addr, line.subno), idx + line.LEADER_SIZE + len(line.indent))
 646                     if self.cnt % 256 == 0:
 647                         self.ctrl.show_status("Searching: 0x%x" % addr)
 648                     self.cnt += 1
 649                     # Don't accumulate lines
 650                     self._lines = []
 651                     self._addr2line = {}
 652
 653             if key == b"/":
 654                 d = Dialog(4, 4, title="Text Search")
 655                 d.add(1, 1, WLabel("Search for:"))
 656                 entry = WTextEntry(20, self.search_str)
 657                 entry.finish_dialog = ACTION_OK
 658                 d.add(13, 1, entry)
 659                 res = d.loop()
 660                 self.redraw()
 661                 self.search_str = entry.get_text()
 662                 if res != ACTION_OK or not self.search_str:
 663                     return
 664                 addr, subno = self.cur_addr_subno()
 665             else:
 666                 addr, subno = self.next_line_addr_subno()
 667
 668             try:
 669                 engine.render_from(TextSearchModel(self.search_str, self, addr, subno), addr, 10000000)
 670             except FoundException as res:
 671                 self.goto_addr(res.args[0], col=res.args[1], from_addr=self.cur_addr())
 672             else:
 673                 self.show_status("Not found: " + self.search_str)
 674
 675         elif key == MENU_PREFS:
 676             uiprefs.handle(APP)
 677
 678         elif key == MENU_PLUGIN:
 679             res = DTextEntry(30, "", title="Plugin module name:").result()
 680             self.redraw()
 681             if res:
 682                 self.show_status("Running '%s' plugin..." % res)
 683                 call_script(res)
 684                 self.update_model()
 685                 self.show_status("Plugin '%s' ran successfully" % res)
 686         else:
 687             self.show_status("Unbound key: " + repr(key))
 688
 689
 690 CPU_PLUGIN = None
 691 ENTRYPOINTS = []
 692 show_bytes = 0
 693
 694 def filter_config_line(l):
 695     l = re.sub(r"#.*$", "", l)
 696     l = l.strip()
 697     return l
 698
 699 def load_symbols(fname):
 700     with open(fname) as f:
 701         for l in f:
 702             l = filter_config_line(l)
 703             if not l:
 704                 continue
 705             m = re.search(r"\b([A-Za-z_$.][A-Za-z0-9_$.]*)\s*=\s*((0x)?[0-9A-Fa-f]+)", l)
 706             if m:
 707                 #print(m.groups())
 708                 ENTRYPOINTS.append((m.group(1), int(m.group(2), 0)))
 709             else:
 710                 print("Warning: cannot parse entrypoint info from: %r" % l)
 711
 712
 713 # Allow undescores to separate digit groups
 714 def str2int(s):
 715     return int(s.replace("_", ""), 0)
 716
 717
 718 def parse_range(arg):
 719     # name start(len)
 720     # name start-end
 721     if "(" in arg:
 722         m = re.match(r"(.+?)\s*\(\s*(.+?)\s*\)", arg)
 723         start = str2int(m.group(1))
 724         end = start + str2int(m.group(2)) - 1
 725     else:
 726         m = re.match(r"(.+)\s*-\s*(.+)", arg)
 727         start = str2int(m.group(1))
 728         end = str2int(m.group(2))
 729     return start, end
 730
 731
 732 def parse_entrypoints(f):
 733     for l in f:
 734         l = filter_config_line(l)
 735         if not l:
 736             continue
 737         if l[0] == "[":
 738             return l
 739         m = re.match(r'load "(.+?)"', l)
 740         if m:
 741             load_symbols(m.group(1))
 742         else:
 743             label, addr = [v.strip() for v in l.split("=")]
 744             ENTRYPOINTS.append((label, int(addr, 0)))
 745     return ""
 746
 747 def parse_subareas(f):
 748     subareas = []
 749     for l in f:
 750         l = filter_config_line(l)
 751         if not l:
 752             continue
 753         if l[0] == "[":
 754             return l
 755
 756         args = l.split()
 757         assert len(args) == 2
 758         start, end = parse_range(args[1])
 759         engine.ADDRESS_SPACE.add_subarea(start, end, args[0])
 760     engine.ADDRESS_SPACE.finish_subareas()
 761     return ""
 762
 763
 764 def load_target_file(loader, fname):
 765     entry = loader.load(engine.ADDRESS_SPACE, fname)
 766     log.info("Loaded %s, entrypoint: %s", fname, hex(entry) if entry is not None else None)
 767     if entry is not None:
 768         ENTRYPOINTS.append(("_ENTRY_", entry))
 769
 770
 771 def parse_disasm_def(fname):
 772     global CPU_PLUGIN
 773     global show_bytes
 774     with open(fname) as f:
 775         for l in f:
 776             l = filter_config_line(l)
 777             if not l:
 778                 continue
 779             #print(l)
 780             while True:
 781                 if not l:
 782                     #return
 783                     break
 784                 if l[0] == "[":
 785                     section = l[1:-1]
 786                     print("Processing section: %s" % section)
 787                     if section == "entrypoints":
 788                         l = parse_entrypoints(f)
 789                     elif section == "subareas":
 790                         l = parse_subareas(f)
 791                     else:
 792                         assert 0, "Unknown section: " + section
 793                 else:
 794                     break
 795
 796             if not l:
 797                 break
 798
 799             if l.startswith("load"):
 800                 args = l.split()
 801                 if args[2][0] in string.digits:
 802                     addr = int(args[2], 0)
 803                     print("Loading %s @0x%x" % (args[1], addr))
 804                     engine.ADDRESS_SPACE.load_content(open(args[1], "rb"), addr)
 805                 else:
 806                     print("Loading %s (%s plugin)" % (args[1], args[2]))
 807                     loader = __import__(args[2])
 808                     load_target_file(loader, args[1])
 809             elif l.startswith("cpu "):
 810                 args = l.split()
 811                 CPU_PLUGIN = __import__(args[1])
 812                 print("Loading CPU plugin %s" % (args[1]))
 813             elif l.startswith("show bytes "):
 814                 args = l.split()
 815                 show_bytes = int(args[2])
 816             elif l.startswith("area "):
 817                 args = l.split()
 818                 assert len(args) == 4
 819                 start, end = parse_range(args[2])
 820                 a = engine.ADDRESS_SPACE.add_area(start, end, {"name": args[1], "access": args[3].upper()})
 821                 print("Adding area: %s" % engine.str_area(a))
 822             else:
 823                 assert 0, "Unknown directive: " + l
 824
 825
 826 class MainScreen:
 827
 828     def __init__(self):
 829         self.screen_size = Screen.screen_size()
 830         self.e = DisasmViewer(1, 2, self.screen_size[0] - 2, self.screen_size[1] - 4)
 831
 832         menu_file = WMenuBox([
 833             ("Save (Shift+s)", b"S"), ("Write disasm (Shift+w)", b"W"),
 834             ("Write function (Ctrl+w)", b"\x17"),
 835             ("Quit (q)", b"q")
 836         ])
 837         menu_goto = WMenuBox([
 838             ("Follow (Enter)", KEY_ENTER), ("Return (Esc)", KEY_ESC),
 839             ("Goto... (g)", b"g"), ("Search disasm... (/)", b"/"),
 840             ("Search next (Shift+/)", b"?"), ("Next undefined (Ctrl+u)", b"\x15"),
 841             ("Next non-function code (Ctrl+f)", b"\x06"),
 842         ])
 843         menu_edit = WMenuBox([
 844             ("Undefined (u)", b"u"), ("Code (c)", b"c"), ("Data (d)", b"d"),
 845             ("ASCII String (a)", b"a"), ("Filler (f)", b"f"), ("Make label (n)", b"n"),
 846             ("Mark function start (F)", b"F"), ("Add code to function", MENU_ADD_TO_FUNC),
 847             ("Number/Address (o)", b"o"), ("Hex/dec (h)", b"h"),
 848         ])
 849         menu_analysis = WMenuBox([
 850             ("Info (whereami) (i)", b"i"), ("Memory map (Shift+i)", b"I"),
 851             ("Run plugin...", MENU_PLUGIN),
 852             ("Preferences...", MENU_PREFS),
 853         ])
 854         menu_help = WMenuBox([
 855             ("Help (F1)", KEY_F1), ("About...", "about"),
 856         ])
 857         self.menu_bar = WMenuBar([
 858             ("File", menu_file), ("Goto", menu_goto), ("Edit", menu_edit),
 859             ("Analysis", menu_analysis), ("Help", menu_help)
 860         ])
 861         self.menu_bar.permanent = True
 862
 863     def redraw(self, allow_cursor=True):
 864         self.menu_bar.redraw()
 865         self.e.attr_color(C_B_WHITE, C_BLUE)
 866         self.e.draw_box(0, 1, self.screen_size[0], self.screen_size[1] - 2)
 867         self.e.attr_reset()
 868         self.e.redraw()
 869         if allow_cursor:
 870             self.e.cursor(True)
 871
 872     def loop(self):
 873         while 1:
 874             key = self.e.get_input()
 875             if isinstance(key, list):
 876                 x, y = key
 877                 if self.menu_bar.inside(x, y):
 878                     self.menu_bar.focus = True
 879
 880             if self.menu_bar.focus:
 881                 res = self.menu_bar.handle_input(key)
 882                 if res == ACTION_CANCEL:
 883                     self.menu_bar.focus = False
 884                 elif res is not None and res is not True:
 885
 886                     res = self.e.handle_input(res)
 887                     if res is not None and res is not True:
 888                         return res
 889             else:
 890                 if key == KEY_F9:
 891                     self.menu_bar.focus = True
 892                     self.menu_bar.redraw()
 893                     continue
 894
 895                 res = self.e.handle_input(key)
 896
 897                 if res is not None and res is not True:
 898                     return res
 899
 900
 901 def call_script(script):
 902     mod = __import__(script)
 903     main_f = getattr(mod, "main", None)
 904     if main_f:
 905         main_f(APP)
 906
 907
 908 if __name__ == "__main__":
 909
 910     argp = argparse.ArgumentParser(description="ScratchABit interactive disassembler")
 911     argp.add_argument("file", help="Input file (binary or disassembly .def)")
 912     argp.add_argument("--script", action="append", help="Run script from file after loading environment")
 913     argp.add_argument("--save", action="store_true", help="Save after --script and quit; don't show UI")
 914     args = argp.parse_args()
 915
 916     # Plugin dirs are relative to the dir where scratchabit.py resides.
 917     # sys.path[0] below provide absolute path of this dir, resolved for
 918     # symlinks.
 919     plugin_dirs = ["plugins", "plugins/cpu", "plugins/loader"]
 920     for d in plugin_dirs:
 921         sys.path.append(os.path.join(sys.path[0], d))
 922     log.basicConfig(filename="scratchabit.log", format='%(asctime)s %(message)s', level=log.DEBUG)
 923     log.info("Started")
 924
 925     if args.file.endswith(".def"):
 926         parse_disasm_def(args.file)
 927         project_name = args.file.rsplit(".", 1)[0]
 928     else:
 929         import default_plugins
 930         for loader_id in default_plugins.loaders:
 931             loader = __import__(loader_id)
 932             arch_id = loader.detect(args.file)
 933             if arch_id:
 934                 break
 935         if not arch_id:
 936             print("Error: file '%s' not recognized by default loaders" % args.file)
 937             sys.exit(1)
 938         if arch_id not in default_plugins.cpus:
 939             print("Error: no plugin for CPU '%s' as detected for file '%s'" % (arch_id, args.file))
 940             sys.exit(1)
 941         load_target_file(loader, args.file)
 942         CPU_PLUGIN = __import__(default_plugins.cpus[arch_id])
 943         project_name = args.file
 944
 945     p = CPU_PLUGIN.PROCESSOR_ENTRY()
 946     if hasattr(p, "config"):
 947         p.config()
 948     engine.set_processor(p)
 949     if hasattr(p, "help_text"):
 950         help.set_cpu_help(p.help_text)
 951     APP.cpu_plugin = p
 952     APP.aspace = engine.ADDRESS_SPACE
 953     APP.is_ui = False
 954     engine.ADDRESS_SPACE.is_loading = True
 955
 956     engine.DisasmObj.LEADER_SIZE = 8 + 1
 957     if show_bytes:
 958         engine.DisasmObj.LEADER_SIZE += show_bytes * 2 + 1
 959
 960     # Strip suffix if any from def filename
 961     project_dir = project_name + ".scratchabit"
 962
 963     if saveload.save_exists(project_dir):
 964         saveload.load_state(project_dir)
 965     else:
 966         for label, addr in ENTRYPOINTS:
 967             if engine.ADDRESS_SPACE.is_exec(addr):
 968                 engine.add_entrypoint(addr)
 969             engine.ADDRESS_SPACE.make_unique_label(addr, label)
 970         def _progress(cnt):
 971             sys.stdout.write("Performing initial analysis... %d\r" % cnt)
 972         engine.analyze(_progress)
 973         print()
 974
 975     #engine.print_address_map()
 976
 977     if args.script:
 978         for script in args.script:
 979             call_script(script)
 980         if args.save:
 981             saveload.save_state(project_dir)
 982             sys.exit()
 983
 984     addr_stack = []
 985     if os.path.exists(project_dir + "/session.addr_stack"):
 986         addr_stack = saveload.load_addr_stack(project_dir)
 987         print(addr_stack)
 988         show_addr = addr_stack.pop()
 989     else:
 990         if ENTRYPOINTS:
 991             show_addr = ENTRYPOINTS[0][1]
 992         else:
 993             show_addr = engine.ADDRESS_SPACE.min_addr()
 994
 995     t = time.time()
 996     #_model = engine.render()
 997     _model = engine.render_partial_around(show_addr, 0, HEIGHT * 2)
 998     print("Rendering time: %fs" % (time.time() - t))
 999     #print(_model.lines())
1000     #sys.exit()
1001
1002     engine.ADDRESS_SPACE.is_loading = False
1003     engine.ADDRESS_SPACE.changed = False
1004     Screen.init_tty()
1005     try:
1006         Screen.cls()
1007         Screen.enable_mouse()
1008         main_screen = MainScreen()
1009         APP.main_screen = main_screen
1010         APP.is_ui = True
1011         main_screen.e.set_model(_model)
1012         main_screen.e.addr_stack = addr_stack
1013         main_screen.e.goto_addr(show_addr)
1014         Screen.set_screen_redraw(main_screen.redraw)
1015         main_screen.redraw()
1016         main_screen.e.show_status("Press F1 for help, F9 for menus")
1017         main_screen.loop()
1018     except:
1019         log.exception("Unhandled exception")
1020         raise
1021     finally:
1022         Screen.goto(0, main_screen.screen_size[1])
1023         Screen.cursor(True)
1024         Screen.disable_mouse()
1025         Screen.deinit_tty()
1026         Screen.wr("\n\n")
1027         saveload.save_session(project_dir, main_screen.e)