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