smalictrace.py

#!/usr/bin/python2
# -*- coding: UTF-8 -*-
"""
Smali CTrace - a smali call hierarchy tracer
Author: fikr4n
"""

import os
import re
import gtk
from gtk import gdk
import pango
import gobject
import time
import fnmatch
import cgi
import cPickle as pickle
import collections
import threading


ParseResult = collections.namedtuple('ParseResult', ['callers', 'members'])


class Dict(dict):
	_SENTINEL = object()
	def getcreate(self, key, value_type):
		r = self.get(key, self._SENTINEL)
		if r is self._SENTINEL:
			self[key] = r = value_type()
		return r


class CallInfo(object):
	DIRECT, VIRTUAL_DIRECT, VIRTUAL_INDIRECT, SUPER_DIRECT, SUPER_INDIRECT = (
		1, 2, 4, 8, 16)
	def __init__(self):
		self.call_types = 0
	def select_from_call_types(self, seq):
		def fun(x):
			fun.v >>= 1
			return fun.v & self.call_types
		fun.v = 1 << len(seq)
		return filter(fun, seq)


class ClassInfo(object):
	def __init__(self):
		self.superclass = None
		self.subclasses = set()
		self.methods = set()
	def add_subclass(self, x):
		self.subclasses.add(x)
	def add_method(self, x):
		self.methods.add(x)
	def is_declared(self):
		return self.superclass is not None


class Parser(object):
	"""Smali parser"""
	
	# Cannot detect comment or string constant correctly
	_comment_pattern = re.compile(r'(?<!\S)#.*$')
	_general_pattern = re.compile(r'^\s*([^\s#]\S*)(.*\s)(\S+)\s*$')
	_class_pattern = re.compile(r'^L\S+;(?:->\S+)?$')
	
	def __init__(self, progress_handler):
		self._callers = Dict()  # {callee: {caller: call_info}}
		self._members = Dict()  # {declared_entity: class_info or None if not class}
		self._canonicals = {}  # {smali_type: java_type}
		self._handler = progress_handler
	
	def _notify_progress(self, fraction):
		if self._handler:
			self._handler(fraction)
	
	def _list_files(self, path, result):
		"""List files and class names recursively into result list"""
		for f in os.listdir(path):
			self._notify_progress(None)
			child = os.path.join(path, f)
			if os.path.isdir(child):
				self._list_files(child, result)
			elif os.path.isfile(child):
				if f.endswith('.smali'):
					result.append(child)
	
	def parse_dir(self, path):
		files = []
		self._list_files(path, files)
		count = len(files) * 2.0  # 1/2 of overall progress
		for i, path in enumerate(files):
			self.parse_file(path)
			self._notify_progress(i / count)
	
	def parse_file(self, path):
		with open(path) as f:
			class_name = None
			current_scope = None
			for line in f:
				line = self._comment_pattern.sub('', line)  # strip comment first
				m = self._general_pattern.match(line)  # identify the line
				if not m:
					continue
				token_1 = m.group(1)
				token_n = m.group(3)
				if token_1 == '.method':
					current_scope = class_name + '->' + token_n
					self._add_def(current_scope, token_1, cls=class_name,
						simple_name=token_n)
				elif token_1 == '.field':
					self._add_def(class_name + '->' + token_n, token_1)
				elif token_1 == '.class':
					current_scope = class_name = token_n
				elif token_1 == '.super':
					self._add_def(token_n, token_1, cls=class_name)
				elif token_1 == '.implements':
					self._add_def(token_n, token_1, cls=class_name)
				elif token_1 in ('.source', '.registers', '.prologue', '.line',
						'.end'):
					pass  # ignore
				elif self._class_pattern.match(token_n):  # assume as usage
					self._add_use(current_scope, token_n, token_1)
	
	def derive_indirect_calls(self):
		addition = set()
		count = len(self._callers) * 4.0  # 2~3/4 of overall progress
		for i, (callee, callers) in enumerate(self._callers.items()):
			callee_ = callee.split('->', 1)
			for caller, info in callers.items():
				if info.call_types & CallInfo.SUPER_DIRECT:
					self._list_super_indirect_calls(caller, callee_, addition)
				if info.call_types & CallInfo.VIRTUAL_DIRECT:
					self._list_virtual_indirect_calls(caller, callee_, addition)
			if i % 100 == 0:
				self._notify_progress(i / count + 0.5)
		
		count = len(addition) * 4.0  # 3~4/4 of overall progress
		for i, (callee, caller, call_type) in enumerate(addition):
			x = self._callers.getcreate(callee, Dict).getcreate(caller, CallInfo)
			x.call_types |= call_type
			
			# Early canonicalizing
			if callee not in self._canonicals:
				self._canonicals[callee] = Canonical.from_smali(callee)
			
			if i % 100 == 0:
				self._notify_progress(i / count + 0.75)
	
	def _find_declaring_ancestor(self, cls, method):
		while True:
			class_info = self._members.get(cls)
			if class_info is None:
				return None
			if method in class_info.methods:
				return cls + '->' + method
			cls = class_info.superclass
	
	def _list_super_indirect_calls(self, caller, callee_, result):
		callee = self._find_declaring_ancestor(*callee_)
		if callee is not None:
			result.add((callee, caller, CallInfo.SUPER_INDIRECT))
	
	def _list_virtual_indirect_calls(self, caller, callee_, result):
		def add(cls):
			class_info = self._members.get(cls)
			if class_info is None:
				return
			if callee_method in class_info.methods:
				result.add((cls + '->' + callee_method, caller,
					CallInfo.VIRTUAL_INDIRECT))
			for subclass in class_info.subclasses:
				add(subclass)
		
		callee_class, callee_method = callee_
		add(callee_class)
		callee = self._find_declaring_ancestor(callee_class, callee_method)
		if callee is not None:
			result.add((callee, caller, CallInfo.VIRTUAL_INDIRECT))
	
	def _add_use(self, caller, callee, invoke_type):
		call_type = 0
		if invoke_type == 'invoke-virtual' or invoke_type == 'invoke-interface':
			call_type = CallInfo.VIRTUAL_DIRECT
		elif invoke_type == 'invoke-super':
			call_type = CallInfo.SUPER_DIRECT
		elif invoke_type.startswith('invoke-'):
			call_type = CallInfo.DIRECT
		
		callers = self._callers.getcreate(callee, Dict)
		info = callers.getcreate(caller, CallInfo)
		info.call_types |= call_type
		
		# Early canonicalizing
		if callee not in self._canonicals:
			self._canonicals[callee] = Canonical.from_smali(callee)
	
	def _add_def(self, name, kind, cls=None, simple_name=None):
		if kind == '.method':
			self._members[name] = None
			self._members.getcreate(cls, ClassInfo).add_method(simple_name)
		elif kind == '.field':
			self._members[name] = None
		elif kind == '.super':
			self._members.getcreate(cls, ClassInfo).superclass = name
			self._members.getcreate(name, ClassInfo).add_subclass(cls)
		elif kind == '.implements':
			self._members.getcreate(name, ClassInfo).add_subclass(cls)
		
		# Early canonicalizing
		if name not in self._canonicals:
			self._canonicals[name] = Canonical.from_smali(name)
	
	def get_result(self):
		def canonical(x):
			return self._canonicals.get(x) or Canonical.from_smali(x)

		callers = {canonical(k):
			{canonical(kk): vv for kk, vv in v.items()}
			for k, v in self._callers.items()}
		members = set(canonical(i) for i in self._members)
		return ParseResult(callers, members)


class Canonical(object):
	"""Canonicalizer/decanonicalizer"""
	
	_method_pattern = re.compile(r'^(.*)->(.*)\((.*)\)(.*)$')  # X->N(X*)X
	_field_pattern = re.compile(r'^(.*)->(.*):(.*)$')  # X->N:X
	_type_pattern = re.compile(r'(\[*)([ZBCSIJFDV]|L([^;\s]+);)')
	_prims = {'Z': 'boolean', 'B': 'byte', 'C': 'char', 'S': 'short',
		'I': 'int', 'J': 'long', 'F': 'float', 'D': 'double', 'V': 'void'}
	
	@classmethod
	def from_smali(self, signature):
		def tr(s):
			return ', '.join((i.group(3).replace('/', '.') if i.group(3) else
			self._prims[i.group(2)]) + ('[]' * len(i.group(1))) for i in
			self._type_pattern.finditer(s))
		
		m = self._method_pattern.match(signature)
		if m:
			return '%s.%s(%s) : %s' % (tr(m.group(1)), m.group(2), tr(m.group(3)), tr(m.group(4)))
		m = self._field_pattern.match(signature)
		if m:
			return '%s.%s : %s' % (tr(m.group(1)), m.group(2), tr(m.group(3)))
		return tr(signature)


class CallHierarchyView(object):

	def __init__(self, parser, max_tree_depth):
		self.parser = parser
		self._max_tree_depth = max_tree_depth
		self.search_thread = None
		self.search_lock = threading.Lock()
		self.name_entry = name_entry = gtk.Entry()
		name_entry.connect('activate', self.on_name_entry_activate)
		search_button = self.search_button = gtk.Button(stock=gtk.STOCK_FIND)
		search_button.connect('clicked', self.on_search)
		search_button.set_no_show_all(True)
		search_button.show()
		stop_button = self.stop_button = gtk.Button(stock=gtk.STOCK_STOP)
		stop_button.connect('clicked', self.stop_search)
		stop_button.set_no_show_all(True)
		self.progress = progress = gtk.ProgressBar()
		progress.set_size_request(-1, 8)
		name_box = gtk.HBox(spacing=8)
		name_box.set_border_width(8)
		name_box.pack_start(name_entry, True, True)
		name_box.pack_start(search_button, False)
		name_box.pack_start(stop_button, False)
		self.store = gtk.TreeStore(str, str, str, str)  # name, icon, fnstyle, info
		self.treeview = treeview = self._create_treeview(self.store)
		scroller = gtk.ScrolledWindow()
		scroller.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
		scroller.add(treeview)
		settings_view = SettingsView(max_depth_value=max_tree_depth)
		settings_view.on_max_depth_changed = self.on_max_depth_changed
		extension = gtk.expander_new_with_mnemonic('_Settings')
		extension.add(settings_view.view)
		vbox = gtk.VBox()
		vbox.pack_start(name_box, False)
		vbox.pack_start(progress, False)
		vbox.pack_start(scroller, True, True)
		vbox.pack_start(extension, False)
		self.window = window = gtk.Window()
		window.set_title(Application.title)
		window.set_default_size(800, 600)
		window.set_icon(Application.get_icon())
		window.add(vbox)
	
	def on_name_entry_activate(self, widget):
		if self.search_button.get_visible():
			self.search_button.emit('clicked')
		elif self.stop_button.get_visible():
			self.stop_button.emit('clicked')
	
	def on_max_depth_changed(self, value):
		self._max_tree_depth = value
	
	def on_treeview_key_press(self, widget, event):
		if gdk.keyval_name(event.keyval) in ('c', 'C') and event.state | \
				gdk.LOCK_MASK == gdk.CONTROL_MASK | gdk.LOCK_MASK:
			model, pathlist = self.treeview.get_selection().get_selected_rows()
			gtk.clipboard_get().set_text('\n'.join(str(model.get_value(
				model.get_iter(p), 0)) for p in pathlist))
	
	def _create_treeview(self, model):
		cell_text = gtk.CellRendererText()
		cell_icon = gtk.CellRendererPixbuf()
		column = gtk.TreeViewColumn()
		column.pack_start(cell_icon, False)
		column.add_attribute(cell_icon, 'stock_id', 1)
		column.pack_start(cell_text, True)
		column.set_cell_data_func(cell_text, self.render_name_cell)
		treeview = gtk.TreeView(model)
		treeview.append_column(column)
		treeview.set_headers_visible(False)
		treeview.get_selection().set_mode(gtk.SELECTION_MULTIPLE)
		treeview.connect('key_press_event', self.on_treeview_key_press)
		return treeview
	
	_class_pattern = re.compile(r'^([^.(:)]+\.)*([^.(:)]+)\.[^.(:)]+[(:]')
	def render_name_cell(self, column, cell, model, titer):
		base = cgi.escape(str(model.get_value(titer, 0)))
		m = self._class_pattern.match(base)
		if m:  # currently support only Java syntax
			s, e = m.span(2)
			base = '%s<span color="#0d47a1">%s</span>%s' % (base[:s], base[s:e],
				base[e:])
		info = str(model.get_value(titer, 3))
		font = str(model.get_value(titer, 2))
		font_style = 'italic' if 'i' in font else 'normal'
		font_weight = 'bold' if 'b' in font else 'normal'
		markup = '%s<span font_style="%s" font_weight="%s">%s</span>' % (
			info, font_style, font_weight, base)
		cell.set_property('markup', markup)
	
	def on_search_start(self):
		self.search_button.hide()
		self.stop_button.show()
	
	def on_search_stop(self):
		self.search_button.show()
		self.stop_button.hide()
		gobject.idle_add(self.progress.set_fraction, 0.0)
	
	def stop_search(self, *args):
		if self.search_thread and not self.search_thread.event_flag.is_set():
			self.search_thread.event_flag.set()
			self.on_search_stop()

	def on_search(self, widget):
		self.stop_search()
		self.on_search_start()
		
		query = self.name_entry.get_text()
		flag = threading.Event()
		pulse_progress_until(self.progress, flag)
		self.search_thread = thread = threading.Thread(target=self.do_search,
			args=(query, flag))
		thread.daemon = True
		thread.event_flag = flag
		thread.start()
	
	def do_search(self, query, event_flag):
		def get_icon(n):
			if n not in self.parser.members:  # no check for undeclared class
				icon = gtk.STOCK_DELETE
			elif '(' in n:
				icon = gtk.STOCK_EXECUTE
			elif ':' in n:
				icon = gtk.STOCK_SELECT_FONT
			else:
				icon = gtk.STOCK_FILE
			return icon
		
		def append_row(parent, names, depth, call_info):
			
			name = names[0]
			call_info_str = call_info.select_from_call_types(u'⇡↑⇠←⇤') \
							.replace(u'⇡↑', u'↟') \
							.replace(u'⇠←', u'↞')
			it = self.store.append(parent, (
				name,  # class/method/field name
				get_icon(name),  # icon
				'b' if name in toplevel else '',
					# bold (in root) or not
				call_info_str
				))
			
			# Append children
			d = depth + 1
			if d >= self._max_tree_depth:
				self.store.append(it, (' ...maximum depth (%d) reached' %
					self._max_tree_depth, gtk.STOCK_INFO, 'i',
					call_info_str))
				return
			for caller, info in self.parser.callers.get(name, {}).items():
				if event_flag.is_set():  # stopped earlier
					return
				if caller not in names:
					names.appendleft(caller)
					append_row(it, names, d, info)
					names.popleft()
				else:
					self.store.append(it, (' ...recursive ' + caller,
						gtk.STOCK_INFO, 'i', call_info_str))
		
		with self.search_lock:
			gobject.idle_add(self.treeview.set_model, None)
		
			toplevel = set()
			self.store.clear()
			for k in self.parser.callers:
				if fnmatch.fnmatch(k, query):
					toplevel.add(k)
			for i in toplevel:
				append_row(None, collections.deque((i,)), 0, CallInfo())
		
			# Finalization
			if not event_flag.is_set():
				event_flag.set()
				gobject.idle_add(self.treeview.set_model, self.store)
				gobject.idle_add(self.on_search_stop)


class SettingsView(object):
	def __init__(self, max_depth_value):
		self.on_max_depth_changed = None
		max_depth_entry = self.max_depth_entry = gtk.SpinButton(
			gtk.Adjustment(value=max_depth_value, lower=1, upper=100,
			step_incr=1, page_incr=10), digits=0)
		max_depth_entry.get_adjustment().connect('value_changed',
			self._max_depth_changed, max_depth_entry)
		max_depth_entry.set_alignment(0)
		max_depth_label = gtk.Label()
		max_depth_label.set_alignment(1, 0.5)
		max_depth_label.set_text_with_mnemonic('Maximum _Depth')
		max_depth_label.set_mnemonic_widget(max_depth_entry)
		legend = self._create_legend()
		table = self.view = gtk.Table(2, 3)
		table.set_border_width(8)
		table.set_row_spacings(8)
		table.set_col_spacings(8)
		table.attach(max_depth_label, 0, 1, 0, 1, gtk.FILL, gtk.FILL)
		table.attach(max_depth_entry, 1, 2, 0, 1, gtk.SHRINK, gtk.FILL)
		table.attach(gtk.Label(), 2, 3, 0, 1)  # space filler
		table.attach(legend, 0, 3, 1, 2)
	
	def _max_depth_changed(self, widget, spin):
		if self.on_max_depth_changed:
			self.on_max_depth_changed(spin.get_value_as_int())
	
	def _create_legend(self):
		icons = (
			(gtk.STOCK_FILE, 'Class'),
			(gtk.STOCK_EXECUTE, 'Method'),
			(gtk.STOCK_SELECT_FONT, 'Field'),
			(gtk.STOCK_DELETE, 'Declaration not found'),
			(gtk.STOCK_INFO, 'Miscellaneous (info)'),
			(u'⇤', 'Direct call – e.g. static, private, etc'),
			(u'↞', 'Virtual call (⇠ + ←)'),
			(u'⇠', 'Virtual call (indirect) which could be performed because '
				'of polymorphism'),
			(u'←', 'Virtual call (direct only) which does not actually '
				'performed – e.g. interface method'),
			(u'↟', 'Super call (⇡ + ↑)'),
			(u'⇡', 'Super call (indirect) because direct super does not '
				'declare the method'),
			(u'↑', 'Super call (direct only) which does not actually '
				'performed – e.g. not declared here'),
		)
		table = gtk.Table(7, 5)
		table.set_border_width(8)
		table.set_row_spacings(8)
		table.set_col_spacings(8)
		separator = gtk.VSeparator()
		table.attach(separator, 2, 3, 0, 7, 0)
		x, y = 0, 0
		for icon, desc in icons:
			if len(icon) == 1:
				image = gtk.Label(icon)
			else:
				image = gtk.Image()
				image.set_from_stock(icon, gtk.ICON_SIZE_MENU)
			image.set_alignment(1, 0.5)
			label = gtk.Label(desc)
			label.set_alignment(0, 0.5)
			table.attach(image, x + 0, x + 1, y, y + 1, gtk.FILL)
			table.attach(label, x + 1, x + 2, y, y + 1, gtk.FILL)
			y += 1
			if y == 5 and x == 0:
				x, y = 3, 0
		frame = gtk.Frame('Legend')
		frame.add(table)
		return frame


class ProgressView(object):

	def __init__(self):
		progress = self.progress = gtk.ProgressBar()
		progress.set_text('Loading...')
		align = gtk.Alignment(0.5, 0.5, xscale=1.0)
		align.add(progress)
		window = self.window = gtk.Window()
		window.set_title(Application.title)
		window.set_border_width(8)
		window.add(align)
		window.set_default_size(400, 120)
		self._last_update_progress = 0.0
	
	def update_progress(self, fraction=None):
		t = time.time()
		if t - self._last_update_progress < 0.25:
			return
		self._last_update_progress = t
		gobject.idle_add(self.update_progress_immediate, fraction)
	
	def update_progress_immediate(self, fraction=None):
		if fraction is None:
			self.progress.pulse()
		else:
			self.progress.set_fraction(fraction)


def pulse_progress_until(progress, event):
	def do_pulse():
		if not event.is_set():
			progress.pulse()
			return True
	gobject.timeout_add(100, do_pulse)


class Application(object):

	title = 'Smali CTrace'

	def __init__(self, args):
		self._classpath = args.classpath
		self._cache_filename = args.save_cache
		self._max_depth = args.max_depth
		gobject.threads_init()
		self.parser = None
		self._launch_loading()
		gtk.main()
	
	@staticmethod
	def get_icon():
		return gtk.icon_theme_get_default().load_icon(gtk.STOCK_ZOOM_IN, 128, 0)
	
	def _launch_loading(self):
		v = ProgressView()
		v.window.set_icon(self.get_icon())
		v.window.connect('delete_event', gtk.main_quit)
		v.window.show_all()
		self.parser = Parser(v.update_progress)
		delivery = {'progress_view': v}
		thread = threading.Thread(target=self._do_parsing, args=(delivery,))
		thread.daemon = True
		thread.start()

	def _do_parsing(self, delivery):
		try:
			if os.path.isdir(self._classpath):
				self.parser.parse_dir(self._classpath)
				self.parser.derive_indirect_calls()
				event = threading.Event()
				pulse_progress_until(delivery['progress_view'].progress, event)
				parse_result = self.parser.get_result()
				if self._cache_filename:
					try:
						with open(self._cache_filename, 'wb') as f:
							pickle.dump(parse_result, f)
					except IOError as e:
						gobject.idle_add(self._launch_error, e, "Can't save to file")
				event.set()
			else:
				event = threading.Event()
				pulse_progress_until(delivery['progress_view'].progress, event)
				with open(self._classpath, 'rb') as f:
					parse_result = pickle.load(f)
				event.set()
			self.parse_result = parse_result
		
			gobject.idle_add(self._launch_main, delivery)
		except IOError as e:
			gobject.idle_add(self._launch_error, e, "Can't open file", True)
	
	def _launch_main(self, delivery):
		delivery['progress_view'].window.destroy()
		v = CallHierarchyView(self.parse_result, self._max_depth)
		v.window.connect('delete_event', gtk.main_quit)
		v.window.show_all()
	
	def _launch_error(self, exception, message, fatal=False):
		dlg = gtk.MessageDialog(parent=None,
			type=gtk.MESSAGE_ERROR, buttons=gtk.BUTTONS_OK, message_format=
			'%s\n\n%s' % (message, str(exception)))
		dlg.set_title('Error')
		dlg.connect('response', gtk.main_quit if fatal else lambda *x: dlg.destroy())
		dlg.show_all()


if __name__ == '__main__':
	import argparse
	parser = argparse.ArgumentParser(description='help you inspecting smali')
	parser.add_argument('classpath', default='.', help='root directory of the classes or cache file (see "-c")')
	parser.add_argument('-c', '--save-cache', metavar='FILE', help='save cache to file, ignored if input is cache file')
	parser.add_argument('--max-depth', type=int, default=16, metavar='DEPTH', help='maximum tree depth (default to 16)')
	args = parser.parse_args()
	Application(args)

# TODO
# - Move handler to method argument
# - Rename member to declared entity

# fikr4n 2016