class TriageFilePicker(QWidget): def __init__(self, context): super(TriageFilePicker, self).__init__() self.context = context self.actionHandler = UIActionHandler() self.actionHandler.setupActionHandler(self) self.contextMenu = Menu() self.contextMenuManager = ContextMenuManager(self) layout = QVBoxLayout() layout.setContentsMargins(0, 0, 0, 0) self.model = QFileSystemModel() self.model.setRootPath("") self.model.setFilter(QDir.AllEntries | QDir.Hidden | QDir.System) self.tree = QTreeView(self) self.tree.setModel(self.model) self.tree.setSelectionMode(QAbstractItemView.ExtendedSelection) self.tree.setColumnWidth(0, 500) layout.addWidget(self.tree, 1) self.setLayout(layout) self.tree.doubleClicked.connect(self.onDoubleClick) recentFile = QSettings().value("triage/recentFile", os.path.expanduser("~")) while len(recentFile) > 0: f = self.model.index(recentFile) if f.isValid(): self.tree.scrollTo(f) self.tree.setExpanded(f, True) break parentDir = os.path.dirname(recentFile) if parentDir == recentFile: break recentFile = parentDir self.actionHandler.bindAction("Open Selected Files", UIAction( lambda context: self.openSelectedFiles(), lambda context: self.areFilesSelected())) self.contextMenu.addAction("Open Selected Files", "Open") def contextMenuEvent(self, event): self.contextMenuManager.show(self.contextMenu, self.actionHandler) def onDoubleClick(self, index): self.openSelectedFiles() def openSelectedFiles(self): failedToOpen = [] files = set() for index in self.tree.selectionModel().selectedIndexes(): if self.model.fileInfo(index).isFile(): files.add(self.model.fileInfo(index).absoluteFilePath()) for filename in files: QSettings().setValue("triage/recentFile", filename) f = FileContext.openFilename(filename) if not f: failedToOpen.append(filename) continue for data in f.getAllDataViews(): Settings().set_string("analysis.mode", Settings().get_string("triage.analysisMode"), data) Settings().set_bool("triage.preferSummaryView", True, data) if data.view_type != "Raw": linearSweepMode = Settings().get_string("triage.linearSweep") if linearSweepMode == "none": Settings().set_bool("analysis.linearSweep.autorun", False, data) elif linearSweepMode == "partial": Settings().set_bool("analysis.linearSweep.autorun", True, data) Settings().set_bool("analysis.linearSweep.controlFlowGraph", False, data) elif linearSweepMode == "full": Settings().set_bool("analysis.linearSweep.autorun", True, data) Settings().set_bool("analysis.linearSweep.controlFlowGraph", True, data) self.context.openFileContext(f) if len(failedToOpen) > 0: QMessageBox.critical(self, "Error", "Unable to open:\n" + "\n".join(failedToOpen)) def areFilesSelected(self): return self.tree.selectionModel().hasSelection()
class Snippets(QDialog): def __init__(self, context, parent=None): super(Snippets, self).__init__(parent) # Create widgets self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint) 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") indentation = Settings().get_string("snippets.indentation") if Settings().get_bool("snippets.syntaxHighlight"): self.edit = QCodeEditor(SyntaxHighlighter=Pylighter, delimeter=indentation) else: self.edit = QCodeEditor(SyntaxHighlighter=None, delimeter=indentation) 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.setTabStopDistance( 4 * font.horizontalAdvance(' ')) #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, actionText))) 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 askSave(self): return QMessageBox.question( self, self.tr("Save?"), self.tr("Do you want to save changes to {}?").format( self.currentFileLabel.text()), QMessageBox.Yes | QMessageBox.No | QMessageBox.Cancel) def reject(self): self.settings.setValue("ui/snippeteditor/geometry", self.saveGeometry()) if self.snippetChanged(): save = self.askSave() if save == QMessageBox.Yes: self.save() elif save == QMessageBox.No: self.loadSnippet() elif save == QMessageBox.Cancel: 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(): save = self.askSave() if save == QMessageBox.Yes: self.save() elif save == QMessageBox.No: pass elif save == QMessageBox.Cancel: 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: "), flags=self.windowFlags()) 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(): 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()