ScratchABit.py

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