path: root/python/examples/snippets/__init__.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import os
import re
from PySide2.QtWidgets import (QLineEdit, QPushButton, QApplication, QTextEdit, QWidget,
    QVBoxLayout, QHBoxLayout, QDialog, QFileSystemModel, QTreeView, QLabel, QSplitter, 
    QInputDialog, QMessageBox, QHeaderView, QMenu, QAction, QKeySequenceEdit,
    QPlainTextEdit)
from PySide2.QtCore import (QDir, QObject, Qt, QFileInfo, QItemSelectionModel, QSettings)
from PySide2.QtGui import (QFont, QFontMetrics, QDesktopServices, QKeySequence)
from binaryninja import user_plugin_path
from binaryninja.plugin import PluginCommand, MainThreadActionHandler
from binaryninja.mainthread import execute_on_main_thread
from binaryninja.log import (log_error, log_info, log_warn, log_alert, log_debug)
from binaryninjaui import (getMonospaceFont, UIAction, UIActionHandler, Menu)
import numbers

snippetPath = os.path.realpath(os.path.join(user_plugin_path(), "..", "snippets"))
try:
    if not os.path.exists(snippetPath):
        os.mkdir(snippetPath)
except IOError:
    log_error("Unable to create %s" % snippetPath)


def includeWalk(dir, includeExt):
    filePaths = []
    for (root, dirs, files) in os.walk(dir):
        for f in files:
            if os.path.splitext(f)[1] in includeExt:
                filePaths.append(os.path.join(root, f))
    return filePaths

def loadSnippetFromFile(snippetPath):
    try:
        snippetText = open(snippetPath, 'r').readlines()
    except:
        return ("", "", "")
    if (len(snippetText) < 3):
        return ("", "", "")
    else:
        qKeySequence = QKeySequence(snippetText[1].strip()[1:])
        if qKeySequence.isEmpty():
            qKeySequence = None
        return (snippetText[0].strip()[1:], 
                qKeySequence,
                ''.join(snippetText[2:])
        )

def executeSnippet(code, context):
    snippetGlobals = {}
    snippetGlobals['current_view'] = context.binaryView
    snippetGlobals['bv'] = context.binaryView
    snippetGlobals['current_function'] = context.function
    if context.function is not None:
        snippetGlobals['current_basic_block'] = context.function.get_basic_block_at(context.address)
    else:
        snippetGlobals['current_basic_block'] = None
    snippetGlobals['current_address'] = context.address
    snippetGlobals['here'] = context.address
    if context.address is not None and isinstance(context.length, numbers.Integral):
        snippetGlobals['current_selection'] = (context.address, context.address+context.length)
    else:
        snippetGlobals['current_selection'] = None
    snippetGlobals['current_llil'] = context.lowLevelILFunction
    snippetGlobals['current_mlil'] = context.mediumLevelILFunction

    exec("from binaryninja import *", snippetGlobals)
    exec(code, snippetGlobals)
    if snippetGlobals['here'] != context.address:
        context.binaryView.file.navigate(context.binaryView.file.view, snippetGlobals['here'])
    if snippetGlobals['current_address'] != context.address:
        context.binaryView.file.navigate(context.binaryView.file.view, snippetGlobals['current_address'])

def makeSnippetFunction(code):
    return lambda context: executeSnippet(code, context)

class Snippets(QDialog):

    def __init__(self, parent=None):
        super(Snippets, self).__init__(parent)
        # Create widgets
        self.setWindowModality(Qt.ApplicationModal)
        self.title = QLabel(self.tr("Snippet Editor"))
        self.saveButton = QPushButton(self.tr("Save"))
        self.closeButton = QPushButton(self.tr("Close"))
        self.clearHotkeyButton = QPushButton(self.tr("Clear Hotkey"))
        self.setWindowTitle(self.title.text())
        self.newFolderButton = QPushButton("New Folder")
        self.deleteSnippetButton = QPushButton("Delete")
        self.newSnippetButton = QPushButton("New Snippet")
        self.edit = QPlainTextEdit()
        self.edit.setPlaceholderText("python code")
        self.resetting = False
        self.columns = 3

        self.keySequenceEdit = QKeySequenceEdit(self)
        self.currentHotkey = QKeySequence()
        self.currentHotkeyLabel = QLabel("")
        self.currentFileLabel = QLabel()
        self.currentFile = ""
        self.snippetDescription = QLineEdit()
        self.snippetDescription.setPlaceholderText("optional description")

        #Set Editbox Size
        font = getMonospaceFont(self)
        self.edit.setFont(font)
        font = QFontMetrics(font)
        self.edit.setTabStopWidth(4 * font.width(' ')); #TODO, replace with settings API

        #Files
        self.files = QFileSystemModel()
        self.files.setRootPath(snippetPath)
        self.files.setNameFilters(["*.py"])

        #Tree
        self.tree = QTreeView()
        self.tree.setModel(self.files)
        self.tree.setSortingEnabled(True)
        self.tree.hideColumn(2)
        self.tree.sortByColumn(0, Qt.AscendingOrder)
        self.tree.setRootIndex(self.files.index(snippetPath))
        for x in range(self.columns):
            #self.tree.resizeColumnToContents(x)
            self.tree.header().setSectionResizeMode(x, QHeaderView.ResizeToContents) 
        treeLayout = QVBoxLayout()
        treeLayout.addWidget(self.tree)
        treeButtons = QHBoxLayout()
        treeButtons.addWidget(self.newFolderButton)
        treeButtons.addWidget(self.newSnippetButton)
        treeButtons.addWidget(self.deleteSnippetButton)
        treeLayout.addLayout(treeButtons)
        treeWidget = QWidget()
        treeWidget.setLayout(treeLayout)

        # Create layout and add widgets
        buttons = QHBoxLayout()
        buttons.addWidget(self.clearHotkeyButton)
        buttons.addWidget(self.keySequenceEdit)
        buttons.addWidget(self.currentHotkeyLabel)
        buttons.addWidget(self.closeButton)
        buttons.addWidget(self.saveButton)

        description = QHBoxLayout()
        description.addWidget(QLabel(self.tr("Description: ")))
        description.addWidget(self.snippetDescription)

        vlayoutWidget = QWidget()
        vlayout = QVBoxLayout()
        vlayout.addLayout(description)
        vlayout.addWidget(self.edit)
        vlayout.addLayout(buttons)
        vlayoutWidget.setLayout(vlayout)

        hsplitter = QSplitter()
        hsplitter.addWidget(treeWidget)
        hsplitter.addWidget(vlayoutWidget)

        hlayout = QHBoxLayout()
        hlayout.addWidget(hsplitter)

        self.showNormal() #Fixes bug that maximized windows are "stuck"
        self.settings = QSettings("Vector35", "Snippet Editor")
        if self.settings.contains("ui/snippeteditor/geometry"):
            self.restoreGeometry(self.settings.value("ui/snippeteditor/geometry"))
        else:
            self.edit.setMinimumWidth(80 * font.averageCharWidth())
            self.edit.setMinimumHeight(30 * font.lineSpacing())

        # Set dialog layout
        self.setLayout(hlayout)

        # Add signals
        self.saveButton.clicked.connect(self.save)
        self.closeButton.clicked.connect(self.close)
        self.clearHotkeyButton.clicked.connect(self.clearHotkey)
        self.tree.selectionModel().selectionChanged.connect(self.selectFile)
        self.newSnippetButton.clicked.connect(self.newFileDialog)
        self.deleteSnippetButton.clicked.connect(self.deleteSnippet)
        self.newFolderButton.clicked.connect(self.newFolder)

        #Read-only until new snippet
        self.readOnly(True)

    @staticmethod
    def registerAllSnippets():
        for action in list(filter(lambda x: x.startswith("Snippets\\"), UIAction.getAllRegisteredActions())):
            if action == "Snippets\\Snippet Editor...":
                continue
            UIActionHandler.globalActions().unbindAction(action)
            Menu.mainMenu("Tools").removeAction(action)
            UIAction.unregisterAction(action)

        for snippet in includeWalk(snippetPath, ".py"):
            snippetKeys = None
            (snippetDescription, snippetKeys, snippetCode) = loadSnippetFromFile(snippet)
            if not snippetDescription:
                actionText = "Snippets\\" + os.path.basename(snippet).rstrip(".py")
            else:
                actionText = "Snippets\\" + snippetDescription
            if snippetCode:
                if snippetKeys == None:
                    UIAction.registerAction(actionText)
                else:
                    UIAction.registerAction(actionText, snippetKeys)
                UIActionHandler.globalActions().bindAction(actionText, UIAction(makeSnippetFunction(snippetCode)))
                Menu.mainMenu("Tools").addAction(actionText, actionText)

    def clearSelection(self):
        self.keySequenceEdit.clear()
        self.currentHotkey = QKeySequence()
        self.currentHotkeyLabel.setText("")
        self.currentFileLabel.setText("")
        self.snippetDescription.setText("")
        self.edit.setPlainText("")
        self.currentFile = ""

    def reject(self):
        self.settings.setValue("ui/snippeteditor/geometry", self.saveGeometry())

        if self.snippetChanged():
            question = QMessageBox.question(self, self.tr("Discard"), self.tr("You have unsaved changes, quit anyway?"))
            if question != QMessageBox.StandardButton.Yes:
                return
        self.accept()

    def newFolder(self):
        (folderName, ok) = QInputDialog.getText(self, self.tr("Folder Name"), self.tr("Folder Name: "))
        if ok and folderName:
            index = self.tree.selectionModel().currentIndex()
            selection = self.files.filePath(index)
            if QFileInfo(selection).isDir():
                QDir(selection).mkdir(folderName)
            else:
                QDir(snippetPath).mkdir(folderName)    

    def selectFile(self, new, old):
        if (self.resetting):
            self.resetting = False
            return
        newSelection = self.files.filePath(new.indexes()[0])
        if QFileInfo(newSelection).isDir():
            self.readOnly(True)
            self.tree.clearSelection()
            self.currentFile = ""
            return

        if old.length() > 0:
            oldSelection = self.files.filePath(old.indexes()[0])
            if not QFileInfo(oldSelection).isDir() and self.snippetChanged():
                question = QMessageBox.question(self, self.tr("Discard"), self.tr("Snippet changed. Discard changes?"))
                if question != QMessageBox.StandardButton.Yes:
                    self.resetting = True
                    self.tree.selectionModel().select(old, QItemSelectionModel.ClearAndSelect | QItemSelectionModel.Rows)
                    return False

        self.currentFile = newSelection
        self.loadSnippet()

    def loadSnippet(self):
        self.currentFileLabel.setText(QFileInfo(self.currentFile).baseName())
        (snippetDescription, snippetKeys, snippetCode) = loadSnippetFromFile(self.currentFile)
        self.snippetDescription.setText(snippetDescription) if snippetDescription else self.snippetDescription.setText("")
        self.keySequenceEdit.setKeySequence(snippetKeys) if snippetKeys else self.keySequenceEdit.setKeySequence(QKeySequence(""))
        self.edit.setPlainText(snippetCode) if snippetCode else self.edit.setPlainText("")
        self.readOnly(False)

    def newFileDialog(self):
        (snippetName, ok) = QInputDialog.getText(self, self.tr("Snippet Name"), self.tr("Snippet Name: "))
        if ok and snippetName:
            if not snippetName.endswith(".py"):
                snippetName += ".py"
            index = self.tree.selectionModel().currentIndex()
            selection = self.files.filePath(index)
            if QFileInfo(selection).isDir():
                path = os.path.join(selection, snippetName)
            else:
                path = os.path.join(snippetPath, snippetName)
                self.readOnly(False)
            open(path, "w").close()
            self.tree.setCurrentIndex(self.files.index(path))
            log_debug("Snippet %s created." % snippetName)

    def readOnly(self, flag):
        self.keySequenceEdit.setEnabled(not flag)
        self.snippetDescription.setReadOnly(flag)
        self.edit.setReadOnly(flag)
        if flag:
            self.snippetDescription.setDisabled(True)
            self.edit.setDisabled(True)
        else:
            self.snippetDescription.setEnabled(True)
            self.edit.setEnabled(True)

    def deleteSnippet(self):
        selection = self.tree.selectedIndexes()[::self.columns][0] #treeview returns each selected element in the row
        snippetName = self.files.fileName(selection)
        question = QMessageBox.question(self, self.tr("Confirm"), self.tr("Confirm deletion: ") + snippetName)
        if (question == QMessageBox.StandardButton.Yes):
            log_debug("Deleting snippet %s." % snippetName)
            self.clearSelection()
            self.files.remove(selection)
            self.registerAllSnippets()

    def snippetChanged(self):
        if (self.currentFile == "" or QFileInfo(self.currentFile).isDir()):
            return False
        (snippetDescription, snippetKeys, snippetCode) = loadSnippetFromFile(self.currentFile)
        if snippetKeys == None and not self.keySequenceEdit.keySequence().isEmpty():
            return True
        if snippetKeys != None and snippetKeys != self.keySequenceEdit.keySequence().toString():
            return True
        return self.edit.toPlainText() != snippetCode or \
               self.snippetDescription.text() != snippetDescription

    def save(self):
        log_debug("Saving snippet %s" % self.currentFile)
        outputSnippet = open(self.currentFile, "w")
        outputSnippet.write("#" + self.snippetDescription.text() + "\n")
        outputSnippet.write("#" + self.keySequenceEdit.keySequence().toString() + "\n")
        outputSnippet.write(self.edit.toPlainText())
        outputSnippet.close()
        self.registerAllSnippets()

    def clearHotkey(self):
        self.keySequenceEdit.clear()

def launchPlugin(context):
    snippets = Snippets()
    snippets.exec_()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    snippets = Snippets()
    snippets.show()
    sys.exit(app.exec_())
else:
    Snippets.registerAllSnippets()
    UIAction.registerAction("Snippets\\Snippet Editor...")
    UIActionHandler.globalActions().bindAction("Snippets\\Snippet Editor...", UIAction(launchPlugin))
    Menu.mainMenu("Tools").addAction("Snippets\\Snippet Editor...", "Snippet")