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"C":
 409             addr = self.cur_addr()
 410             self.show_status("Analyzing at %x" % addr)
 411             self.model.AS.make_alt_code(addr)
 412             engine.add_entrypoint(addr, False)
 413             engine.analyze(self.analyze_status)
 414             self.update_model()
 415
 416         elif key == b"F":
 417             addr = self.cur_addr()
 418             fl = self.model.AS.get_flags(addr, 0xff)
 419             if not self.require_non_func(fl):
 420                 return
 421             self.show_status("Retracing as a function...")
 422             self.model.AS.make_label("fun_", addr)
 423             engine.add_entrypoint(addr, True)
 424             engine.analyze(self.analyze_status)
 425             self.update_model()
 426             self.show_status("Retraced as a function")
 427
 428         elif key == MENU_ADD_TO_FUNC:
 429             addr = self.cur_addr()
 430             if actions.add_code_to_func(APP, addr):
 431                 self.update_model()
 432
 433         elif key == b"d":
 434             addr = self.cur_addr()
 435             fl = self.model.AS.get_flags(addr)
 436             if not self.expect_flags(fl, (self.model.AS.DATA, self.model.AS.UNK)):
 437                 return
 438             if fl == self.model.AS.UNK:
 439                 self.model.AS.set_flags(addr, 1, self.model.AS.DATA, self.model.AS.DATA_CONT)
 440             else:
 441                 sz = self.model.AS.get_unit_size(addr)
 442                 self.model.undefine_unit(addr)
 443                 sz *= 2
 444                 if sz > 4: sz = 1
 445                 self.model.AS.set_flags(addr, sz, self.model.AS.DATA, self.model.AS.DATA_CONT)
 446             self.update_model()
 447         elif key == b"f":
 448             addr = self.cur_addr()
 449             fl = self.model.AS.get_flags(addr)
 450             if not self.expect_flags(fl, (self.model.AS.UNK,)):
 451                 return
 452
 453             off, area = self.model.AS.addr2area(self.cur_addr())
 454             # Don't cross area boundaries with filler
 455             remaining = area[engine.END] - addr + 1
 456             sz = 0
 457             while remaining:
 458                 try:
 459                     fl = self.model.AS.get_flags(addr)
 460                 except engine.InvalidAddrException:
 461                     break
 462                 if fl != self.model.AS.UNK:
 463                     break
 464                 b = self.model.AS.get_byte(addr)
 465                 if b not in (0, 0xff):
 466                     self.show_status("Filler must consist of 0x00 or 0xff")
 467                     return
 468                 sz += 1
 469                 addr += 1
 470                 remaining -= 1
 471             if sz > 0:
 472                 self.model.AS.make_filler(self.cur_addr(), sz)
 473                 self.update_model()
 474
 475         elif key == b"u":
 476             addr = self.cur_addr()
 477             self.model.undefine_unit(addr)
 478             self.update_model()
 479
 480         elif key == b"h":
 481             op_no = self.cur_operand_no(self.get_cur_line())
 482             if op_no >= 0:
 483                 addr = self.cur_addr()
 484                 subtype = self.model.AS.get_arg_prop(addr, op_no, "subtype")
 485                 if subtype != engine.IMM_ADDR:
 486                     next_subtype = {
 487                         engine.IMM_UHEX: engine.IMM_UDEC,
 488                         engine.IMM_UDEC: engine.IMM_UHEX,
 489                     }
 490                     self.model.AS.set_arg_prop(addr, op_no, "subtype", next_subtype[subtype])
 491                     self.redraw()
 492                     self.show_status("Changed arg #%d to %s" % (op_no, next_subtype[subtype]))
 493         elif key == b"o":
 494             addr = self.cur_addr()
 495             line = self.get_cur_line()
 496             o = line.get_operand_addr()
 497             if not o:
 498                 self.show_status("Cannot convert operand to offset")
 499                 return
 500             if o.type != idaapi.o_imm or not self.model.AS.is_valid_addr(o.get_addr()):
 501                 self.show_status("Cannot convert operand to offset: #%s: %s" % (o.n, o.type))
 502                 return
 503
 504             if self.model.AS.get_arg_prop(addr, o.n, "subtype") == engine.IMM_ADDR:
 505                 self.model.AS.unmake_arg_offset(addr, o.n, o.get_addr())
 506             else:
 507                 self.model.AS.make_arg_offset(addr, o.n, o.get_addr())
 508             self.update_model(True)
 509         elif key == b";":
 510             addr = self.cur_addr()
 511             comment = self.model.AS.get_comment(addr) or ""
 512             res = DMultiEntry(60, 5, comment.split("\n"), title="Comment:").result()
 513             if res != ACTION_CANCEL:
 514                 res = "\n".join(res).rstrip("\n")
 515                 self.model.AS.set_comment(addr, res)
 516                 self.update_model()
 517             else:
 518                 self.redraw()
 519         elif key == b"n":
 520             addr = self.cur_addr()
 521             label = self.model.AS.get_label(addr)
 522             def_label = self.model.AS.get_default_label(addr)
 523             s = label or def_label
 524             while True:
 525                 res = DTextEntry(30, s, title="New label:").result()
 526                 if not res:
 527                     break
 528                 if res == def_label:
 529                     res = addr
 530                 else:
 531                     if self.model.AS.label_exists(res):
 532                         s = res
 533                         self.show_status("Duplicate label")
 534                         continue
 535                 self.model.AS.set_label(addr, res)
 536                 if not label:
 537                     # If it's new label, we need to add it to model
 538                     self.update_model()
 539                     return
 540                 break
 541             self.redraw()
 542
 543         elif key == editor.KEY_F1:
 544             help.help(self)
 545             self.redraw()
 546         elif key == b"S":
 547             self.show_status("Saving...")
 548             saveload.save_state(project_dir)
 549             self.model.AS.changed = False
 550             self.show_status("Saved.")
 551         elif key == b"\x11":  # ^Q
 552             class IssueList(WListBox):
 553                 def render_line(self, l):
 554                     return "%08x %s" % l
 555             d = Dialog(4, 4, title="Problems list")
 556             lw = IssueList(40, 16, self.model.AS.get_issues())
 557             lw.finish_dialog = ACTION_OK
 558             d.add(1, 1, lw)
 559             res = d.loop()
 560             self.redraw()
 561             if res == ACTION_OK:
 562                 val = lw.get_cur_line()
 563                 if val:
 564                     self.goto_addr(val[0], from_addr=self.cur_addr())
 565
 566         elif key == b"i":
 567             off, area = self.model.AS.addr2area(self.cur_addr())
 568             props = area[engine.PROPS]
 569             percent = 100 * off / (area[engine.END] - area[engine.START] + 1)
 570             func = self.model.AS.lookup_func(self.cur_addr())
 571             func = self.model.AS.get_label(func.start) if func else None
 572             status = "Area: 0x%x %s (%s): %.1f%%, func: %s" % (
 573                 area[engine.START], props.get("name", "noname"), props["access"], percent, func
 574             )
 575             subarea = self.model.AS.lookup_subarea(self.cur_addr())
 576             if subarea:
 577                 status += ", subarea: " + subarea[2]
 578             self.show_status(status)
 579         elif key == b"I":
 580             from scratchabit import memmap
 581             addr = memmap.show(self.model.AS, self.cur_addr())
 582             if addr is not None:
 583                 self.goto_addr(addr, from_addr=self.cur_addr())
 584             self.redraw()
 585         elif key == b"W":
 586             out_fname = "out.lst"
 587             with open(out_fname, "w") as f:
 588                 engine.render_partial(actions.TextSaveModel(f, self), 0, 0, 10000000)
 589             self.show_status("Disassembly listing written: " + out_fname)
 590         elif key == MENU_WRITE_ALL_HTML:
 591             out_fname = "out.html"
 592             with open(out_fname, "w") as f:
 593                 f.write("<pre>\n")
 594                 m = actions.HTMLSaveModel(f, self)
 595                 m.aspace = self.model.AS
 596                 engine.render_partial(m, 0, 0, 10000000)
 597                 f.write("</pre>\n")
 598             self.show_status("Disassembly HTML listing written: " + out_fname)
 599         elif key == b"\x17":  # Ctrl+W
 600             outfile = actions.write_func_by_addr(APP, self.cur_addr(), feedback_obj=self)
 601             if outfile:
 602                 self.show_status("Wrote file: %s" % outfile)
 603         elif key == b"\x15":  # Ctrl+U
 604             # Next undefined
 605             addr = self.cur_addr()
 606             flags = self.model.AS.get_flags(addr)
 607             if flags == self.model.AS.UNK:
 608                 # If already on undefined, skip the current stride of them,
 609                 # as they indeed go in batches.
 610                 while True:
 611                     flags = self.model.AS.get_flags(addr)
 612                     if flags != self.model.AS.UNK:
 613                         break
 614                     addr = self.model.AS.next_addr(addr)
 615                     if addr is None:
 616                         break
 617
 618             if addr is not None:
 619                 while True:
 620                     flags = self.model.AS.get_flags(addr)
 621                     if flags == self.model.AS.UNK:
 622                         self.goto_addr(addr, from_addr=self.cur_addr())
 623                         break
 624                     addr = self.model.AS.next_addr(addr)
 625                     if addr is None:
 626                         break
 627
 628             if addr is None:
 629                 self.show_status("There're no further undefined strides")
 630
 631         elif key == b"\x06":  # Ctrl+F
 632             # Next non-function
 633             addr = self.cur_addr()
 634             flags = self.model.AS.get_flags(addr, ~ADDRESS_SPACE.ALT_CODE)
 635             if flags == self.model.AS.CODE:
 636                 # If already on non-func code, skip the current stride of it,
 637                 # as it indeed go in batches.
 638                 while True:
 639                     flags = self.model.AS.get_flags(addr, ~ADDRESS_SPACE.ALT_CODE)
 640                     self.show_status("fl=%x" % flags)
 641                     if flags not in (self.model.AS.CODE, self.model.AS.CODE_CONT):
 642                         break
 643                     addr = self.model.AS.next_addr(addr)
 644                     if addr is None:
 645                         break
 646
 647             if addr is not None:
 648                 while True:
 649                     flags = self.model.AS.get_flags(addr, ~ADDRESS_SPACE.ALT_CODE)
 650                     if flags == self.model.AS.CODE:
 651                         self.goto_addr(addr, from_addr=self.cur_addr())
 652                         break
 653                     addr = self.model.AS.next_addr(addr)
 654                     if addr is None:
 655                         break
 656
 657             if addr is None:
 658                 self.show_status("There're no further non-function code strides")
 659
 660         elif key in (b"/", b"?"):  # "/" and Shift+"/"
 661
 662             class FoundException(Exception): pass
 663
 664             class TextSearchModel(engine.Model):
 665                 def __init__(self, substr, ctrl, this_addr, this_subno):
 666                     super().__init__()
 667                     self.search = substr
 668                     self.ctrl = ctrl
 669                     self.this_addr = this_addr
 670                     self.this_subno = this_subno
 671                     self.cnt = 0
 672                 def add_object(self, addr, line):
 673                     super().add_object(addr, line)
 674                     # Skip virtual lines before the line from which we started
 675                     if addr == self.this_addr and line.subno < self.this_subno:
 676                         return
 677                     txt = line.render()
 678                     idx = txt.find(self.search)
 679                     if idx != -1:
 680                         raise FoundException((addr, line.subno), idx + line.LEADER_SIZE + len(line.indent))
 681                     if self.cnt % 256 == 0:
 682                         self.ctrl.show_status("Searching: 0x%x" % addr)
 683                     self.cnt += 1
 684                     # Don't accumulate lines
 685                     self._lines = []
 686                     self._addr2line = {}
 687
 688             if key == b"/":
 689                 d = Dialog(4, 4, title="Text Search")
 690                 d.add(1, 1, WLabel("Search for:"))
 691                 entry = WTextEntry(20, self.search_str)
 692                 entry.finish_dialog = ACTION_OK
 693                 d.add(13, 1, entry)
 694                 res = d.loop()
 695                 self.redraw()
 696                 self.search_str = entry.get_text()
 697                 if res != ACTION_OK or not self.search_str:
 698                     return
 699                 addr, subno = self.cur_addr_subno()
 700             else:
 701                 addr, subno = self.next_line_addr_subno()
 702
 703             try:
 704                 engine.render_from(TextSearchModel(self.search_str, self, addr, subno), addr, 10000000)
 705             except FoundException as res:
 706                 self.goto_addr(res.args[0], col=res.args[1], from_addr=self.cur_addr())
 707             else:
 708                 self.show_status("Not found: " + self.search_str)
 709
 710         elif key == MENU_PREFS:
 711             uiprefs.handle(APP)
 712
 713         elif key == MENU_PLUGIN:
 714             res = DTextEntry(30, "", title="Plugin module name:").result()
 715             self.redraw()
 716             if res:
 717                 self.show_status("Running '%s' plugin..." % res)
 718                 call_script(res)
 719                 self.update_model()
 720                 self.show_status("Plugin '%s' ran successfully" % res)
 721         else:
 722             self.show_status("Unbound key: " + repr(key))
 723
 724
 725 ACTION_MAP = {
 726     b"g": DisasmViewer.action_goto,
 727     b"a": DisasmViewer.action_make_ascii,
 728 }
 729
 730
 731 CPU_PLUGIN = None
 732 ENTRYPOINTS = []
 733 APP.show_bytes = 4
 734
 735 def filter_config_line(l):
 736     l = re.sub(r"#.*$", "", l)
 737     l = l.strip()
 738     return l
 739
 740 def load_symbols(fname):
 741     with open(fname) as f:
 742         for l in f:
 743             l = filter_config_line(l)
 744             if not l:
 745                 continue
 746             m = re.search(r"\b([A-Za-z_$.][A-Za-z0-9_$.]*)\s*=\s*((0x)?[0-9A-Fa-f]+)", l)
 747             if m:
 748                 #print(m.groups())
 749                 ENTRYPOINTS.append((m.group(1), int(m.group(2), 0)))
 750             else:
 751                 print("Warning: cannot parse entrypoint info from: %r" % l)
 752
 753
 754 # Allow undescores to separate digit groups
 755 def str2int(s):
 756     return int(s.replace("_", ""), 0)
 757
 758
 759 def parse_range(arg):
 760     # name start(len)
 761     # name start-end
 762     if "(" in arg:
 763         m = re.match(r"(.+?)\s*\(\s*(.+?)\s*\)", arg)
 764         start = str2int(m.group(1))
 765         end = start + str2int(m.group(2)) - 1
 766     else:
 767         m = re.match(r"(.+)\s*-\s*(.+)", arg)
 768         start = str2int(m.group(1))
 769         end = str2int(m.group(2))
 770     return start, end
 771
 772
 773 def parse_entrypoints(f):
 774     for l in f:
 775         l = filter_config_line(l)
 776         if not l:
 777             continue
 778         if l[0] == "[":
 779             return l
 780         m = re.match(r'load "(.+?)"', l)
 781         if m:
 782             load_symbols(m.group(1))
 783         else:
 784             label, addr = [v.strip() for v in l.split("=")]
 785             ENTRYPOINTS.append((label, int(addr, 0)))
 786     return ""
 787
 788 def parse_subareas(f):
 789     subareas = []
 790     for l in f:
 791         l = filter_config_line(l)
 792         if not l:
 793             continue
 794         if l[0] == "[":
 795             return l
 796
 797         args = l.split()
 798         assert len(args) == 2
 799         start, end = parse_range(args[1])
 800         engine.ADDRESS_SPACE.add_subarea(start, end, args[0])
 801     engine.ADDRESS_SPACE.finish_subareas()
 802     return ""
 803
 804
 805 def load_target_file(loader, fname):
 806     entry = loader.load(engine.ADDRESS_SPACE, fname)
 807     log.info("Loaded %s, entrypoint: %s", fname, hex(entry) if entry is not None else None)
 808     if entry is not None:
 809         ENTRYPOINTS.append(("_ENTRY_", entry))
 810
 811
 812 def parse_disasm_def(fname):
 813     global CPU_PLUGIN
 814     with open(fname) as f:
 815         for l in f:
 816             l = filter_config_line(l)
 817             if not l:
 818                 continue
 819             #print(l)
 820             while True:
 821                 if not l:
 822                     #return
 823                     break
 824                 if l[0] == "[":
 825                     section = l[1:-1]
 826                     print("Processing section: %s" % section)
 827                     if section == "entrypoints":
 828                         l = parse_entrypoints(f)
 829                     elif section == "subareas":
 830                         l = parse_subareas(f)
 831                     else:
 832                         assert 0, "Unknown section: " + section
 833                 else:
 834                     break
 835
 836             if not l:
 837                 break
 838
 839             if l.startswith("load"):
 840                 args = l.split()
 841                 if args[2][0] in string.digits:
 842                     addr = int(args[2], 0)
 843                     print("Loading %s @0x%x" % (args[1], addr))
 844                     engine.ADDRESS_SPACE.load_content(open(args[1], "rb"), addr)
 845                 else:
 846                     print("Loading %s (%s plugin)" % (args[1], args[2]))
 847                     loader = __import__(args[2])
 848                     load_target_file(loader, args[1])
 849             elif l.startswith("cpu "):
 850                 args = l.split()
 851                 CPU_PLUGIN = __import__(args[1])
 852                 if hasattr(CPU_PLUGIN, "arch_id"):
 853                     engine.set_arch_id(CPU_PLUGIN.arch_id)
 854                 print("Loading CPU plugin %s" % (args[1]))
 855             elif l.startswith("show bytes "):
 856                 args = l.split()
 857                 APP.show_bytes = int(args[2])
 858             elif l.startswith("area "):
 859                 args = l.split()
 860                 assert len(args) == 4
 861                 start, end = parse_range(args[2])
 862                 a = engine.ADDRESS_SPACE.add_area(start, end, {"name": args[1], "access": args[3].upper()})
 863                 print("Adding area: %s" % engine.str_area(a))
 864             else:
 865                 assert 0, "Unknown directive: " + l
 866
 867
 868 class MainScreen:
 869
 870     def __init__(self):
 871         self.screen_size = Screen.screen_size()
 872         self.e = DisasmViewer(1, 2, self.screen_size[0] - 2, self.screen_size[1] - 4)
 873
 874         menu_file = WMenuBox([
 875             ("Save (Shift+s)", b"S"),
 876             ("Write disasm (Shift+w)", b"W"),
 877             ("Write disasm in HTML", MENU_WRITE_ALL_HTML),
 878             ("Write function (Ctrl+w)", b"\x17"),
 879             ("Quit (q)", b"q")
 880         ])
 881         menu_goto = WMenuBox([
 882             ("Follow (Enter)", KEY_ENTER), ("Return (Esc)", KEY_ESC),
 883             ("Goto... (g)", b"g"), ("Search disasm... (/)", b"/"),
 884             ("Search next (Shift+/)", b"?"), ("Next undefined (Ctrl+u)", b"\x15"),
 885             ("Next non-function code (Ctrl+f)", b"\x06"),
 886         ])
 887         menu_edit = WMenuBox([
 888             ("Undefined (u)", b"u"),
 889             ("Code (c)", b"c"),
 890             ("Alt code (Shift+c)", b"C"),
 891             ("Data (d)", b"d"),
 892             ("ASCII String (a)", b"a"), ("Filler (f)", b"f"), ("Make label (n)", b"n"),
 893             ("Mark function start (F)", b"F"), ("Add code to function", MENU_ADD_TO_FUNC),
 894             ("Number/Address (o)", b"o"), ("Hex/dec (h)", b"h"),
 895         ])
 896         menu_analysis = WMenuBox([
 897             ("Info (whereami) (i)", b"i"), ("Memory map (Shift+i)", b"I"),
 898             ("Run plugin...", MENU_PLUGIN),
 899             ("Preferences...", MENU_PREFS),
 900         ])
 901         menu_help = WMenuBox([
 902             ("Help (F1)", KEY_F1), ("About...", "about"),
 903         ])
 904         self.menu_bar = WMenuBar([
 905             ("File", menu_file), ("Goto", menu_goto), ("Edit", menu_edit),
 906             ("Analysis", menu_analysis), ("Help", menu_help)
 907         ])
 908         self.menu_bar.permanent = True
 909
 910     def redraw(self, allow_cursor=True):
 911         self.menu_bar.redraw()
 912         self.e.attr_color(C_B_WHITE, C_BLUE)
 913         self.e.draw_box(0, 1, self.screen_size[0], self.screen_size[1] - 2)
 914         self.e.attr_reset()
 915         self.e.redraw()
 916         if allow_cursor:
 917             self.e.cursor(True)
 918
 919     def loop(self):
 920         while 1:
 921             key = self.e.get_input()
 922             if isinstance(key, list):
 923                 x, y = key
 924                 if self.menu_bar.inside(x, y):
 925                     self.menu_bar.focus = True
 926
 927             if self.menu_bar.focus:
 928                 res = self.menu_bar.handle_input(key)
 929                 if res == ACTION_CANCEL:
 930                     self.menu_bar.focus = False
 931                 elif res is not None and res is not True:
 932
 933                     res = self.e.handle_input(res)
 934                     if res is not None and res is not True:
 935                         return res
 936             else:
 937                 if key == KEY_F9:
 938                     self.menu_bar.focus = True
 939                     self.menu_bar.redraw()
 940                     continue
 941
 942                 res = self.e.handle_input(key)
 943
 944                 if res is not None and res is not True:
 945                     return res
 946
 947
 948 def call_script(script, progress_func=None):
 949     mod = __import__(script)
 950     main_f = getattr(mod, "main", None)
 951     if main_f:
 952         main_f(APP)
 953     # A script might have queues some entrypoints, etc.
 954     # Analyze them now.
 955     engine.analyze(progress_func)
 956
 957
 958 if __name__ == "__main__":
 959
 960     argp = argparse.ArgumentParser(description="ScratchABit interactive disassembler")
 961     argp.add_argument("file", help="Input file (binary or disassembly .def)")
 962     argp.add_argument("--script", action="append", help="Run script from file after loading environment")
 963     argp.add_argument("--save", action="store_true", help="Save after --script and quit; don't show UI")
 964     args = argp.parse_args()
 965
 966     # Plugin dirs are relative to the dir where scratchabit.py resides.
 967     # sys.path[0] below provide absolute path of this dir, resolved for
 968     # symlinks.
 969     plugin_dirs = ["plugins", "plugins/cpu", "plugins/loader"]
 970     for d in plugin_dirs:
 971         sys.path.append(os.path.join(sys.path[0], d))
 972     log.basicConfig(filename="scratchabit.log", format='%(asctime)s %(message)s', level=log.DEBUG)
 973     log.info("Started")
 974
 975     if args.file.endswith(".def"):
 976         parse_disasm_def(args.file)
 977         project_name = args.file.rsplit(".", 1)[0]
 978     else:
 979         import default_plugins
 980         for loader_id in default_plugins.loaders:
 981             loader = __import__(loader_id)
 982             arch_id = loader.detect(args.file)
 983             if arch_id:
 984                 break
 985         if not arch_id:
 986             print("Error: file '%s' not recognized by default loaders" % args.file)
 987             sys.exit(1)
 988         if arch_id not in default_plugins.cpus:
 989             print("Error: no plugin for CPU '%s' as detected for file '%s'" % (arch_id, args.file))
 990             sys.exit(1)
 991
 992         engine.set_arch_id(arch_id)
 993         load_target_file(loader, args.file)
 994         CPU_PLUGIN = __import__(default_plugins.cpus[arch_id])
 995         project_name = args.file
 996
 997     p = CPU_PLUGIN.PROCESSOR_ENTRY()
 998     if hasattr(p, "config"):
 999         p.config()
1000     engine.set_processor(p)
1001     if hasattr(p, "help_text"):
1002         help.set_cpu_help(p.help_text)
1003     APP.cpu_plugin = p
1004     APP.aspace = engine.ADDRESS_SPACE
1005     APP.is_ui = False
1006     engine.ADDRESS_SPACE.is_loading = True
1007
1008     # Calc various offset based on show_bytes value
1009     APP.set_show_bytes(APP.show_bytes)
1010
1011     # Strip suffix if any from def filename
1012     project_dir = project_name + ".scratchabit"
1013
1014     if saveload.save_exists(project_dir):
1015         saveload.load_state(project_dir)
1016     else:
1017         for label, addr in ENTRYPOINTS:
1018             if engine.arch_id == "arm_32" and addr & 1:
1019                 addr &= ~1
1020                 engine.ADDRESS_SPACE.make_alt_code(addr)
1021             if engine.ADDRESS_SPACE.is_exec(addr):
1022                 engine.add_entrypoint(addr)
1023             engine.ADDRESS_SPACE.make_unique_label(addr, label)
1024         def _progress(cnt):
1025             sys.stdout.write("Performing initial analysis... %d\r" % cnt)
1026         engine.analyze(_progress)
1027         print()
1028
1029     #engine.print_address_map()
1030
1031     if args.script:
1032         def _progress(cnt):
1033             sys.stdout.write("Performing analysis after running script(s)... %d\r" % cnt)
1034         for script in args.script:
1035             call_script(script, _progress)
1036
1037     if args.save:
1038         saveload.save_state(project_dir)
1039         sys.exit()
1040
1041     addr_stack = []
1042     if os.path.exists(project_dir + "/session.addr_stack"):
1043         addr_stack = saveload.load_addr_stack(project_dir)
1044         print(addr_stack)
1045         show_addr = addr_stack.pop()
1046     else:
1047         if ENTRYPOINTS:
1048             show_addr = ENTRYPOINTS[0][1]
1049             if engine.arch_id == "arm_32":
1050                  show_addr &= ~1
1051         else:
1052             show_addr = engine.ADDRESS_SPACE.min_addr()
1053
1054     t = time.time()
1055     #_model = engine.render()
1056     _model = engine.render_partial_around(show_addr, 0, HEIGHT * 2)
1057     print("Rendering time: %fs" % (time.time() - t))
1058     #print(_model.lines())
1059     #sys.exit()
1060
1061     engine.ADDRESS_SPACE.is_loading = False
1062     engine.ADDRESS_SPACE.changed = False
1063     Screen.init_tty()
1064     try:
1065         Screen.cls()
1066         Screen.enable_mouse()
1067         main_screen = MainScreen()
1068         APP.main_screen = main_screen
1069         APP.is_ui = True
1070         main_screen.e.set_model(_model)
1071         main_screen.e.addr_stack = addr_stack
1072         main_screen.e.goto_addr(show_addr)
1073         Screen.set_screen_redraw(main_screen.redraw)
1074         main_screen.redraw()
1075         main_screen.e.show_status("Press F1 for help, F9 for menus")
1076         main_screen.loop()
1077     except:
1078         log.exception("Unhandled exception")
1079         raise
1080     finally:
1081         Screen.goto(0, main_screen.screen_size[1])
1082         Screen.cursor(True)
1083         Screen.disable_mouse()
1084         Screen.deinit_tty()
1085         Screen.wr("\n\n")
1086         saveload.save_session(project_dir, main_screen.e)