#!/usr/bin/env python
# Inspired by IDAscope.


from pygrap import graph_free

import idagrap.ui.helpers.QtShim as QtShim
import idc
import idaapi
from idagrap.config.General import config
from idagrap.patterns.Modules import MODULES
from idagrap.ui.widgets.EditorWidget import EditorWidget
import idagrap.ui.helpers.QtGrapSyntax as syntax
import os

QMainWindow = QtShim.get_QMainWindow()


class PatternGenerationWidget(QMainWindow):
    def __init__(self, parent):
        """Initialization."""
        # Initialization
        self.cc = parent.cc
        self.cc.QMainWindow.__init__(self)
        # print "[|] loading PatternGenerationWidget"
        
        # Enable access to shared IDAscope modules
        self.parent = parent
        self.name = "Pattern Generation"
        self.icon = self.cc.QIcon(config['icons_path'] + "icons8-plus.png")
        self.color = False

        # This widget relies on the crypto identifier
        self.central_widget = self.cc.QWidget()
        self.setCentralWidget(self.central_widget)
        self._createGui()

        self.actionsDefined = False
        self.real_time_option = True
        
        

    def _createGui(self):
        """
        Setup function for the full GUI of this widget.
        """
        # Toolbar
        self._createToolbar()
        
        # Quick pattern text
        self._createQuickPatternTextWidget()

        # Text pattern
        self._createTextWidget()

        # Options widgets
        self._createOptionsWidgets()

        # Layout and fill the widget
        generation_layout = self.cc.QVBoxLayout()

        for options_widget in self.options_widgets:
            generation_layout.addWidget(options_widget)

        hbox = self.cc.QHBoxLayout()
        hbox.addWidget(self.text_qp_widget)
        hbox.addWidget(self.toolbar_qp)
        generation_layout.addLayout(hbox)
        
        generation_layout.addWidget(self.text_widget)

        self.central_widget.setLayout(generation_layout)
        
    def showEvent(self, QShowEvent):
        # Update the UI if the graph is defined
        if not self.actionsDefined and self.cc.PatternGenerator.graph.graph:
            self._createContextActions()
            self._updateContextMenus()
    
    def _createToolbar(self):
        """
        Creates the toolbar, containing buttons to control the widget.
        """
        self.toolbar = self.addToolBar('Pattern Generation Toolbar')
        self.toolbar.setMovable(False)

        self._createLoadGraphAction()
        self.toolbar.addAction(self.loadGraphAction)

        self._createGenerateAction()
        self.toolbar.addAction(self.generateAction)

        self._createFuncAction()
        self.toolbar.addAction(self.funcAction)
        
        self._createResetAction()
        self.toolbar.addAction(self.resetAction)
        
        self._createSaveAction()
        self.toolbar.addAction(self.saveAction)

        self._createOpenAction()
        self.toolbar.addAction(self.openAction)
        
    def _createQuickPatternTextWidget(self):
        self.text_qp_widget = self.cc.QLineEdit()
        self.text_qp_widget.setReadOnly(False)
        
        self.toolbar_qp = self.addToolBar('Pattern Generation Toolbar')
        self._createGenerateQuickPatternAction()
        self.toolbar_qp.addAction(self.generateQuickPatternAction)
        self.text_qp_widget.returnPressed.connect(self._onGenerateQuickPatternButtonClicked)

    def _createTextWidget(self):
        self.text_widget = self.cc.QTextEdit()
        self.text_widget.setReadOnly(False)
        self.text_widget.setFontFamily("Monospace")
        self.highlight = syntax.PythonHighlighter(self.text_widget.document())

    def _createOptionsWidgets(self):
        self.options_widgets = []

        self.real_time_check = self.cc.QCheckBox("Automatically update the pattern")
        self.real_time_check.setChecked(True)
        self.real_time_check.stateChanged.connect(self._real_time_check_option_trigger)
        self.options_widgets.append(self.real_time_check)
        
        self.generic_arguments_check = self.cc.QCheckBox("Generic arguments")
        self.generic_arguments_check.stateChanged.connect(self._generic_arguments_option_trigger)
        self.options_widgets.append(self.generic_arguments_check)

        self.lighten_memory_ops_check = self.cc.QCheckBox("Lighten memory handling operations")
        self.lighten_memory_ops_check.stateChanged.connect(self._lighten_memory_ops_option_trigger)
        self.options_widgets.append(self.lighten_memory_ops_check)

        self.std_jmp_check = self.cc.QCheckBox("Standardize jump operations")
        self.std_jmp_check.stateChanged.connect(self._std_jmp_check_option_trigger)
        self.options_widgets.append(self.std_jmp_check)

        self.factorize_check = self.cc.QCheckBox("Factorize")
        self.factorize_check.stateChanged.connect(self._factorize_check_option_trigger)
        self.options_widgets.append(self.factorize_check)

    def _createOpenAction(self):
        """
        Create an action for the open button of the toolbar and connect it.
        """
        # Action
        self.openAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-edit-property-52.png"),
            "Open pattern file for editing",
            self
        )

        self.openAction.triggered.connect(self._onOpenClicked)

    def _onOpenClicked(self):
        options = self.cc.QFileDialog.Options()
        filename, _ = self.cc.QFileDialog.getOpenFileName(self, "Save pattern file (.grapp files in %APPDATA%\IDAgrap\patterns will be parsed as patterns)", self.default_filepath(), "Grap pattern (*.grapp)", options=options)
        if filename:            
            editorWidget = EditorWidget(self.parent, filename)
            basename=os.path.basename(filename)
            self.parent.tabs.addTab(editorWidget, editorWidget.icon, basename)
 
    def _generic_arguments_option_trigger(self, state):
        self.cc.PatternGenerator.generic_arguments_option = (state == 2)
        self._render()

    def _lighten_memory_ops_option_trigger(self, state):
        self.cc.PatternGenerator.lighten_memory_ops_option = (state == 2)
        self._render()

    def _std_jmp_check_option_trigger(self, state):
        self.cc.PatternGenerator.std_jmp_option = (state == 2)
        self._render()

    def _factorize_check_option_trigger(self, state):
        self.cc.PatternGenerator.factorize_option = (state == 2)
        self._render()

    def _real_time_check_option_trigger(self, state):
        self.real_time_option = (state == 2)
        if self.real_time_option:
            self._render()
            self._enable_options()
        self.generateAction.setEnabled(not self.real_time_option)

    def _createLoadGraphAction(self):
        """
        Create an action for the load graph button of the toolbar and connect it.
        """
        # Action
        self.loadGraphAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-fingerprint-scan.png"),
            "Load the Control Flow Graph from IDA (might take some time)",
            self
        )

        self.loadGraphAction.triggered.connect(self._onLoadGraphButtonClickedThread)

    def _createGenerateAction(self):
        # Action
        self.generateAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-workflow.png"),
            "Generate a pattern (enabled only if you disable the \"Auto update\" option)",
            self
        )
        self.generateAction.setEnabled(False)

        self.generateAction.triggered.connect(self._onGenerateButtonClicked)
        
    def _createGenerateQuickPatternAction(self):
        # Action
        self.generateQuickPatternAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-workflow.png"),
            "Generate a pattern from this short pattern field (for instance: xor->add->xor)",
            self
        )

        self.generateQuickPatternAction.triggered.connect(self._onGenerateQuickPatternButtonClicked)

    def _createFuncAction(self):
        # Action
        self.funcAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-function-mac-32.png"),
            "Target whole current function",
            self
        )

        self.funcAction.triggered.connect(self._onFuncButtonClicked)

    def _createResetAction(self):
        # Action
        self.resetAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-delete.png"),
            "Reset the pattern",
            self
        )

        self.resetAction.triggered.connect(self._onResetButtonClicked)
        
    def _createSaveAction(self):
        # Action
        self.saveAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-add-file.png"),
            "Save the pattern to disk",
            self
        )

        self.saveAction.triggered.connect(self._onSaveButtonClicked)

    def _createContextActions(self): 
        actions = [
            ("grap:pg:set_root", None, "[grap] Set root node", self._onSetRootNode),
            ("grap:pg:add_target", None, "[grap] Add target node", self._onAddTargetNode),
            ("grap:pg:match_default", config['icons_path'] + "icons8-asterisk-24.png", "[grap] Default match (apply options)", self._onSetMatchDefault),
            ("grap:pg:match_full", None, "[grap] Full match", self._onSetMatchFull),
            ("grap:pg:match_opcode_arg1", None, "[grap] Opcode+arg1", self._onSetMatchOpcodeArg1),
            ("grap:pg:match_opcode_arg2", None, "[grap] Opcode+arg2", self._onSetMatchOpcodeArg2),
            ("grap:pg:match_opcode_arg3", None, "[grap] Opcode+arg3", self._onSetMatchOpcodeArg3),
            ("grap:pg:match_opcode", None, "[grap] Opcode", self._onSetMatchOpcode),
            ("grap:pg:match_wildcard", None, "[grap] Wildcard: *", self._onSetMatchWildcard),
            ("grap:pg:remove_target", config['icons_path'] + "icons8-delete.png", "[grap] Remove target node", self._onRemoveTargetNode)
        ]

        for actionId, icon_path, text, method in (a for a in actions):
            if icon_path is not None and icon_path != "":
                icon_number = idaapi.load_custom_icon(icon_path)
                # Describe the action
                action_desc = idaapi.action_desc_t(
                    actionId,  # The action name. This acts like an ID and must be unique
                    text,  # The action text.
                    PatternGenerationHandler(method), # The action handler.
                    None,
                    None,
                    icon_number)  
            else:
                # Describe the action
                action_desc = idaapi.action_desc_t(
                    actionId,  # The action name. This acts like an ID and must be unique
                    text,  # The action text.
                    PatternGenerationHandler(method)) # The action handler.  

            # Register the action
            idaapi.register_action(action_desc)

        self.actionsDefined = True

    def _updateContextMenus(self):
        self.hooks = PatternGenerationHooks(self.cc)
        self.hooks.hook()

    def _render(self):
        self.updateWantedName()
        self.text_widget.setText(self.cc.PatternGenerator.generate(auto=True))
        
    def _render_if_real_time(self):
        if self.real_time_option:
            self._render()
            self._enable_options()

    def _onSetRootNode(self):
        try:
            self.cc.PatternGenerator.setRootNode(idc.get_screen_ea())
        except:
            self.cc.PatternGenerator.setRootNode(idc.ScreenEA())

        self._render_if_real_time()

    def _onAddTargetNode(self):
        try:
            self.cc.PatternGenerator.addTargetNode(idc.get_screen_ea())
        except:
            self.cc.PatternGenerator.addTargetNode(idc.ScreenEA())

        self._render_if_real_time()
    
    def setMatchType(self, type):
        try:
            selection, begin, end = None, None, None
            err = idaapi.read_selection(selection, begin, end)
            if err and selection:
                for ea in range(begin, end+1):
                    self.cc.PatternGenerator.setMatchType(ea, type)
            else:
                self.cc.PatternGenerator.setMatchType(idc.get_screen_ea(), type)  
        except:
            self.cc.PatternGenerator.setMatchType(idc.ScreenEA(), type)

        self._render_if_real_time() 
    
    def _onSetMatchDefault(self):
        self.setMatchType("match_default")
    
    def _onSetMatchFull(self):
        self.setMatchType("match_full")
        
    def _onSetMatchOpcodeArg1(self):
        self.setMatchType("match_opcode_arg1")
        
    def _onSetMatchOpcodeArg2(self):
        self.setMatchType("match_opcode_arg2")
        
    def _onSetMatchOpcodeArg3(self):
        self.setMatchType("match_opcode_arg3")
        
    def _onSetMatchOpcode(self):
        self.setMatchType("match_opcode")
        
    def _onSetMatchWildcard(self):
        self.setMatchType("match_wildcard") 
     
    def _onRemoveTargetNode(self):
        try:
            self.cc.PatternGenerator.removeTargetNode(idc.get_screen_ea())
        except:
            self.cc.PatternGenerator.removeTargetNode(idc.ScreenEA())

        self._render_if_real_time()

    def _onLoadGraphButtonClickedThread(self):
        self._onLoadGraphButtonClicked()

    def _onLoadGraphButtonClicked(self):
        existing = False
        if self.cc.PatternGenerator.graph.graph:
            existing = True
    
        # Analyzing
        self.cc.PatternGenerator.graph.force_extract()

        # Update the UI
        if not self.actionsDefined:
            self._createContextActions()
            self._updateContextMenus()
            
        # UI information
        if existing:
            print("[I] CFG updated. You can now define your pattern's root node and target nodes (right click on an instruction in IDA View).")
        else:
            print("[I] CFG loaded. You can now define your pattern's root node and target nodes (right click on an instruction in IDA View).")

    def _onGenerateQuickPatternButtonClicked(self):
        print("[I] Generation of quick pattern")
        self.text_widget.setText(self.cc.PatternGenerator.generate_quick_pattern(self.text_qp_widget.text()))
        self.generateAction.setEnabled(True)
        self._disable_options()
        
    def _onGenerateButtonClicked(self):
        print("[I] Generation of pattern")
        self._render()
        self._enable_options()

    def _onFuncButtonClicked(self):
        if not self.cc.PatternGenerator.graph.graph:
            print("WARNING: Unloaded CFG. Make sure to first \"Load the CFG\"")
            return

        ea = idaapi.get_screen_ea()
        if ea:
            func = idaapi.ida_funcs.get_func(ea)
            if func:
                if self.cc.PatternGenerator.rootNode is None:
                    print("[I] Adding root node as function entrypoint: %x", func.start_ea)
                    self.cc.PatternGenerator.setRootNode(func.start_ea)

                print("[I] Adding nodes to cover whole function")
                flowchart = idaapi.FlowChart(func)
                for bb in flowchart:
                    last_inst_addr = idc.prev_head(bb.end_ea)
                    self.cc.PatternGenerator.addTargetNode(last_inst_addr)

                self._render_if_real_time() 

    def _onResetButtonClicked(self):
        print("[I] Reset pattern")
        self.cc.PatternGenerator.resetPattern()
        self.text_widget.clear()
        self._enable_options()

    def updateWantedName(self):
        pattern_text = self.text_widget.toPlainText()
        lines = pattern_text.split("\n")

        if len(lines) >= 1:
            l = lines[0]
            s = l.strip().split(" ")
            if len(s) >= 2:
                if "graph" in s[0].lower():
                    fn = s[1]
                    if len(fn) >= 1:
                        self.cc.PatternGenerator.wantedName = str(s[1])
        

    def default_filepath(self):
        if "user_patterns_path" in config:
            default_path = config["user_patterns_path"]
        else:
            default_path = config["patterns_path"] + os.path.sep + "test"+ os.path.sep + "misc" + os.path.sep + "files"
        default_filepath = default_path + os.path.sep + self.cc.PatternGenerator.wantedName + ".grapp"
        return default_filepath
 
    def _onSaveButtonClicked(self):
        self.updateWantedName()
        pattern_text = self.text_widget.toPlainText()
        
        if len(pattern_text.strip()) == 0:
            print("WARNING: Pattern is empty.")
            return
    
        #print("[I] Saving pattern")
        options = self.cc.QFileDialog.Options()
        #options |= self.cc.QFileDialog.DontUseNativeDialog
           
        filename, _ = self.cc.QFileDialog.getSaveFileName(self, "Save pattern file (.grapp files in %APPDATA%\IDAgrap\patterns will be parsed as patterns)", self.default_filepath(), "Grap pattern (*.grapp)", options=options)
        if filename:            
            try:
                f = open(filename, "w")
                f.write(pattern_text)
                f.close()
            except Exception as e:
                print("WARNING:", e)
                
    def _disable_options(self):
        self.real_time_check.setEnabled(False)
        self.generic_arguments_check.setEnabled(False)
        self.lighten_memory_ops_check.setEnabled(False)
        self.std_jmp_check.setEnabled(False)
        self.factorize_check.setEnabled(False)

    def _enable_options(self):    
        self.real_time_check.setEnabled(True)
        self.generic_arguments_check.setEnabled(True)
        self.lighten_memory_ops_check.setEnabled(True)
        self.std_jmp_check.setEnabled(True)
        self.factorize_check.setEnabled(True)


class PatternGenerationHandler(idaapi.action_handler_t):
    def __init__(self, callback):
        idaapi.action_handler_t.__init__(self)
        self.callback = callback

    def activate(self, ctx):
        self.callback()

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


class PatternGenerationHooks(idaapi.UI_Hooks):
    def __init__(self, cc):
        idaapi.UI_Hooks.__init__(self)
        self.cc = cc
        self.selected_icon_number = idaapi.load_custom_icon(config['icons_path'] + "icons8-asterisk-24.png")

    def populating_widget_popup(self, form, popup):
        pass

    def finish_populating_widget_popup(self, form, popup):
        try:
            b = idaapi.get_widget_type(form) == idaapi.BWN_DISASM
        except:
            b = idaapi.get_tform_type(form) == idaapi.BWN_DISASM
    
        if b:
            # Add separator
            idaapi.attach_action_to_popup(form, popup, None, None)

            # Add actions
            try:
                currentAddress = idc.get_screen_ea()
            except:
                currentAddress = idc.ScreenEA()

            #if currentAddress in [node.node_id for node in self.cc.PatternGenerator.targetNodes]:
            if currentAddress in self.cc.PatternGenerator.coloredNodes:
                idaapi.attach_action_to_popup(form, popup, "grap:pg:match_default", None)
                idaapi.attach_action_to_popup(form, popup, "grap:pg:match_full", None)
                idaapi.update_action_label("grap:pg:match_full", self.cc.PatternGenerator.preview_match(currentAddress, "[grap] Full match", "match_full"))
                idaapi.attach_action_to_popup(form, popup, "grap:pg:match_opcode_arg1", None)
                idaapi.update_action_label("grap:pg:match_opcode_arg1", self.cc.PatternGenerator.preview_match(currentAddress, "[grap] Opcode+arg1", "match_opcode_arg1"))
                idaapi.attach_action_to_popup(form, popup, "grap:pg:match_opcode_arg2", None)
                idaapi.update_action_label("grap:pg:match_opcode_arg2", self.cc.PatternGenerator.preview_match(currentAddress, "[grap] Opcode+arg2", "match_opcode_arg2"))
                idaapi.attach_action_to_popup(form, popup, "grap:pg:match_opcode_arg3", None)
                idaapi.update_action_label("grap:pg:match_opcode_arg3", self.cc.PatternGenerator.preview_match(currentAddress, "[grap] Opcode+arg3", "match_opcode_arg3"))
                idaapi.attach_action_to_popup(form, popup, "grap:pg:match_opcode", None)
                idaapi.update_action_label("grap:pg:match_opcode", self.cc.PatternGenerator.preview_match(currentAddress, "[grap] Opcode", "match_opcode"))
                idaapi.attach_action_to_popup(form, popup, "grap:pg:match_wildcard", None)
                idaapi.attach_action_to_popup(form, popup, "grap:pg:remove_target", None)
                
                for type in ["match_default", "match_full", "match_opcode_arg1", "match_opcode_arg2", "match_opcode_arg3", "match_opcode", "match_wildcard"]:
                    idaapi.update_action_icon("grap:pg:"+type, -1)
                
                if currentAddress not in self.cc.PatternGenerator.targetNodeType:
                    type = "match_default"
                else:
                    type = self.cc.PatternGenerator.targetNodeType[currentAddress]
                idaapi.update_action_icon("grap:pg:"+type, self.selected_icon_number)
                    
            elif self.cc.PatternGenerator.rootNode is None or currentAddress != self.cc.PatternGenerator.rootNode.node_id:
                idaapi.attach_action_to_popup(form, popup, "grap:pg:set_root", None)
                idaapi.attach_action_to_popup(form, popup, "grap:pg:add_target", None)