import networkx as nx
from awesome.context import ignored
import sark
import sark.qt
from sark.qt import QtGui, QtCore
import logging
import logging.handlers as handlers
import os
from time import ctime
from idaapi import plugin_t
import idaapi
import idautils
import idc
import sark.ui
from DIE.Lib.IDAConnector import *
import DIE.Lib.DieConfig
import DIE.Lib.DIEDb
from DIE.Lib import DebugAPI
#import DIE.UI.BPView
import DIE.UI.FunctionViewEx
import DIE.UI.ValueViewEx
import DIE.UI.ParserView
import DIE.UI.BPView
import DIE.UI.SetupView
import DIE.UI.Die_Icons
from DIE.UI.FuncScopeChooser import ScopeChooser
from DIE.UI.AboutScreen import AboutWindow
from DIE.Lib.DIE_Exceptions import DbFileMismatch
class MenuHelperException(Exception):
pass
class DieManager:
"""
Manage the DIE framework
"""
def __init__(self, is_dbg_log=True, is_dbg_pause=False, is_dbg_profile=False):
### Logging ###
log_filename = os.path.join(os.getcwd(), "DIE.log")
self._menus_names = []
#TODO: Fix logging to include rotating_file_handler \ console_logging
if is_dbg_log:
logging.basicConfig(filename=log_filename,
level=logging.DEBUG,
format='[%(asctime)s] [%(levelname)s] [%(name)s][%(filename)s:%(lineno)s] : %(message)s')
else:
logging.basicConfig(filename=log_filename,
level=logging.INFO,
format='[%(asctime)s] [%(levelname)s] [%(name)s][%(filename)s:%(lineno)s] : %(message)s')
idaapi.msg("Logfile created at %s\n" % log_filename)
self.logger = logging.getLogger(__name__)
### DIE Configuration ###
self.config_file_name = os.path.join(os.getcwd(), "DIE.cfg")
DIE.Lib.DieConfig.initialize()
config = DIE.Lib.DieConfig.get_config()
try:
config.load(self.config_file_name)
except IOError:
pass
except:
import traceback
idaapi.msg(traceback.format_exc())
self.die_config = config
self.icon_list = {}
self.debugAPI = DebugAPI.DebugHooker(is_dbg_pause, is_dbg_profile)
DIE.Lib.DIEDb.initialize_db()
self.die_db = DIE.Lib.DIEDb.get_db()
self.is_marked = False
DIE.UI.FunctionViewEx.initialize()
DIE.UI.ValueViewEx.initialize()
DIE.UI.BPView.initialize()
DIE.UI.ParserView.initialize()
DIE.UI.Die_Icons.initlialize()
self.function_view = DIE.UI.FunctionViewEx.get_view()
self.value_view = DIE.UI.ValueViewEx.get_view()
self.bp_view = DIE.UI.BPView.get_view()
self.parser_view = DIE.UI.ParserView.get_view()
self.load_icons()
return
###########################################################################
# Icons
def load_icon(self, icon_filename, icon_key_name):
"""
Load a single custom icon
@param icon_filename: Icon file name
@param icon_key_name: The key value to store the icon with in the icon_list.
"""
try:
icons_path = self.die_config.icons_path
icon_filename = os.path.join(icons_path, icon_filename)
icon_num = idaapi.load_custom_icon(icon_filename)
self.icon_list[icon_key_name.lower()] = icon_num
return True
except Exception as ex:
self.logger.error("Failed to load icon %s: %s", icon_filename, ex)
return False
def load_icons(self):
"""
Load custom DIE Icons
"""
self.load_icon("dbg.png", "debug")
self.load_icon("dbg_all.png", "debug_all")
self.load_icon("dbg_custom.png", "debug_scope")
self.load_icon("die.png", "die")
self.load_icon("funcview.png", "function_view")
self.load_icon("valueview.png", "value_view")
self.load_icon("stop.png", "exception_view")
self.load_icon("settings.png", "settings")
self.load_icon("plugins.png", "plugins")
self.load_icon("save.png", "save")
self.load_icon("load.png", "load")
###########################################################################
# Menu Items
# Handlers
class DIE_gohere_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.go_here()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_goall_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.go_all()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_goscope_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.show_scope_chooser()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_show_funcview_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.show_function_view()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_show_valview_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.show_value_view()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_show_parsersview_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.show_parser_view()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_show_exceptionsview_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.show_breakpoint_view()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_load_db_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.load_db()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_save_db_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.save_db()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_show_cfg_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.show_cfg()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_mark_execflow_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.mark_exec_flow()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_show_about_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.show_about()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
class DIE_show_settings_Handler(idaapi.action_handler_t):
def __init__(self, outer_instance):
idaapi.action_handler_t.__init__(self)
self.outer_instance = outer_instance
def activate(self, ctx):
self.outer_instance.show_settings()
return True
def update(self, ctx):
return idaapi.AST_ENABLE_ALWAYS
def add_menu_item_helper(self, name, text, tooltip, handler, icon, shortcut):
description = idaapi.action_desc_t(name, text, handler, shortcut, tooltip, icon)
idaapi.register_action(description)
idaapi.attach_action_to_menu("DIE/" + text, name, idaapi.SETMENU_APP)
self._menus_names.append(name)
def add_menu_items(self):
# Add root level menu
idaapi.create_menu("DIE", "DIE")
# Save DieDB
self.add_menu_item_helper("DIE:savedb", "Save DieDB", 'DIE Save DieDB',
self.DIE_save_db_Handler(self), self.icon_list["save"], '')
# Load DieDB
self.add_menu_item_helper("DIE:loaddb", "Load DieDB", 'DIE Load DieDB',
self.DIE_load_db_Handler(self), self.icon_list["load"], '')
# Show complete execution CFG
self.add_menu_item_helper("DIE:showcfg", "Show CFG", 'DIE Show CFG',
self.DIE_show_cfg_Handler(self), -1, '')
# Mark\Unmark Execution Flow
self.add_menu_item_helper("DIE:markexecutionflow", "Mark\Unmark Execution Flow",
'DIE Mark\Unmark Execution Flow', self.DIE_mark_execflow_Handler(self), -1, '')
# Parser View
self.add_menu_item_helper("DIE:parsersview", "Parsers View", 'DIE Parsers View',
self.DIE_show_parsersview_Handler(self), self.icon_list["plugins"], '')
# Exceptions View
self.add_menu_item_helper("DIE:exceptionsview", "Exceptions View", 'DIE Exceptions View',
self.DIE_show_exceptionsview_Handler(self), self.icon_list["exception_view"], '')
# Value View
self.add_menu_item_helper("DIE:valueview", "Value View", 'DIE Value View',
self.DIE_show_valview_Handler(self), self.icon_list["value_view"], '')
# Function View
self.add_menu_item_helper("DIE:functionview", "Function View", 'DIE Function View',
self.DIE_show_funcview_Handler(self), self.icon_list["function_view"], '')
# Debug a custom scope
self.add_menu_item_helper("DIE:debugcustomscope", "Debug a custom scope", 'DIE Debug a custom scope',
self.DIE_goscope_Handler(self), self.icon_list["debug_scope"], 'Ctrl+Alt+c')
# Debug entire code
self.add_menu_item_helper("DIE:debugentirecode", "Debug entire code", 'DIE Debug entire code',
self.DIE_goall_Handler(self), self.icon_list["debug_scope"], 'Ctrl+Alt+g')
# Debug from current location
self.add_menu_item_helper("DIE:gofromcurrentlocation", "Debug from current location", 'DIE from current location',
self.DIE_gohere_Handler(self), self.icon_list["debug"], 'Ctrl+Alt+f')
# Settings
self.add_menu_item_helper("DIE:settings", "Settings", 'DIE Settings',
self.DIE_show_settings_Handler(self), self.icon_list["settings"], '')
# About
self.add_menu_item_helper("DIE:about", "About", 'DIE About',
self.DIE_show_about_Handler(self), -1, '')
def del_menu_items(self):
for menu_name in self._menus_names:
idaapi.unregister_action(menu_name)
idaapi.delete_menu("DIE")
###########################################################################
# Debugging
def go_here(self):
self.debugAPI.start_debug(idc.here(), None, auto_start=True)
def go_all(self):
self.debugAPI.start_debug(None, None, auto_start=True)
def show_scope_chooser(self):
global chooser
functions = get_functions()
func_list = functions.keys()
chooser = ScopeChooser(func_list)
chooser.Compile()
ok = chooser.Execute()
if ok == 1:
start_func = func_list[chooser.cbStartFunction.value]
start_func_ea = functions[start_func]
end_func = func_list[chooser.cbEndFunction.value]
end_func_ea = functions[end_func]
self.debugAPI.start_debug(start_func_ea, end_func_ea, True)
chooser.Free()
###########################################################################
# DIE DB
def save_db(self):
db_file = idc.AskFile(1, "*.ddb", "Save DIE Db File")
if db_file is None:
return
self.die_db.save_db(db_file)
def load_db(self):
try:
db_file = idc.AskFile(0, "*.ddb", "Load DIE Db File")
if db_file is not None:
self.die_db.load_db(db_file)
if self.die_db is not None:
self.show_db_details()
except DbFileMismatch as mismatch:
idaapi.msg("Error while loading DIE DB: %s\n" % mismatch)
except Exception as ex:
logging.exception("Error while loading DB: %s", ex)
return False
###########################################################################
# Function View
def show_function_view(self):
self.function_view.Show()
###########################################################################
# Value View
def show_value_view(self):
self.value_view.Show()
###########################################################################
# Parser View
def show_parser_view(self):
self.parser_view.Show()
###########################################################################
# Exceptions View
def show_breakpoint_view(self):
self.bp_view.Show()
###########################################################################
# About
def show_about(self):
AboutWindow().exec_()
###########################################################################
# Settings View
def show_settings(self):
DIE.UI.SetupView.Show(self.config_file_name)
###########################################################################
# Show DB Details
def show_db_details(self):
"""
Print DB details
"""
(start_time,
end_time,
filename,
num_of_functions,
num_of_threads,
numof_parsed_val) = self.die_db.get_run_info()
idaapi.msg("Die DB Loaded.\n")
idaapi.msg("Start Time: %s, End Time %s\n" % (ctime(start_time), ctime(end_time)))
idaapi.msg("Functions: %d, Threads: %d\n" % (num_of_functions, num_of_threads))
idaapi.msg("Parsed Values: %d\n" % numof_parsed_val)
###########################################################################
# Mark\Unmark Execution Flow
def mark_exec_flow(self):
"""
Mark \ Unmark execution flow
"""
color = 0x123456
if self.is_marked:
color = None
with sark.ui.Update():
for func_ea in self.die_db.get_function_counter():
try:
sark.Function(func_ea).color = color
except sark.exceptions.SarkNoFunction:
pass
# Swap is_marked value
self.is_marked = not self.is_marked
###########################################################################
# Show CFG
def show_cfg(self):
"""
Show execution Call flow graph
"""
cfg = self.die_db.get_call_graph_complete()
graph = nx.DiGraph()
if not cfg:
idaapi.msg("No CFG to display")
return
for ctxt_node in cfg:
(from_address, to_address) = ctxt_node
graph.add_edge(from_address, to_address)
viewer = sark.ui.NXGraph(graph, "Callgraph for {}".format("Exection CFG"), handler=sark.ui.AddressNodeHandler())
viewer.Show()
def show_logo(self):
"""
Show DIE Logo
"""
idaapi.msg('-----------------------------------------------------\n'
' _________-----_____ \n'
' _____------ __ ----_ \n'
' ___---- ___------ /\ \n'
' ----________ ---- \ \n'
' -----__ | _____) \n'
' __- / /\ \n'
' _______----- ___-- \ /)\ \n'
' ------_______ ---____ \__/ / \n'
' -----__ \ -- _ //\ \n'
' --__--__ \_____/ \_/\ \n'
' ----| / | \n'
' Dynamic | |___________| \n'
' IDA | | ((_(_)| )_) \n'
' Enrichment | \_((_(_)|/(_) \n'
' \ ( \n'
' \_____________) \n'
' D.I.E v0.1 is now loaded, enjoy. \n'
'-----------------------------------------------------\n'
)
class die_plugin_t(plugin_t):
flags = idaapi.PLUGIN_PROC
comment = "Dynamic IDA Enrichment plugin (aka. DIE)"
help = "Help is a matter of trust."
wanted_name = "DIE"
wanted_hotkey = ""
def init(self):
try:
# For Debugging:
#self.die_manager = DieManager(is_dbg_log=True, is_dbg_pause=False, is_dbg_profile=True)
self.die_manager = DieManager()
self.die_manager.add_menu_items()
self.die_manager.show_logo()
return idaapi.PLUGIN_KEEP
except Exception as ex:
idaapi.msg("Failed to initialize DIE. {}\n".format(ex))
self.die_manager.del_menu_items()
del self.die_manager
idaapi.msg("Errors and fun!\n")
return idaapi.PLUGIN_SKIP
def run(self, arg):
pass
def term(self):
with ignored(AttributeError):
if not self.die_manager.die_db.is_saved:
response = idc.AskYN(1, "DIE DB was not saved, Would you like to save it now?")
if response == 1:
self.die_manager.save_db()
self.die_manager.del_menu_items()
def PLUGIN_ENTRY():
return die_plugin_t()
