class MainWindow(QMainWindow):
def __init__(self, parent=None, flags=Qt.WindowFlags()):
super(MainWindow, self).__init__(parent, flags)
self.dayWidth = 70
self.ui = Ui_MainWindow()
self.ui.setupUi(self)
self.initModel()
self.initActions()
self.initItemDelegate()
self.initGrid()
leftView = self.ui.ganttView.leftView()
leftView.setColumnHidden(1, True)
leftView.setColumnHidden(2, True)
leftView.setColumnHidden(3, True)
leftView.setColumnHidden(4, True)
leftView.setColumnHidden(5, True)
leftView.header().setStretchLastSection(True)
self.ui.ganttView.leftView().customContextMenuRequested.connect(
self.showContextMenu)
self.ui.ganttView.selectionModel().selectionChanged.connect(
self.enableActions)
self.ui.ganttView.graphicsView().clicked.connect(self.slotClicked)
self.ui.ganttView.graphicsView().qrealClicked.connect(
self.slotDoubleClicked)
def initModel(self):
self.model = QStandardItemModel(0, 6, self)
self.model.setHeaderData(0, Qt.Horizontal, "Task")
self.ui.ganttView.setModel(self.model)
self.l = QWidget(self)
self.l.setWindowTitle("Legend")
self.l.show()
##self.l.setModel(self.model)
self.constraintModel = ConstraintModel(self)
self.ui.ganttView.setConstraintModel(self.constraintModel)
def initActions(self):
self.newEntryAction = QAction("New entry", self)
self.newEntryAction.setShortcut(QKeySequence.New)
self.newEntryAction.triggered.connect(self.addNewEntry)
self.removeEntryAction = QAction("Remove entry", self)
self.removeEntryAction.setShortcut(QKeySequence.Delete)
self.removeEntryAction.triggered.connect(self.removeEntry)
self.demoAction = QAction("Demo entry", self)
self.demoAction.triggered.connect(self.addDemoEntry)
self.printAction = QAction("Print Preview...", self)
self.printAction.triggered.connect(self.printPreview)
self.zoomInAction = QAction("Zoom In", self)
self.zoomInAction.setShortcut(QKeySequence.ZoomIn)
self.zoomInAction.triggered.connect(self.zoomIn)
self.zoomOutAction = QAction("Zoom Out", self)
self.zoomOutAction.setShortcut(QKeySequence.ZoomOut)
self.zoomOutAction.triggered.connect(self.zoomOut)
self.ui.ganttView.leftView().setContextMenuPolicy(Qt.CustomContextMenu)
self.ui.ganttView.leftView().addAction(self.newEntryAction)
self.ui.ganttView.leftView().addAction(self.removeEntryAction)
menuBar = QMenuBar(self)
self.setMenuBar(menuBar)
entryMenu = menuBar.addMenu("Entry")
entryMenu.addAction(self.newEntryAction)
entryMenu.addAction(self.removeEntryAction)
entryMenu.addSeparator()
entryMenu.addAction(self.demoAction)
entryMenu.addSeparator()
entryMenu.addAction(self.printAction)
zoomMenu = menuBar.addMenu("Zoom")
zoomMenu.addAction(self.zoomInAction)
zoomMenu.addAction(self.zoomOutAction)
##self.enableActions(QItemSelection())
def initItemDelegate(self):
delegate = EntryDelegate(self.constraintModel, self)
self.ui.ganttView.leftView().setItemDelegate(delegate)
def initGrid(self):
self.grid = DateTimeGrid()
self.grid.setDayWidth(self.dayWidth)
self.ui.ganttView.setGrid(self.grid)
def showContextMenu(self, pos):
if not self.ui.ganttView.leftView().indexAt(pos).isValid():
self.ui.ganttView.selectionModel().clearSelection()
menu = QMenu(self.ui.ganttView.leftView())
menu.addAction(self.newEntryAction)
menu.addAction(self.removeEntryAction)
menu.exec_(self.ui.ganttView.leftView().viewport().mapToGlobal(pos))
def enableActions(self, selected):
if len(selected.indexes()) == 0:
self.newEntryAction.setEnabled(True)
self.removeEntryAction.setEnabled(False)
return
selectedIndex = selected.indexes()[0]
dataType = self.model.data(self.model.index(selectedIndex.row(), 1))
if dataType in [KDGantt.TypeEvent, KDGantt.TypeTask]:
self.newEntryAction.setEnabled(False)
self.removeEntryAction.setEnabled(True)
return
self.newEntryAction.setEnabled(True)
self.removeEntryAction.setEnabled(True)
def addNewEntry(self):
dialog = EntryDialog(self.model)
dialog.setWindowTitle("New Entry")
if dialog.exec_() == QDialog.Rejected:
dialog = None
return
selectedIndexes = self.ui.ganttView.selectionModel().selectedIndexes()
if len(selectedIndexes) > 0:
parent = selectedIndexes[0]
else:
parent = QModelIndex()
if not self.model.insertRow(self.model.rowCount(parent), parent):
return
row = self.model.rowCount(parent) - 1
if row == 0 and parent.isValid():
self.model.insertColumns(self.model.columnCount(paren), 5, parent)
self.model.setData(self.model.index(row, 0, parent), dialog.name())
self.model.setData(self.model.index(row, 1, parent), dialog.type())
if dialog.type() != KDGantt.TypeSummary:
self.model.setData(self.model.index(row, 2, parent),
dialog.startDate(), KDGantt.StartTimeRole)
self.model.setData(self.model.index(row, 3, parent),
dialog.endDate(), KDGantt.EndTimeRole)
self.model.setData(self.model.index(row, 4, parent),
dialog.completion())
self.model.setData(self.model.index(row, 5, parent), dialog.legend())
self.addConstraint(dialog.depends(), self.model.index(row, 0, parent))
self.setReadOnly(self.model.index(row, 0, parent), dialog.readOnly())
dialog = None
def setReadOnly(self, index, readOnly):
row = index.row()
parent = index.parent()
for column in range(0, 5):
item = self.model.itemFromIndex(
self.model.index(row, column, parent))
flags = None
if readOnly:
flags = item.flags() & ~Qt.ItemIsEditable
else:
flags = item.flags() | Qt.ItemIsEditable
item.setFlags(flags)
def addConstraint(self, index1, index2):
if not index1.isValid() or not index2.isValid():
return
c = Constraint(index1, index2)
self.ui.ganttView.constraintModel().addConstraint(c)
def addConstraintFromItem(self, item1, item2):
self.addConstraint(self.model.indexFromItem(item1),
self.model.indexFromItem(item2))
def removeEntry(self):
selectedIndexes = self.ui.ganttView.selectionModel().selectedIndexes()
if len(selectedIndexes) > 0:
index = selectedIndexes[0]
else:
index = QModelIndex()
if not index.isValid():
return
self.model.removeRow(index.row(), index.parent())
def addDemoEntry(self):
softwareRelease = MyStandardItem("Software Release")
codeFreeze = MyStandardItem("Code Freeze")
codeFreeze.setData(KDGantt.TextPositionRole,
StyleOptionGanttItem.Right)
packaging = MyStandardItem("Packaging")
upload = MyStandardItem("Upload")
testing = MyStandardItem("Testing")
updateDocumentation = MyStandardItem("Update Documentation")
now = QDateTime.currentDateTime()
softwareRelease.appendRow([
codeFreeze,
MyStandardItem(KDGantt.TypeEvent),
MyStandardItem(now, KDGantt.StartTimeRole)
])
softwareRelease.appendRow([
packaging,
MyStandardItem(KDGantt.TypeTask),
MyStandardItem(now.addDays(5), KDGantt.StartTimeRole),
MyStandardItem(now.addDays(10), KDGantt.EndTimeRole)
])
softwareRelease.appendRow([
upload,
MyStandardItem(KDGantt.TypeTask),
MyStandardItem(
now.addDays(10).addSecs(2 * 60 * 60), KDGantt.StartTimeRole),
MyStandardItem(now.addDays(11), KDGantt.EndTimeRole)
])
softwareRelease.appendRow([
testing,
MyStandardItem(KDGantt.TypeTask),
MyStandardItem(now.addSecs(3 * 60 * 60), KDGantt.StartTimeRole),
MyStandardItem(now.addDays(5), KDGantt.EndTimeRole)
])
softwareRelease.appendRow([
updateDocumentation,
MyStandardItem(KDGantt.TypeTask),
MyStandardItem(now.addSecs(3 * 60 * 60), KDGantt.StartTimeRole),
MyStandardItem(now.addDays(3), KDGantt.EndTimeRole)
])
self.model.appendRow(
[softwareRelease,
MyStandardItem(KDGantt.TypeSummary)])
self.addConstraintFromItem(codeFreeze, packaging)
self.addConstraintFromItem(codeFreeze, testing)
self.addConstraintFromItem(codeFreeze, updateDocumentation)
self.addConstraintFromItem(packaging, upload)
self.addConstraintFromItem(testing, packaging)
self.addConstraintFromItem(updateDocumentation, packaging)
def zoomIn(self):
self.dayWidth += 10
if self.dayWidth > 400:
self.grid.setScale(DateTimeGrid.ScaleHour)
self.grid.setDayWidth(self.dayWidth)
def zoomOut(self):
self.dayWidth -= 10
if self.dayWidth < 10:
self.dayWidth = 10
if self.dayWidth <= 400:
self.grid.setScale(DateTimeGrid.ScaleDay)
self.grid.setDayWidth(self.dayWidth)
def printPreview(self):
preview = QLabel(self, Qt.Window)
preview.setAttribute(Qt.WA_DeleteOnClose)
preview.setScaledContents(True)
preview.setWindowTitle("Print Preview")
pix = QPixmap(1000, 300)
pix.fill(Qt.white)
p = QPainter(pix)
p.setRenderHints(QPainter.Antialiasing)
self.ui.ganttView.print_(p, pix.rect())
preview.setPixmap(pix)
preview.show()
def slotClicked(self, index):
self.statusBar().showMessage(
"(%d,%d,_,%s) clicked" %
(index.row(), index.column(), index.model()))
def slotDoubleClicked(self, index):
self.statusBar().showMessage(
"(%d,%d,_,%s) qreal clicked" %
(index.row(), index.column(), index.model()))
class App(QMainWindow):
def __init__(self):
super(App, self).__init__()
self.treeView = None
self.model = None
self.pattern_delegate = None
self.callee_delegate = None
self.sig_trie = None
self.searchResults = None
self.searchIndex = -1
self.findNextAction = None
self.findPrevAction = None
# these two maps are used to make the hyperlinks work
# mapping from href to FunctionNode
self.hrefs_to_funcs = {}
# mapping from FunctionNode to tree view element (QStandardItem)
self.func_node_items = {}
self.init_ui()
def init_ui(self):
self.setWindowTitle('Signature Explorer')
self.resize(1000, 640)
app_icon = QIcon()
app_icon.addFile('icon.ico', QSize(48, 48))
self.setWindowIcon(app_icon)
self.pattern_delegate = PatternDelegate()
self.callee_delegate = CalleesDelegate()
self.treeView = TrieView()
# self.treeView.setAlternatingRowColors(True)
self.model = QStandardItemModel(0, 7, self.treeView)
self.model.setHeaderData(0, Qt.Horizontal, 'Signature')
self.model.setHeaderData(1, Qt.Horizontal, 'Function')
self.model.setHeaderData(2, Qt.Horizontal, 'Callees')
self.model.setHeaderData(3, Qt.Horizontal, 'Offset Extra Pattern')
self.model.setHeaderData(4, Qt.Horizontal, 'Extra Pattern')
self.model.setHeaderData(5, Qt.Horizontal, 'Source Binary')
self.model.setHeaderData(6, Qt.Horizontal, 'ID')
self.treeView.setModel(self.model)
self.treeView.setSelectionBehavior(QAbstractItemView.SelectRows)
self.treeView.setColumnWidth(0, 400)
self.treeView.setColumnWidth(1, 200)
self.treeView.setColumnWidth(2, 250)
self.treeView.setColumnWidth(3, 25)
self.treeView.setColumnWidth(4, 100)
self.treeView.setColumnWidth(5, 200)
self.treeView.setColumnWidth(6, 75)
self.treeView.setItemDelegateForColumn(0, self.pattern_delegate)
self.treeView.setItemDelegateForColumn(2, self.callee_delegate)
self.treeView.setItemDelegateForColumn(4, self.pattern_delegate)
self.treeView.horizontalScrollBar().setEnabled(True)
self.treeView.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)
self.treeView.setEditTriggers(QAbstractItemView.NoEditTriggers)
self.treeView.linkActivated.connect(self.on_func_link_clicked)
# self.treeView.expanded.connect(lambda x: self.treeView.resizeColumnToContents(1))
# self.treeView.collapsed.connect(lambda x: self.treeView.resizeColumnToContents(1))
main_layout = QVBoxLayout()
main_layout.addWidget(self.treeView)
panel = QWidget()
panel.setLayout(main_layout)
self.setCentralWidget(panel)
menuBar = self.menuBar()
fileMenu = QMenu("File")
openAction = QAction("&Open", self)
openAction.setShortcuts(QKeySequence.Open)
openAction.triggered.connect(self.open_file)
fileMenu.addAction(openAction)
closeAction = QAction("&Close", self)
closeAction.setShortcuts(QKeySequence.Close)
closeAction.triggered.connect(self.close_file)
fileMenu.addAction(closeAction)
saveAsAction = QAction("Save As...", self)
saveAsAction.setShortcuts(QKeySequence.Save)
saveAsAction.triggered.connect(self.save_as)
fileMenu.addAction(saveAsAction)
menuBar.addMenu(fileMenu)
editMenu = QMenu("Edit")
findAction = QAction("&Find", self)
findAction.setShortcuts(QKeySequence.Find)
findAction.triggered.connect(self.search)
editMenu.addAction(findAction)
self.findNextAction = QAction("&Find Next", self)
self.findNextAction.setShortcuts(QKeySequence.FindNext)
self.findNextAction.triggered.connect(self.select_next)
self.findNextAction.setEnabled(False)
editMenu.addAction(self.findNextAction)
self.findPrevAction = QAction("&Find Prev", self)
self.findPrevAction.setShortcuts(QKeySequence.FindPrevious)
self.findPrevAction.triggered.connect(self.select_prev)
self.findPrevAction.setEnabled(False)
editMenu.addAction(self.findPrevAction)
menuBar.addMenu(editMenu)
viewMenu = QMenu("View")
expandAction = QAction("&Expand All", self)
expandAction.triggered.connect(self.treeView.expandAll)
viewMenu.addAction(expandAction)
collapseAction = QAction("&Collapse All", self)
collapseAction.triggered.connect(self.treeView.collapseAll)
viewMenu.addAction(collapseAction)
menuBar.addMenu(viewMenu)
def search(self):
query_string, ok = QInputDialog.getText(self, 'Find in Trie',
'Function name')
if not ok or not query_string:
return
self.searchResults = self.model.findItems(
query_string, Qt.MatchContains | Qt.MatchRecursive, 1)
if self.searchResults:
self.findNextAction.setEnabled(True)
self.findPrevAction.setEnabled(True)
self.searchIndex = 0
self.select_next()
else:
self.findNextAction.setEnabled(False)
self.findPrevAction.setEnabled(False)
self.searchIndex = -1
QMessageBox.warning(self, 'Find in Trie', 'No results found')
def select_next(self):
next_item = self.searchResults[self.searchIndex]
self.searchIndex = (self.searchIndex + 1) % len(self.searchResults)
self.select_tree_item(next_item)
def select_prev(self):
prev_item = self.searchResults[self.searchIndex]
self.searchIndex = (self.searchIndex - 1) % len(self.searchResults)
self.select_tree_item(prev_item)
def select_tree_item(self, item):
path = []
while item:
path.insert(0, self.model.indexFromItem(item))
item = item.parent()
# print(path)
for index in path:
self.treeView.setExpanded(index, True)
self.treeView.selectionModel().select(
path[-1],
QItemSelectionModel.ClearAndSelect | QItemSelectionModel.Rows)
self.treeView.scrollTo(path[-1])
def close_file(self):
self.model.removeRows(0, self.model.rowCount())
self.sig_trie = None
self.hrefs_to_funcs = {}
self.func_node_items = {}
def open_file(self):
sig_filter = 'Signature library (*.sig)'
json_zlib_filter = 'Compressed JSON signature library (*.json.zlib)'
json_filter = 'JSON signature library (*.json)'
pkl_filter = 'Pickled signature library (*.pkl)'
fname, filter = QFileDialog.getOpenFileName(self,
'Open file',
filter=';;'.join([
sig_filter,
json_zlib_filter,
json_filter, pkl_filter
]))
if filter and fname:
print('Opening signature library %s' % (fname, ))
if filter == json_zlib_filter:
with open(fname, 'rb') as f:
json_trie = zlib.decompress(f.read()).decode('utf-8')
sig_trie = sig_serialize_json.deserialize(json.loads(json_trie))
elif filter == json_filter:
with open(fname, 'r') as f:
json_trie = f.read()
sig_trie = sig_serialize_json.deserialize(json.loads(json_trie))
elif filter == sig_filter:
with open(fname, 'rb') as f:
fb_trie = f.read()
sig_trie = sig_serialize_fb.SignatureLibraryReader().deserialize(
fb_trie)
elif filter == pkl_filter:
with open(fname, 'rb') as f:
sig_trie = pickle.load(f)
else:
return
self.open_trie(sig_trie, os.path.basename(fname))
def save_as(self):
sig_filter = 'Signature library (*.sig)'
json_zlib_filter = 'Compressed JSON signature library (*.json.zlib)'
json_filter = 'JSON signature library (*.json)'
pkl_filter = 'Pickled signature library (*.pkl)'
fname, filter = QFileDialog.getSaveFileName(self,
'Open file',
filter=';;'.join([
sig_filter,
json_zlib_filter,
json_filter, pkl_filter
]))
if filter == json_zlib_filter:
with open(fname, 'wb') as f:
f.write(
zlib.compress(
sig_serialize_json.serialize(
self.sig_trie).encode('utf-8')))
elif filter == json_filter:
with open(fname, 'w') as f:
json.dump(sig_serialize_json.serialize(self.sig_trie),
f,
indent=4)
elif filter == sig_filter:
with open(fname, 'wb') as f:
f.write(sig_serialize_fb.SignatureLibraryWriter().serialize(
self.sig_trie))
elif filter == pkl_filter:
with open(fname, 'wb') as f:
pickle.dump(self.sig_trie, f)
else:
return
print('Saved as ' + fname)
@staticmethod
def generate_href(func):
return str(id(func))
def get_func_name(self, func_node):
if func_node is None:
return '<missing>'
else:
return '<a href="' + self.generate_href(
func_node) + '">' + func_node.name + '</a>'
# handles when the user clicks on a hyperlink to a function node
def on_func_link_clicked(self, link):
print('Hyperlink clicked: ' + link)
self.select_tree_item(self.func_node_items[self.hrefs_to_funcs[link]])
# Generate treeview row for function (leaf) node in the trie
def add_func_node(self, parent, pattern_col_item, func):
self.hrefs_to_funcs[self.generate_href(func)] = func
self.func_node_items[func] = pattern_col_item
if not func.callees: func.callees = {}
callees_text = '<br />'.join([
str(k) + ': ' + self.get_func_name(v)
for k, v in func.callees.items()
])
callees_item = QStandardItem(callees_text)
cols = [
pattern_col_item,
QStandardItem(func.name), callees_item,
QStandardItem(str(func.pattern_offset) if func.pattern else ''),
QStandardItem(str(func.pattern) if func.pattern else ''),
QStandardItem(func.source_binary),
QStandardItem(self.generate_href(func))
]
boldface = cols[1].font()
boldface.setBold(True)
cols[1].setFont(boldface)
parent.appendRow(cols)
# Recursively add rows for this trie node and its children
def add_trie_node(self, parent, pattern_text, node):
left_item = QStandardItem(pattern_text)
if not node.value: # Stem node
parent.appendRow([left_item, QStandardItem('')])
else: # Leaf node
self.add_func_node(parent, left_item, node.value[0])
for func in node.value[1:]:
self.add_func_node(parent, QStandardItem(''), func)
pairs = map(lambda node: (str(node.pattern), node),
node.children.values())
pairs = sorted(pairs, key=lambda kv: kv[0].replace('?', '\xff'))
for text, child in pairs:
self.add_trie_node(left_item, text, child)
return left_item
# Add bridge nodes to a special node at the root
def add_bridge_nodes(self, parent, sig_trie):
bridge_item = QStandardItem('(bridge)')
parent.appendRow([bridge_item, QStandardItem('')])
def visit(func, visited):
if func is None or func in visited: return
visited.add(func)
if func.is_bridge:
self.add_func_node(bridge_item, QStandardItem(''), func)
for callee in func.callees.values():
visit(callee, visited)
visited = set()
for func in sig_trie.all_values():
visit(func, visited)
def open_trie(self, sig_trie, filename):
self.close_file()
self.sig_trie = sig_trie
root_node = self.add_trie_node(self.model, filename, sig_trie)
self.add_bridge_nodes(root_node, sig_trie)