<commit_before># Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
import re
from collections import Counter

import idc
from PyQt5 import QtGui, QtCore, QtWidgets

import capa.rules
import capa.engine
import capa.ida.helpers
import capa.features.basicblock
from capa.ida.plugin.item import CapaExplorerFunctionItem
from capa.ida.plugin.model import CapaExplorerDataModel

MAX_SECTION_SIZE = 750

# default colors used in views
COLOR_GREEN_RGB = (79, 121, 66)
COLOR_BLUE_RGB = (37, 147, 215)


def calc_level_by_indent(line, prev_level=0):
    """ """
    if not len(line.strip()):
        # blank line, which may occur for comments so we simply use the last level
        return prev_level
    stripped = line.lstrip()
    if stripped.startswith("description"):
        # need to adjust two spaces when encountering string description
        line = line[2:]
    # calc line level based on preceding whitespace
    return len(line) - len(stripped)


def parse_feature_for_node(feature):
    """ """
    description = ""
    comment = ""

    if feature.startswith("- count"):
        # count is weird, we need to handle special
        # first, we need to grab the comment, if exists
        # next, we need to check for an embedded description
        feature, _, comment = feature.partition("#")
        m = re.search(r"- count\(([a-zA-Z]+)\((.+)\s+=\s+(.+)\)\):\s*(.+)", feature)
        if m:
            # reconstruct count without description
            feature, value, description, count = m.groups()
            feature = "- count(%s(%s)): %s" % (feature, value, count)
    elif not feature.startswith("#"):
        feature, _, comment = feature.partition("#")
        feature, _, description = feature.partition("=")

    return map(lambda o: o.strip(), (feature, description, comment))


def parse_node_for_feature(feature, description, comment, depth):
    """ """
    depth = (depth * 2) + 4
    display = ""

    if feature.startswith("#"):
        display += "%s%s\n" % (" " * depth, feature)
    elif description:
        if feature.startswith(("- and", "- or", "- optional", "- basic block", "- not")):
            display += "%s%s" % (" " * depth, feature)
            if comment:
                display += " # %s" % comment
            display += "\n%s- description: %s\n" % (" " * (depth + 2), description)
        elif feature.startswith("- string"):
            display += "%s%s" % (" " * depth, feature)
            if comment:
                display += " # %s" % comment
            display += "\n%sdescription: %s\n" % (" " * (depth + 2), description)
        elif feature.startswith("- count"):
            # count is weird, we need to format description based on feature type, so we parse with regex
            # assume format - count(<feature_name>(<feature_value>)): <count>
            m = re.search(r"- count\(([a-zA-Z]+)\((.+)\)\): (.+)", feature)
            if m:
                name, value, count = m.groups()
                if name in ("string",):
                    display += "%s%s" % (" " * depth, feature)
                    if comment:
                        display += " # %s" % comment
                    display += "\n%sdescription: %s\n" % (" " * (depth + 2), description)
                else:
                    display += "%s- count(%s(%s = %s)): %s" % (
                        " " * depth,
                        name,
                        value,
                        description,
                        count,
                    )
                    if comment:
                        display += " # %s\n" % comment
        else:
            display += "%s%s = %s" % (" " * depth, feature, description)
            if comment:
                display += " # %s\n" % comment
    else:
        display += "%s%s" % (" " * depth, feature)
        if comment:
            display += " # %s\n" % comment

    return display if display.endswith("\n") else display + "\n"


def yaml_to_nodes(s):
    level = 0
    for line in s.splitlines():
        feature, description, comment = parse_feature_for_node(line.strip())

        o = QtWidgets.QTreeWidgetItem(None)

        # set node attributes
        setattr(o, "capa_level", calc_level_by_indent(line, level))

        if feature.startswith(("- and:", "- or:", "- not:", "- basic block:", "- optional:")):
            setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_expression())
        elif feature.startswith("#"):
            setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_comment())
        else:
            setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_feature())

        # set node text
        for (i, v) in enumerate((feature, description, comment)):
            o.setText(i, v)

        yield o


def iterate_tree(o):
    """ """
    itr = QtWidgets.QTreeWidgetItemIterator(o)
    while itr.value():
        yield itr.value()
        itr += 1


def calc_item_depth(o):
    """ """
    depth = 0
    while True:
        if not o.parent():
            break
        depth += 1
        o = o.parent()
    return depth


def build_action(o, display, data, slot):
    """ """
    action = QtWidgets.QAction(display, o)

    action.setData(data)
    action.triggered.connect(lambda checked: slot(action))

    return action


def build_context_menu(o, actions):
    """ """
    menu = QtWidgets.QMenu()

    for action in actions:
        if isinstance(action, QtWidgets.QMenu):
            menu.addMenu(action)
        else:
            menu.addAction(build_action(o, *action))

    return menu


class CapaExplorerRulgenPreview(QtWidgets.QTextEdit):
    def __init__(self, parent=None):
        """ """
        super(CapaExplorerRulgenPreview, self).__init__(parent)

        self.setFont(QtGui.QFont("Courier", weight=QtGui.QFont.Bold))
        self.setLineWrapMode(QtWidgets.QTextEdit.NoWrap)
        self.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarAsNeeded)

    def reset_view(self):
        """ """
        self.clear()

    def load_preview_meta(self, ea, author, scope):
        """ """
        metadata_default = [
            "# generated using capa explorer for IDA Pro",
            "rule:",
            "  meta:",
            "    name: <insert_name>",
            "    namespace: <insert_namespace>",
            "    author: %s" % author,
            "    scope: %s" % scope,
            "    references: <insert_references>",
            "    examples:",
            "      - %s:0x%X" % (capa.ida.helpers.get_file_md5().upper(), ea)
            if ea
            else "      - %s" % (capa.ida.helpers.get_file_md5().upper()),
            "  features:",
        ]
        self.setText("\n".join(metadata_default))

    def keyPressEvent(self, e):
        """ """
        if e.key() == QtCore.Qt.Key_Tab:
            self.insertPlainText(" " * 2)
        else:
            super(CapaExplorerRulgenPreview, self).keyPressEvent(e)


class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):

    updated = QtCore.pyqtSignal()

    def __init__(self, preview, parent=None):
        """ """
        super(CapaExplorerRulgenEditor, self).__init__(parent)

        self.preview = preview

        self.setHeaderLabels(["Feature", "Description", "Comment"])
        self.header().setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents)
        self.header().setStretchLastSection(False)
        self.setExpandsOnDoubleClick(False)
        self.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
        self.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
        self.setStyleSheet("QTreeView::item {padding-right: 15 px;padding-bottom: 2 px;}")<commit_msg>adding support for multi-line tab and SHIFT + Tab<commit_after># Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
import re
from collections import Counter

import idc
from PyQt5 import QtGui, QtCore, QtWidgets

import capa.rules
import capa.engine
import capa.ida.helpers
import capa.features.basicblock
from capa.ida.plugin.item import CapaExplorerFunctionItem
from capa.ida.plugin.model import CapaExplorerDataModel

MAX_SECTION_SIZE = 750

# default colors used in views
COLOR_GREEN_RGB = (79, 121, 66)
COLOR_BLUE_RGB = (37, 147, 215)


def calc_level_by_indent(line, prev_level=0):
    """ """
    if not len(line.strip()):
        # blank line, which may occur for comments so we simply use the last level
        return prev_level
    stripped = line.lstrip()
    if stripped.startswith("description"):
        # need to adjust two spaces when encountering string description
        line = line[2:]
    # calc line level based on preceding whitespace
    return len(line) - len(stripped)


def parse_feature_for_node(feature):
    """ """
    description = ""
    comment = ""

    if feature.startswith("- count"):
        # count is weird, we need to handle special
        # first, we need to grab the comment, if exists
        # next, we need to check for an embedded description
        feature, _, comment = feature.partition("#")
        m = re.search(r"- count\(([a-zA-Z]+)\((.+)\s+=\s+(.+)\)\):\s*(.+)", feature)
        if m:
            # reconstruct count without description
            feature, value, description, count = m.groups()
            feature = "- count(%s(%s)): %s" % (feature, value, count)
    elif not feature.startswith("#"):
        feature, _, comment = feature.partition("#")
        feature, _, description = feature.partition("=")

    return map(lambda o: o.strip(), (feature, description, comment))


def parse_node_for_feature(feature, description, comment, depth):
    """ """
    depth = (depth * 2) + 4
    display = ""

    if feature.startswith("#"):
        display += "%s%s\n" % (" " * depth, feature)
    elif description:
        if feature.startswith(("- and", "- or", "- optional", "- basic block", "- not")):
            display += "%s%s" % (" " * depth, feature)
            if comment:
                display += " # %s" % comment
            display += "\n%s- description: %s\n" % (" " * (depth + 2), description)
        elif feature.startswith("- string"):
            display += "%s%s" % (" " * depth, feature)
            if comment:
                display += " # %s" % comment
            display += "\n%sdescription: %s\n" % (" " * (depth + 2), description)
        elif feature.startswith("- count"):
            # count is weird, we need to format description based on feature type, so we parse with regex
            # assume format - count(<feature_name>(<feature_value>)): <count>
            m = re.search(r"- count\(([a-zA-Z]+)\((.+)\)\): (.+)", feature)
            if m:
                name, value, count = m.groups()
                if name in ("string",):
                    display += "%s%s" % (" " * depth, feature)
                    if comment:
                        display += " # %s" % comment
                    display += "\n%sdescription: %s\n" % (" " * (depth + 2), description)
                else:
                    display += "%s- count(%s(%s = %s)): %s" % (
                        " " * depth,
                        name,
                        value,
                        description,
                        count,
                    )
                    if comment:
                        display += " # %s\n" % comment
        else:
            display += "%s%s = %s" % (" " * depth, feature, description)
            if comment:
                display += " # %s\n" % comment
    else:
        display += "%s%s" % (" " * depth, feature)
        if comment:
            display += " # %s\n" % comment

    return display if display.endswith("\n") else display + "\n"


def yaml_to_nodes(s):
    level = 0
    for line in s.splitlines():
        feature, description, comment = parse_feature_for_node(line.strip())

        o = QtWidgets.QTreeWidgetItem(None)

        # set node attributes
        setattr(o, "capa_level", calc_level_by_indent(line, level))

        if feature.startswith(("- and:", "- or:", "- not:", "- basic block:", "- optional:")):
            setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_expression())
        elif feature.startswith("#"):
            setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_comment())
        else:
            setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_feature())

        # set node text
        for (i, v) in enumerate((feature, description, comment)):
            o.setText(i, v)

        yield o


def iterate_tree(o):
    """ """
    itr = QtWidgets.QTreeWidgetItemIterator(o)
    while itr.value():
        yield itr.value()
        itr += 1


def calc_item_depth(o):
    """ """
    depth = 0
    while True:
        if not o.parent():
            break
        depth += 1
        o = o.parent()
    return depth


def build_action(o, display, data, slot):
    """ """
    action = QtWidgets.QAction(display, o)

    action.setData(data)
    action.triggered.connect(lambda checked: slot(action))

    return action


def build_context_menu(o, actions):
    """ """
    menu = QtWidgets.QMenu()

    for action in actions:
        if isinstance(action, QtWidgets.QMenu):
            menu.addMenu(action)
        else:
            menu.addAction(build_action(o, *action))

    return menu


class CapaExplorerRulgenPreview(QtWidgets.QTextEdit):

    INDENT = " " * 2

    def __init__(self, parent=None):
        """ """
        super(CapaExplorerRulgenPreview, self).__init__(parent)

        self.setFont(QtGui.QFont("Courier", weight=QtGui.QFont.Bold))
        self.setLineWrapMode(QtWidgets.QTextEdit.NoWrap)
        self.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarAsNeeded)

    def reset_view(self):
        """ """
        self.clear()

    def load_preview_meta(self, ea, author, scope):
        """ """
        metadata_default = [
            "# generated using capa explorer for IDA Pro",
            "rule:",
            "  meta:",
            "    name: <insert_name>",
            "    namespace: <insert_namespace>",
            "    author: %s" % author,
            "    scope: %s" % scope,
            "    references: <insert_references>",
            "    examples:",
            "      - %s:0x%X" % (capa.ida.helpers.get_file_md5().upper(), ea)
            if ea
            else "      - %s" % (capa.ida.helpers.get_file_md5().upper()),
            "  features:",
        ]
        self.setText("\n".join(metadata_default))

    def keyPressEvent(self, e):
        """intercept key press events"""
        if e.key() in (QtCore.Qt.Key_Tab, QtCore.Qt.Key_Backtab):
            # apparently it's not easy to implement tabs as spaces, or multi-line tab or SHIFT + Tab
            # so we need to implement it ourselves so we can retain properly formatted capa rules
            # when a user uses the Tab key
            if self.textCursor().selection().isEmpty():
                # single line, only worry about Tab
                if e.key() == QtCore.Qt.Key_Tab:
                    self.insertPlainText(self.INDENT)
            else:
                # multi-line tab or SHIFT + Tab
                cur = self.textCursor()
                select_start_ppos = cur.selectionStart()
                select_end_ppos = cur.selectionEnd()

                scroll_ppos = self.verticalScrollBar().sliderPosition()

                # determine lineno for first selected line, and column
                cur.setPosition(select_start_ppos)
                start_lineno = self.count_previous_lines_from_block(cur.block())
                start_lineco = cur.columnNumber()

                # determine lineno for last selected line
                cur.setPosition(select_end_ppos)
                end_lineno = self.count_previous_lines_from_block(cur.block())

                # now we need to indent or dedent the selected lines. for now, we read the text, modify
                # the lines between start_lineno and end_lineno accordingly, and then reset the view
                # this might not be the best solution, but it avoids messing around with cursor positions
                # to determine the beginning of lines

                plain = self.toPlainText().splitlines()

                if e.key() == QtCore.Qt.Key_Tab:
                    # user Tab, indent selected lines
                    lines_modified = end_lineno - start_lineno
                    first_modified = True
                    change = [self.INDENT + line for line in plain[start_lineno : end_lineno + 1]]
                else:
                    # user SHIFT + Tab, dedent selected lines
                    lines_modified = 0
                    first_modified = False
                    change = []
                    for (lineno, line) in enumerate(plain[start_lineno : end_lineno + 1]):
                        if line.startswith(self.INDENT):
                            if lineno == 0:
                                # keep track if first line is modified, so we can properly display
                                # the text selection later
                                first_modified = True
                            lines_modified += 1
                            line = line[len(self.INDENT) :]
                        change.append(line)

                # apply modifications, and reset view
                plain[start_lineno : end_lineno + 1] = change
                self.setPlainText("\n".join(plain) + "\n")

                # now we need to properly adjust the selection positions, so users don't have to
                # re-select when indenting or dedenting the same lines repeatedly
                if e.key() == QtCore.Qt.Key_Tab:
                    # user Tab, increase increment selection positions
                    select_start_ppos += len(self.INDENT)
                    select_end_ppos += (lines_modified * len(self.INDENT)) + len(self.INDENT)
                elif lines_modified:
                    # user SHIFT + Tab, decrease selection positions
                    if start_lineco not in (0, 1) and first_modified:
                        # only decrease start position if not in first column
                        select_start_ppos -= len(self.INDENT)
                    select_end_ppos -= lines_modified * len(self.INDENT)

                # apply updated selection and restore previous scroll position
                self.set_selection(select_start_ppos, select_end_ppos, len(self.toPlainText()))
                self.verticalScrollBar().setSliderPosition(scroll_ppos)
        else:
            super(CapaExplorerRulgenPreview, self).keyPressEvent(e)

    def count_previous_lines_from_block(self, block):
        """calculate number of lines preceding block"""
        count = 0
        while True:
            block = block.previous()
            if not block.isValid():
                break
            count += block.lineCount()
        return count

    def set_selection(self, start, end, max):
        """set text selection"""
        cursor = self.textCursor()
        cursor.setPosition(start)
        cursor.setPosition(end if end < max else max, QtGui.QTextCursor.KeepAnchor)
        self.setTextCursor(cursor)


class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):

    updated = QtCore.pyqtSignal()

    def __init__(self, preview, parent=None):
        """ """
        super(CapaExplorerRulgenEditor, self).__init__(parent)

        self.preview = preview

        self.setHeaderLabels(["Feature", "Description", "Comment"])
        self.header().setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents)
        self.header().setStretchLastSection(False)
        self.setExpandsOnDoubleClick(False)
        self.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
        self.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
        self.setStyleSheet("QTreeView::item {padding-right: 15 px;padding-bottom: 2 px;}")
{'fraction_non_alphanumeric': 0.11859582542694497, 'fraction_numerical': 0.004506641366223909, 'mean_word_length': 2.0618736383442267, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2095701', 'n_tokens_mistral': 6110, 'n_tokens_neox': 5838, 'n_words': 1579}