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