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

from pygrap import graph_free

import idagrap.ui.helpers.QtShim as QtShim
import idc
import os
from idagrap.config.General import config
from idagrap.patterns.Modules import MODULES
from idagrap.ui.widgets.EditorWidget import EditorWidget
from PyQt5.QtWidgets import QMenu

QMainWindow = QtShim.get_QMainWindow()


class CryptoIdentificationWidget(QMainWindow):
    """Cryptographic identification Widget.

    This is the core of the the cryptographic identification widget.

    Attributes:
        cc (ClassCollection): Collection of many classes.
        parent (QWidget): The parent QWidget.
        central_widget (QWidget): QWidget of this widget.
        signature_widget (QWidget): Table for the found patterns.
        name (str): Name of the widget.
        icon (QIcon): Icon for the widget.
        scanGraphAction (QAction): Toolbar action.
        signature_tree (QTreeWidget): Tree of the signature table.
        qtreewidgetitems_to_addresses (QTreeWidgetItem:int dict): Dictionary
                    that links together QtTreeWidgetItem and the start
                    address of the related match pattern.

    Arguments:
        parent (QWidget): The parent QWidget.

    """

    def __init__(self, parent):
        """Initialization."""
        # Initialization
        self.cc = parent.cc
        self.cc.QMainWindow.__init__(self)
        #print "[|] loading CryptoIdentificationWidget"

        # Enable access to shared IDAscope modules
        self.parent = parent
        self.name = "Pattern Search"
        self.icon = self.cc.QIcon(config['icons_path'] + "icons8-search.png")
        self.color = False
        self.show_all = False

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

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

        # Signature widget
        self._createSignatureWidget()

        # Layout and fill the widget
        crypto_layout = self.cc.QVBoxLayout()
        splitter = self.cc.QSplitter(self.cc.QtCore.Qt.Vertical)
        q_clean_style = self.cc.QStyleFactory.create('Plastique')
        splitter.setStyle(q_clean_style)
        splitter.addWidget(self.signature_widget)
        crypto_layout.addWidget(splitter)

        self.central_widget.setLayout(crypto_layout)

    def _createToolbar(self):
        """
        Creates the toolbar, containing buttons to control the widget.
        """
        self._createScanGraphAction()
        self._createMatchGraphAction()
        self._createShowAllAction()
        self._createColoringAction()

        self.toolbar = self.addToolBar('Crypto Identification Toolbar')
        self.toolbar.setMovable(False)

        self.toolbar.addAction(self.scanGraphAction)
        self.toolbar.addAction(self.matchGraphAction)
        self.toolbar.addAction(self.showAllAction)
        self.toolbar.addAction(self.coloringAction)

    def _createScanGraphAction(self):
        """
        Create an action for the scan button of the toolbar and connect it.
        """
        # Action
        self.scanGraphAction = 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.scanGraphAction.triggered.connect(self._onScanGraphBouttonClickedThread)
        
    def _createMatchGraphAction(self):
        """
        Create an action for the match button of the toolbar and connect it.
        """
        # Action
        self.matchGraphAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-search.png"),
            "Match the CFG against patterns",
            self
        )

        self.matchGraphAction.triggered.connect(self._onMatchGraphBouttonClickedThread)

    def _createShowAllAction(self):
        """
        Create an action for the coloring button of the toolbar and connect it.
        """
        # Action
        self.showAllAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-eye-50.png"),
            "Show also unmached patterns",
            self
        )
    
        self.showAllAction.setCheckable(True)
        self.showAllAction.toggled.connect(self._showAllBouttonToggled)

    def _showAllBouttonToggled(self, boolean):
        """Handle the different states of the show all button.

        Arguments:
            boolean (bool): State of the button.
        """
        if boolean:
            self.show_all = True
        else:
            self.show_all = False

        self.populateSignatureTree()

    def _createColoringAction(self):
        """
        Create an action for the coloring button of the toolbar and connect it.
        """
        # Action
        self.coloringAction = self.cc.QAction(
            self.cc.QIcon(config['icons_path'] + "icons8-color-palette.png"),
            "Color matches",
            self
        )

        self.coloringAction.setCheckable(True)
        self.coloringAction.toggled.connect(self._coloringBouttonToggled)

    def _coloringBouttonToggled(self, boolean):
        """Handle the different states of the coloring button.

        Arguments:
            boolean (bool): State of the button.
        """
        if boolean:
            self._activateColoringBoutton()
            self.color = True
        else:
            self._deactivateColoringBoutton()
            self.color = False

    def _activateColoringBoutton(self):
        """Action to execute when the coloring button is activated."""
        # Colors generation
        self.cc.CryptoColor.n_assigned_colors = 0
        for ana in self.cc.CryptoIdentifier.get_analyzed_patterns():
            found_patterns = ana.get_found_patterns()
            pcolors = self.cc.CryptoColor.get_patterns_colors()

            # If there is 1 or more matches
            if len(found_patterns) > 0:
                for match_dict_list in found_patterns:
                    for pattern_id, match_dicts in match_dict_list.items():

                        if pattern_id not in pcolors:
                            self.cc.CryptoColor.add_pattern(pattern_id)

                        for match in iter(match_dicts.values()):
                            self.cc.CryptoColor.add_match(match)

        # Highlight matches
        self.cc.CryptoColor.highlight_matches()

        # Update the UI
        self.populateSignatureTree()

    def _deactivateColoringBoutton(self):
        """Action to execute when the coloring button is deactivated."""
        # Remove all patterns colors
        self.cc.CryptoColor.clear()

        # Update the UI
        self.populateSignatureTree()

    def _onScanGraphBouttonClickedThread(self):
        self._onScanGraphBouttonClicked()
        
    def _onMatchGraphBouttonClickedThread(self):
        self._onMatchGraphBouttonClicked()

    def _onScanGraphBouttonClicked(self):
        """
        The logic of the scan button from the toolbar.
        Uses the scanning functions of in CryptoIdentifier and updates the
        elements displaying the results.
        """
        #
        # Crypto Widget
        #

        # Analyzing
        self.cc.CryptoIdentifier.graph.force_extract()
            
            
    def _onMatchGraphBouttonClicked(self):
        """
        The logic of the scan button from the toolbar.
        Uses the scanning functions of in CryptoIdentifier and updates the
        elements displaying the results.
        """
        #
        # Crypto Widget
        #

        # Analyzing
        self.cc.CryptoIdentifier.analyzing()

        # Update the UI
        if self.color:
            # Simulate unclick then click on color button
            self.cc.CryptoColor.clear()
            self._activateColoringBoutton()
        else:
            self.populateSignatureTree()
            

    def _createSignatureWidget(self):
        """
        Create the widget for the signature part.
        """
        # Initizalition of the table
        self.signature_widget = self.cc.QWidget()
        signature_layout = self.cc.QVBoxLayout()
        self.signature_tree = self.cc.QTreeWidget()
        self.signature_tree.setColumnCount(1)
        self.signature_tree.setHeaderLabels(["Found patterns"])

        # Action
        self.signature_tree.itemDoubleClicked.connect(self._onSignatureTreeItemDoubleClicked)

        # Context Menu
        self.signature_tree.setContextMenuPolicy(self.cc.QtCore.Qt.CustomContextMenu)
        self.signature_tree.customContextMenuRequested.connect(self._onSignatureTreeRightClicked)

        signature_layout.addWidget(self.signature_tree)
        self.signature_widget.setLayout(signature_layout)

    def _onSignatureTreeRightClicked(self, position):
        indexes = self.signature_tree.selectedIndexes()
        if len(indexes) > 0:
            level = 0
            index = indexes[0]
            while index.parent().isValid():
                index = index.parent()
                level += 1

        editAction = self.cc.QAction("Edit pattern", self)
        editAction.triggered.connect(self._onSignatureTreeEditPattern)
        self.addAction(editAction)
        
        printAction = self.cc.QAction("Print pattern path", self)
        printAction.triggered.connect(self._onSignatureTreePrintPatternPath)
        self.addAction(printAction)

        menu = QMenu()

        if level == 0:
            menu.addAction(editAction)
            menu.addAction(printAction)

        menu.exec_(self.signature_tree.viewport().mapToGlobal(position))

    def populateSignatureTree(self):
        """
        populate the TreeWidget for display of the signature scanning results.
        """
        # Initialization
        self.signature_tree.clear()
        self.signature_tree.setSortingEnabled(True)
        self.qtreewidgetitems_to_addresses = dict()

        shown_patterns = set()
        # For each analyzed patterns
        for ana in self.cc.CryptoIdentifier.get_analyzed_patterns():
            tree_matches = ana._patterns._matches
            found_patterns = ana.get_found_patterns()
            algo = ana.get_algo()
            patterns = ana.get_patterns()
            colors = self.cc.CryptoColor.get_patterns_colors()

            # If there is 1 or more matches
            if len(found_patterns) > 0:
                if patterns._perform_analysis:
                    algo_info = self.cc.QTreeWidgetItem(self.signature_tree)
                    algo_info.setText(0, algo.get_name())
                    patterns_info = self.cc.QTreeWidgetItem(algo_info)
                else:
                    patterns_info = self.cc.QTreeWidgetItem(self.signature_tree)
                pattern_name = patterns.get_name()
                txt = "%s (%d matches)" % (pattern_name, len(found_patterns))
                patterns_info.setText(0, txt)
                shown_patterns.add(pattern_name)

                for match_dict_list in found_patterns:
                    if patterns._perform_analysis:
                        matches_info = self.cc.QTreeWidgetItem(patterns_info)

                    for pattern_id, match_dicts in match_dict_list.items():
                        if patterns._perform_analysis:
                            pattern_info = self.cc.QTreeWidgetItem(matches_info)
                            pattern_info.setText(0, "%s (%d matches)" % (
                                patterns.get_pattern_name(pattern_id),
                                len(match_dicts.values())
                            ))

                            if pattern_id in colors:
                                pattern_info.setForeground(0, self.cc.QBrush(self.cc.QColor(colors[pattern_id])))

                        for match in iter(match_dicts.values()):
                            if match.get_start_address() is not None:
                                if patterns._perform_analysis:
                                    match_info = self.cc.QTreeWidgetItem(pattern_info)
                                    match_info.setText(0, "0x%x (%d instructions)" % (
                                        match.get_start_address(),
                                        match.get_num_insts()
                                        ))
                                else:
                                    match_info = self.cc.QTreeWidgetItem(patterns_info)
                                
                                if pattern_id in colors and not patterns._perform_analysis:
                                    patterns_info.setForeground(0, self.cc.QBrush(self.cc.QColor(colors[pattern_id])))

                                if patterns._perform_analysis:
                                    matches_info.setText(0, "%s" % idc.GetFunctionName(match.get_start_address()))
                                else:
                                        try:
                                            func_name = idc.get_func_name(match.get_start_address())
                                        except:
                                            func_name = idc.GetFunctionName(match.get_start_address())
                                
                                        match_info.setText(0, "0x%x in %s (%d instructions)" % (
                                        match.get_start_address(),
                                        func_name,
                                        match.get_num_insts()
                                        ))

                                # Add the start address of the match
                                match_info.setData(1, 0, True)
                                match_info.setData(2, 0, hex(match.get_start_address()))
                                self.signature_tree.expandItem(match_info)
                            else:
                                match_info = self.cc.QTreeWidgetItem(patterns_info)
                                match_info.setText(0, "One pattern matched with no extracted node (no getid option)")
                                self.signature_tree.collapseItem(match_info)
                        if patterns._perform_analysis:
                            self.signature_tree.expandItem(pattern_info)
                    if patterns._perform_analysis:
                        self.signature_tree.expandItem(matches_info)
                    
                if len(found_patterns) <= 5:
                    self.signature_tree.expandItem(patterns_info)
                if patterns._perform_analysis:
                    self.signature_tree.expandItem(algo_info)
        
        # Sort by first column, by ascending (usual) order
        self.signature_tree.sortByColumn(0, 0)

        # Adding unmatched patterns
        if self.show_all:
            for pattern_name in self.cc.CryptoIdentifier.pattern_to_path:
                if pattern_name not in shown_patterns:
                    patterns_info = self.cc.QTreeWidgetItem(self.signature_tree)
                    txt = "%s (unmatched)" % (pattern_name)
                    patterns_info.setText(0, txt)

    def _onSignatureTreeItemDoubleClicked(self, item, column):
        """Action for the double clicked.

        Arguments:
            item (QTreeWidgetItem): Item that was clicked.
            column (int): Selected column.
        """
        # Jump to the match address

        if item.data(1, 0):
            addr = int(item.data(2, 0), 16)
            idc.jumpto(addr)


    def _onSignatureTreePrintPatternPath(self, event) :
        getSelected = self.signature_tree.selectedItems()
        pattern_name = getSelected[0].text(0).split(" ")[0]
        if pattern_name in self.cc.CryptoIdentifier.pattern_to_path:
            print("%s: %s" % (pattern_name, self.cc.CryptoIdentifier.pattern_to_path[pattern_name]))
        else:
            print("%s: Could not find where pattern is defined" % pattern_name)

    def _onSignatureTreeEditPattern(self, event) :
        getSelected = self.signature_tree.selectedItems()
        pattern_name = getSelected[0].text(0).split(" ")[0]
        if pattern_name in self.cc.CryptoIdentifier.pattern_to_path:
            path = self.cc.CryptoIdentifier.pattern_to_path[pattern_name]
            editorWidget = EditorWidget(self.parent, path)
            basename=os.path.basename(path)
            self.parent.tabs.addTab(editorWidget, editorWidget.icon, basename)
        else:
            print("%s: Could not find where pattern is defined" % pattern_name)