diff options
Diffstat (limited to 'python/examples/snippets')
| -rw-r--r-- | python/examples/snippets/__init__.py | 444 |
1 files changed, 444 insertions, 0 deletions
diff --git a/python/examples/snippets/__init__.py b/python/examples/snippets/__init__.py new file mode 100644 index 00000000..0f551a06 --- /dev/null +++ b/python/examples/snippets/__init__.py @@ -0,0 +1,444 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +import sys +import os +import re +import codecs +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, QUrl) +from PySide2.QtGui import (QFont, QFontMetrics, QDesktopServices, QKeySequence, QIcon) +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_debug) +from binaryninjaui import (getMonospaceFont, UIAction, UIActionHandler, Menu, DockHandler, + getThemeColor, ThemeColor) +import numbers +from .QCodeEditor import QCodeEditor, PythonHighlighter + +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 = codecs.open(snippetPath, 'r', "utf-8").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 actionFromSnippet(snippetName, snippetDescription): + if not snippetDescription: + shortName = os.path.basename(snippetName) + if shortName.endswith('.py'): + shortName = shortName[:-3] + return "Snippets\\" + shortName + else: + return "Snippets\\" + snippetDescription + + +def executeSnippet(code, context): + snippetGlobals = {} + if context.binaryView == None: + dock = DockHandler.getActiveDockHandler() + if not dock: + log_error("Snippet triggered with no context and no dock handler. This should not happen. Please report reproduction steps if possible.") + return + viewFrame = dock.getViewFrame() + if not viewFrame: + log_error("Snippet triggered with no context and no view frame. Snippets require at least one open binary.") + return + viewInterface = viewFrame.getCurrentViewInterface() + context.binaryView = viewInterface.getData() + snippetGlobals['current_view'] = context.binaryView + snippetGlobals['bv'] = context.binaryView + if not context.function: + if not context.lowLevelILFunction: + if not context.mediumLevelILFunction: + snippetGlobals['current_mlil'] = None + snippetGlobals['current_function'] = None + snippetGlobals['current_llil'] = None + else: + snippetGlobals['current_mlil'] = context.mediumLevelILFunction + snippetGlobals['current_function'] = context.mediumLevelILFunction.source_function + snippetGlobals['current_llil'] = context.mediumLevelILFunction.source_function.llil + else: + snippetGlobals['current_llil'] = context.lowLevelILFunction + snippetGlobals['current_function'] = context.lowLevelILFunction.source_function + snippetGlobals['current_mlil'] = context.lowLevelILFunction.source_function.mlil + else: + snippetGlobals['current_function'] = context.function + snippetGlobals['current_mlil'] = context.function.mlil + snippetGlobals['current_llil'] = context.function.llil + snippetGlobals['current_token'] = context.function.llil + + 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, int): + snippetGlobals['current_selection'] = (context.address, context.address+context.length) + else: + snippetGlobals['current_selection'] = None + snippetGlobals['uicontext'] = context + + 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, context, 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.saveButton.setShortcut(QKeySequence(self.tr("Ctrl+S"))) + self.runButton = QPushButton(self.tr("&Run")) + self.runButton.setShortcut(QKeySequence(self.tr("Ctrl+R"))) + self.closeButton = QPushButton(self.tr("Close")) + self.clearHotkeyButton = QPushButton(self.tr("Clear Hotkey")) + self.setWindowTitle(self.title.text()) + #self.newFolderButton = QPushButton("New Folder") + self.browseButton = QPushButton("Browse Snippets") + self.browseButton.setIcon(QIcon.fromTheme("edit-undo")) + self.deleteSnippetButton = QPushButton("Delete") + self.newSnippetButton = QPushButton("New Snippet") + self.edit = QCodeEditor(HIGHLIGHT_CURRENT_LINE=False, SyntaxHighlighter=PythonHighlighter) + self.edit.setPlaceholderText("python code") + self.resetting = False + self.columns = 3 + self.context = context + + 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.browseButton) + 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.runButton) + 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" + #Because you can't trust QT to do the right thing here + if (sys.platform == "darwin"): + self.settings = QSettings("Vector35", "Snippet Editor") + else: + self.settings = QSettings("Vector 35", "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.runButton.clicked.connect(self.run) + 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) + self.browseButton.clicked.connect(self.browseSnippets) + + if self.settings.contains("ui/snippeteditor/selected"): + selectedName = self.settings.value("ui/snippeteditor/selected") + self.tree.selectionModel().select(self.files.index(selectedName), QItemSelectionModel.ClearAndSelect | QItemSelectionModel.Rows) + if self.tree.selectionModel().hasSelection(): + self.selectFile(self.tree.selectionModel().selection(), None) + self.edit.setFocus() + cursor = self.edit.textCursor() + cursor.setPosition(self.edit.document().characterCount()-1) + self.edit.setTextCursor(cursor) + else: + self.readOnly(True) + else: + 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) + actionText = actionFromSnippet(snippet, 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, "Snippets") + + def clearSelection(self): + self.keySequenceEdit.clear() + self.currentHotkey = QKeySequence() + self.currentHotkeyLabel.setText("") + self.currentFileLabel.setText("") + self.snippetDescription.setText("") + self.edit.clear() + self.tree.clearSelection() + 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 browseSnippets(self): + url = QUrl.fromLocalFile(snippetPath) + QDesktopServices.openUrl(url); + + 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 + if len(new.indexes()) == 0: + self.clearSelection() + self.currentFile = "" + self.readOnly(True) + return + newSelection = self.files.filePath(new.indexes()[0]) + self.settings.setValue("ui/snippeteditor/selected", newSelection) + if QFileInfo(newSelection).isDir(): + self.readOnly(True) + self.clearSelection() + self.currentFile = "" + return + + if old and 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 = codecs.open(self.currentFile, "w", "utf-8") + outputSnippet.write("#" + self.snippetDescription.text() + "\n") + outputSnippet.write("#" + self.keySequenceEdit.keySequence().toString() + "\n") + outputSnippet.write(self.edit.toPlainText()) + outputSnippet.close() + self.registerAllSnippets() + + def run(self): + if self.context == None: + log_warn("Cannot run snippets outside of the UI at this time.") + return + if self.snippetChanged(): + question = QMessageBox.question(self, self.tr("Confirm"), self.tr("You have unsaved changes, must save first. Save?")) + if (question == QMessageBox.StandardButton.No): + return + else: + self.save() + actionText = actionFromSnippet(self.currentFile, self.snippetDescription.text()) + UIActionHandler.globalActions().executeAction(actionText, self.context) + + log_debug("Saving snippet %s" % self.currentFile) + outputSnippet = codecs.open(self.currentFile, "w", "utf-8") + 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(context) + snippets.exec_() + + +if __name__ == '__main__': + app = QApplication(sys.argv) + snippets = Snippets(None) + 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") |