class Grid(object):
def __init__(self, data):
self.data = data
def create_widget(self, parent=None):
self.tree_view = QTreeView(parent)
self.model = ListModel(self.data)
self.tree_view.setModel(self.model)
return self.tree_view
class getFilesDlg(QDialog):
# sendPaths is a signal emitted by getPaths containing a list of file paths from
# the users selection via this dialog.
sendPaths = Signal(list)
def __init__(self, parent=None):
super(getFilesDlg, self).__init__(parent)
self.setMinimumSize(500,500)
self.fileDlgPaths = []
layout = QVBoxLayout()
self.btnBox = QDialogButtonBox(QDialogButtonBox.Ok |
QDialogButtonBox.Cancel)
self.btnBox.accepted.connect(self.getPaths)
self.btnBox.rejected.connect(self.close)
self.fsModel = QFileSystemModel()
self.treeView = QTreeView()
self.treeView.setSelectionMode(QTreeView.ExtendedSelection)
self.treeView.setModel(self.fsModel)
self.treeView.setColumnWidth(0, 361)
self.treeView.setColumnHidden(1, True)
self.treeView.setColumnHidden(3, True)
layout.addWidget(self.treeView)
layout.addWidget(self.btnBox)
self.setLayout(layout)
self.fsModel.setRootPath(environ['HOMEPATH'])
# self.treeView.setRootIndex(self.fsModel.index("\\"))
self.treeView.expand(self.treeView.rootIndex())
def getPaths(self):
# For some reason duplicates were being returned when they weren't supposed to.
# This obtains the selected files from the dialog and only returns individual
# paths.
indexes = self.treeView.selectedIndexes()
if indexes:
self.fileDlgPaths = []
for i in indexes:
# Possible permission error occuring here
# unable to replicate at this time
path = self.fsModel.filePath(i)
if path not in self.fileDlgPaths:
self.fileDlgPaths.append(path.replace('/','\\'))
self.close() # To close the dialog on an accept signal
self.sendPaths.emit(self.fileDlgPaths)
class DictTreeView(object):
def __init__(self, data):
self.data = data
self.treeView = QTreeView()
self.treeView.setModel(DictModel(self.data))
def set_on_clicked(self, callback):
def execute(index):
callback(index.internalPointer().data)
self.treeView.clicked.connect(execute)
class MainWindow(QMainWindow):
def __init__(self):
super(MainWindow, self).__init__()
self.model = Project()
self.view = QTreeView(self)
self.view.setModel(self.model)
self.view.setDragEnabled(True)
self.view.setDragDropMode(QAbstractItemView.InternalMove)
self.setCentralWidget(self.view)
def mousePressEvent(self, e):
print e.x(), e.y()
return QMainWindow.mousePressEvent(self, e)
class MainWindow(QMainWindow):
def __init__(self):
super(MainWindow, self).__init__()
self.model = Project()
self.view = QTreeView(self)
self.view.setModel(self.model)
self.view.setDragEnabled(True)
self.view.setDragDropMode(QAbstractItemView.InternalMove)
self.setCentralWidget(self.view)
def mousePressEvent(self, e):
print e.x(), e.y()
return QMainWindow.mousePressEvent(self, e)
class NestedListTreeView(object):
def __init__(self, data, header=None):
self.treeView = QTreeView()
self.header = header
self.model = None
self.set_data(data)
# Expand the root item. I'm not sure if this is a good idea in
# general, but it works well on some examples. Maybe a
# heuristic is needed.
root_index = self.model.index(0, 0, QModelIndex())
self.treeView.setExpanded(root_index, True)
# Size all the columns to the size of their current
# contents. Note that this ignores contents for values in the
# tree that are collapsed by default, so it's of limited
# use. However, it's better than the default sizing.
for i in range(self.model.columnCount(None)):
self.treeView.resizeColumnToContents(i)
def set_on_clicked(self, callback):
def execute(index):
callback(index.internalPointer().data)
self.treeView.clicked.connect(execute)
def set_data(self, data):
self.data = data
self.model = ListModel(self.data, header=self.header)
self.treeView.setModel(self.model)
def refresh_data(self):
"""Refresh the UI's view of all data in the tree view. This may be
inefficient with very large data sets.
"""
root_index = self.model.index(0, 0, QModelIndex())
self.treeView.dataChanged(root_index, root_index)
class AutoREView(idaapi.PluginForm):
ADDR_ROLE = QtCore.Qt.UserRole + 1
def __init__(self, data):
super(AutoREView, self).__init__()
self._data = data
self.tv = None
self._model = None
def Show(self):
return idaapi.PluginForm.Show(self, 'AutoRE', options=idaapi.PluginForm.FORM_PERSIST)
def OnCreate(self, form):
if HAS_PYSIDE:
self.parent = self.FormToPySideWidget(form)
else:
self.parent = self.FormToPyQtWidget(form)
self._idp_hooks = AutoReIDPHooks(self)
if not self._idp_hooks.hook():
print 'IDP_Hooks.hook() failed'
self.tv = QTreeView()
self.tv.setExpandsOnDoubleClick(False)
root_layout = QVBoxLayout(self.parent)
# self.le_filter = QLineEdit(self.parent)
# root_layout.addWidget(self.le_filter)
root_layout.addWidget(self.tv)
self.parent.setLayout(root_layout)
self._model = QtGui.QStandardItemModel()
self._init_model()
self.tv.setModel(self._model)
self.tv.setColumnWidth(0, 200)
self.tv.setColumnWidth(1, 300)
self.tv.header().setStretchLastSection(True)
self.tv.expandAll()
self.tv.doubleClicked.connect(self.on_navigate_to_method_requested)
# self.le_filter.textChanged.connect(self.on_filter_text_changed)
self.tv.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
self.tv.customContextMenuRequested.connect(self._tree_customContextMenuRequesssted)
rename_action = QAction('Rename...', self.tv)
rename_action.setShortcut('n')
rename_action.triggered.connect(self._tv_rename_action_triggered)
self.tv.addAction(rename_action)
# def __event_filter(self, source, event):
# if event.type() == QtCore.QEvent.
def _tree_customContextMenuRequesssted(self, pos):
idx = self.tv.indexAt(pos)
if not idx.isValid():
return
addr = idx.data(role=self.ADDR_ROLE)
if not addr:
return
name_idx = idx.sibling(idx.row(), 1)
old_name = name_idx.data()
menu = QMenu()
rename_action = menu.addAction('Rename `%s`...' % old_name)
rename_action.setShortcut('n')
action = menu.exec_(self.tv.mapToGlobal(pos))
if action == rename_action:
return self._rename_ea_requested(addr, name_idx)
def _tv_rename_action_triggered(self):
selected = self.tv.selectionModel().selectedIndexes()
if not selected:
return
idx = selected[0]
if not idx.isValid():
return
addr = idx.data(role=self.ADDR_ROLE)
if not addr:
return
name_idx = idx.sibling(idx.row(), 1)
if not name_idx.isValid():
return
return self._rename_ea_requested(addr, name_idx)
def _rename_ea_requested(self, addr, name_idx):
old_name = name_idx.data()
if idaapi.IDA_SDK_VERSION >= 700:
new_name = idaapi.ask_str(str(old_name), 0, 'New name:')
else:
new_name = idaapi.askstr(0, str(old_name), 'New name:')
if new_name is None:
return
self._rename(addr, new_name)
renamed_name = idaapi.get_ea_name(addr)
name_idx.model().setData(name_idx, renamed_name)
@classmethod
def _rename(cls, ea, new_name):
if not ea or ea == idaapi.BADADDR:
return
if idaapi.IDA_SDK_VERSION >= 700:
return idaapi.force_name(ea, new_name, idaapi.SN_NOCHECK)
return idaapi.do_name_anyway(ea, new_name, 0)
def OnClose(self, form):
if self._idp_hooks:
self._idp_hooks.unhook()
def _tv_init_header(self, model):
item_header = QtGui.QStandardItem("EA")
item_header.setToolTip("Address")
model.setHorizontalHeaderItem(0, item_header)
item_header = QtGui.QStandardItem("Function name")
model.setHorizontalHeaderItem(1, item_header)
item_header = QtGui.QStandardItem("API called")
model.setHorizontalHeaderItem(2, item_header)
# noinspection PyMethodMayBeStatic
def _tv_make_tag_item(self, name):
rv = QtGui.QStandardItem(name)
rv.setEditable(False)
return [rv, QtGui.QStandardItem(), QtGui.QStandardItem()]
def _tv_make_ref_item(self, tag, ref):
ea_item = QtGui.QStandardItem(('%0' + get_addr_width() + 'X') % ref['ea'])
ea_item.setEditable(False)
ea_item.setData(ref['ea'], self.ADDR_ROLE)
name_item = QtGui.QStandardItem(ref['name'])
name_item.setEditable(False)
name_item.setData(ref['ea'], self.ADDR_ROLE)
apis = ', '.join(ref['tags'][tag])
api_name = QtGui.QStandardItem(apis)
api_name.setEditable(False)
api_name.setData(ref['ea'], self.ADDR_ROLE)
api_name.setToolTip(apis)
return [ea_item, name_item, api_name]
def _init_model(self):
self._model.clear()
root_node = self._model.invisibleRootItem()
self._tv_init_header(self._model)
for tag, refs in self._data.items():
item_tag_list = self._tv_make_tag_item(tag)
item_tag = item_tag_list[0]
root_node.appendRow(item_tag_list)
for ref in refs:
ref_item_list = self._tv_make_ref_item(tag, ref)
item_tag.appendRow(ref_item_list)
def on_navigate_to_method_requested(self, index):
addr = index.data(role=self.ADDR_ROLE)
if addr is not None:
idaapi.jumpto(addr)
class BrowserWindow(QMainWindow):
MO_ROLE = Qt.UserRole+1
def __init__(self, conn):
super(BrowserWindow, self).__init__()
self._conn = conn
self._resolver = AsyncResolver()
self._resolver.object_resolved.connect(self._data_resolved)
self._resolver.start()
self._init_models()
self._init_gui()
self._init_data()
self._init_connections()
def __del__(self):
self._resolver.stop_work()
self._resolver.terminate()
def _init_models(self):
self._hierarchy_model = QStandardItemModel()
self._hierarchy_model.setColumnCount(2)
self._hierarchy_model.setHorizontalHeaderLabels(['class', 'dn'])
self._details_model = QStandardItemModel()
self._details_model.setColumnCount(2)
self._details_model.setHorizontalHeaderLabels(['Property', 'Value'])
def _init_gui(self):
self._widget = QSplitter(self, Qt.Horizontal)
self._hierarchy_view = QTreeView(self._widget)
self._details_view = QTableView(self._widget)
self._widget.addWidget(self._hierarchy_view)
self._widget.addWidget(self._details_view)
self._widget.setStretchFactor(0, 2)
self._widget.setStretchFactor(1, 1)
self.setCentralWidget(self._widget)
self._hierarchy_view.setModel(self._hierarchy_model)
self._details_view.setModel(self._details_model)
self._hierarchy_view.expanded.connect(self._mo_item_expand)
def _init_data(self):
item = self._row_for_mo(self._conn.resolve_dn(''))
self._hierarchy_model.insertRow(0, item)
def _init_connections(self):
self.connect(self._resolver,
SIGNAL('object_resolved(QVariant)'),
self,
SLOT('_data_resolved(QVariant)'))
self._hierarchy_view.activated.connect(self._item_activated)
#self.connect(self._hierarchy_view.selectionModel(),
# SIGNAL('currentChanged(QModelIndex,QModelIndex)'),
# self,
# SLOT('_current_changed(QModelIndex, QModelIndex)'))
self.connect(self._hierarchy_view.selectionModel(),
SIGNAL('activated(QModelIndex)'),
self,
SLOT('_item_activated(QModelIndex)'))
def _row_for_mo(self, mo):
row = [QStandardItem(mo.ucs_class), QStandardItem(mo.dn)]
for item in row:
item.setEditable(False)
row[0].appendColumn([QStandardItem('Loading...')])
row[0].setData(mo, self.MO_ROLE)
return row
def _add_mo_in_tree(self, mo, index=QtCore.QModelIndex()):
item = None
if index.isValid():
item = self._hierarchy_model.itemFromIndex(index)
else:
item = self._get_item_for_dn(self._parent_dn(mo.dn))
if item:
item.appendColumn([self._row_for_mo(mo)[0]])
self.auto_width()
def _add_mos_in_tree(self, mos, index=QtCore.QModelIndex()):
item = None
if index.isValid():
item = self._hierarchy_model.itemFromIndex(index)
else:
if not mos:
return
item = self._get_item_for_dn(self._parent_dn(mos[0].dn))
while item.columnCount():
item.removeColumn(0)
items = map(self._row_for_mo, mos)
if items:
for x in xrange(len(items[0])):
item.appendColumn([row[x] for row in items])
self.auto_width()
@staticmethod
def _parent_dn(dn):
parent_dn, _, rn = dn.rpartition('/')
return parent_dn
def _get_item_for_dn(self, dn):
parent_dn = dn
items = self._hierarchy_model.findItems(parent_dn, column=1)
if items:
return self._hierarchy_model.item(items[0].row())
return None
@QtCore.Slot('_data_resolved(QVariant)')
def _data_resolved(self, datav):
print 'Data resolved: ', datav
index, data = datav
if isinstance(data, UcsmObject):
self._add_mo_in_tree(data, index=index)
else:
self._add_mos_in_tree(data, index=index)
@QtCore.Slot('_current_changed(QModelIndex,QModelIndex)')
def _current_changed(self, curr, prev):
self._item_activated(curr)
@QtCore.Slot('_item_activated(QModelIndex)')
def _item_activated(self, index):
print 'Activated: %s data %s' % (index, index.data(self.MO_ROLE))
if index.sibling(0, 0).isValid():
index = index.sibling(0, 0)
data = index.data(self.MO_ROLE)
self.set_detail_object(data)
def _mo_item_expand(self, index):
obj = index.data(self.MO_ROLE)
print 'Expanded object: %s' % obj
try:
self._resolver.add_task(lambda: (index,
self._conn.resolve_children(obj.dn)))
except (KeyError, AttributeError):
QtGui.QMessageBox.critical(0, 'Error', 'Object does not have dn')
def auto_width(self):
for view in [self._hierarchy_view, self._details_view]:
for col in xrange(view.model().columnCount()):
view.resizeColumnToContents(col)
def set_detail_object(self, object):
self._details_model.removeRows(0, self._details_model.rowCount())
for k, v in object.attributes.iteritems():
row = [QStandardItem(k), QStandardItem(v)]
for item in row:
item.setEditable(False)
self._details_model.appendRow(row)
self.auto_width()
class FilterTab(QWidget):
"""This class is the GUI for creating filters for the FilterBox
module. This GUI will be added as a tab to the ModuleFrame
tab dialog.
"""
applySignal = Signal(Clause)
def __init__(self, parent, mframe, existing_filters):
"""Create a FilterTab with the given parent TabDialog, logical
parent ModuleFrame mframe, and existing_filters list of Clause
objects.
"""
super(FilterTab, self).__init__(parent)
self.mframe = mframe
self.parent = parent
self.attributes = self.mframe.agent.datatree.generateAttributeList()
self.clause_list = list()
self.clause_dict = dict()
# Right now we only look at the first passed in filter.
# TODO: At GUI to switch between existing filters
if existing_filters is not None and len(existing_filters) > 0:
for clause in existing_filters[0].conditions.clauses:
self.clause_list.append(str(clause))
self.clause_dict[str(clause)] = clause
self.clause_model = QStringListModel(self.clause_list)
layout = QVBoxLayout(self)
self.sidesplitter = QSplitter(Qt.Horizontal)
# You can only select one attribute at a time to build the
# filter clauses
self.data_view = QTreeView(self)
self.data_view.setModel(self.mframe.agent.datatree)
self.data_view.setDragEnabled(True)
self.data_view.setDropIndicatorShown(True)
self.data_view.expandAll()
self.sidesplitter.addWidget(self.data_view)
self.sidesplitter.setStretchFactor(1,1)
self.filter_widget = self.buildFilterWidget()
self.sidesplitter.addWidget(self.filter_widget)
self.sidesplitter.setStretchFactor(1,0)
layout.addWidget(self.sidesplitter)
# Apply buttons
buttonWidget = QWidget()
buttonLayout = QHBoxLayout(buttonWidget)
self.applyButton = QPushButton("Apply")
self.applyButton.clicked.connect(self.applyFilter)
self.closeButton = QPushButton("Apply & Close")
self.closeButton.clicked.connect(self.applyCloseFilter)
buttonLayout.addWidget(self.applyButton)
buttonLayout.addWidget(self.closeButton)
buttonWidget.setLayout(buttonLayout)
layout.addWidget(buttonWidget)
self.setLayout(layout)
def applyFilter(self):
"""Emits the applySignal with the Clause object currently
represented by this FilterTab.
"""
num_clauses = len(self.clause_list)
if num_clauses == 0:
self.applySignal.emit(None)
else:
self.applySignal.emit(Clause("and", *self.clause_dict.values()))
def applyCloseFilter(self):
"""Calls applyFilter and then closes the containing TabDialog."""
self.applyFilter()
self.parent.close()
def buildFilterWidget(self):
"""Creates the filter portion of the widget by laying out
the subwidgets for relations, workspace and existing
clauses.
"""
filter_widget = QWidget()
filter_layout = QVBoxLayout(filter_widget)
filter_layout.addWidget(self.buildRelationsWidget())
filter_layout.addItem(QSpacerItem(5,5))
filter_layout.addWidget(self.buildWorkFrame())
filter_layout.addItem(QSpacerItem(5,5))
filter_layout.addWidget(self.buildFilterListView())
filter_widget.setLayout(filter_layout)
return filter_widget
def buildFilterListView(self):
"""Creates the QListView that contains all of the basic Clause
objects.
"""
groupBox = QGroupBox("Clauses")
layout = QVBoxLayout(groupBox)
self.list_view = QListView(groupBox)
self.list_view.setModel(self.clause_model)
layout.addWidget(self.list_view)
layout.addItem(QSpacerItem(5,5))
self.delButton = QPushButton("Remove Selected Clause")
self.delButton.clicked.connect(self.deleteClause)
layout.addWidget(self.delButton)
groupBox.setLayout(layout)
return groupBox
def buildWorkFrame(self):
"""Creates the grouped set of widgets that allow users to build
basic Clause objects.
"""
groupBox = QGroupBox("Clause Workspace")
layout = QHBoxLayout(groupBox)
attributeCompleter = QCompleter(self.attributes)
attributeCompleter.setCompletionMode(QCompleter.InlineCompletion)
self.dropAttribute = DropLineEdit(self, self.mframe.agent.datatree, "",
attributeCompleter)
self.dropRelation = DropTextLabel("__")
self.dropValue = FilterValueLineEdit(groupBox,
self.mframe.agent.datatree, self.dropAttribute)
# Clear dropValue when dropAttribute changes
self.dropAttribute.textChanged.connect(self.dropValue.clear)
# Enter in dropValue works like addButton
self.dropValue.returnPressed.connect(self.addClause)
self.addButton = QPushButton("Add", groupBox)
self.addButton.clicked.connect(self.addClause)
layout.addWidget(self.dropAttribute)
layout.addItem(QSpacerItem(5,5))
layout.addWidget(self.dropRelation)
layout.addItem(QSpacerItem(5,5))
layout.addWidget(self.dropValue)
layout.addItem(QSpacerItem(5,5))
layout.addWidget(self.addButton)
groupBox.setLayout(layout)
return groupBox
def buildRelationsWidget(self):
"""Creates the set of draggable relations. These relations are
whatever is available in the relations dict of Table.
"""
relations_widget = QWidget()
layout = QHBoxLayout(relations_widget)
for relation in Table.relations:
layout.addWidget(DragTextLabel(relation))
relations_widget.setLayout(layout)
return relations_widget
def addClause(self):
"""Adds a basic Clause to the current filter."""
if self.dropRelation.text() in Table.relations \
and len(self.dropValue.text()) > 0 \
and len(self.dropAttribute.text()) > 0:
clause = Clause(self.dropRelation.text(),
TableAttribute(self.dropAttribute.text()),
self.dropValue.text())
# Guard double add
if str(clause) not in self.clause_dict:
self.clause_list.append(str(clause))
self.clause_dict[str(clause)] = clause
self.clause_model.setStringList(self.clause_list)
def deleteClause(self):
"""Removes the selected basic Clause objects from the current
filter.
"""
clause = self.clause_model.data(
self.list_view.selectedIndexes()[0], Qt.DisplayRole)
if clause is not None and clause in self.clause_list:
self.clause_list.remove(clause)
del self.clause_dict[clause]
self.clause_model.setStringList(self.clause_list)
if result:
if index:
tv.setCurrentIndex(index)
return
tv.clearSelection()
app = QApplication([])
model = treeModel()
dialog = QDialog()
dialog.setMinimumSize(300, 150)
layout = QVBoxLayout(dialog)
tv = QTreeView(dialog)
tv.setModel(model)
tv.setAlternatingRowColors(True)
layout.addWidget(tv)
label = QLabel("Search for the following person")
layout.addWidget(label)
buts = []
frame = QFrame(dialog)
layout2 = QHBoxLayout(frame)
for person in model.people:
but = QPushButton(person.fname, frame)
buts.append(but)
layout2.addWidget(but)
QObject.connect(but, SIGNAL("clicked()"), but_clicked)
class AutoREView(idaapi.PluginForm):
ADDR_ROLE = QtCore.Qt.UserRole + 1
def __init__(self, data):
super(AutoREView, self).__init__()
self._data = data
def Show(self):
return idaapi.PluginForm.Show(self,
'AutoRE',
options=idaapi.PluginForm.FORM_PERSIST)
def OnCreate(self, form):
# if HAS_PYSIDE:
# self.parent = self.FormToPySideWidget(form)
# else:
self.parent = self.FormToPyQtWidget(form)
self.tv = QTreeView()
self.tv.setExpandsOnDoubleClick(False)
root_layout = QVBoxLayout(self.parent)
# self.le_filter = QLineEdit(self.parent)
# root_layout.addWidget(self.le_filter)
root_layout.addWidget(self.tv)
self.parent.setLayout(root_layout)
self._model = QtGui.QStandardItemModel()
self._init_model()
self.tv.setModel(self._model)
self.tv.setColumnWidth(0, 200)
self.tv.setColumnWidth(1, 300)
self.tv.header().setStretchLastSection(True)
self.tv.expandAll()
self.tv.doubleClicked.connect(self.on_navigate_to_method_requested)
# self.le_filter.textChanged.connect(self.on_filter_text_changed)
def OnClose(self, form):
# print 'TODO: OnClose(): clear the pointer to form in the plugin'
pass
def _tv_init_header(self, model):
item_header = QtGui.QStandardItem("EA")
item_header.setToolTip("Address")
model.setHorizontalHeaderItem(0, item_header)
item_header = QtGui.QStandardItem("Function name")
model.setHorizontalHeaderItem(1, item_header)
item_header = QtGui.QStandardItem("API called")
model.setHorizontalHeaderItem(2, item_header)
def _tv_make_tag_item(self, name):
rv = QtGui.QStandardItem(name)
rv.setEditable(False)
return [rv, QtGui.QStandardItem(), QtGui.QStandardItem()]
def _tv_make_ref_item(self, tag, ref):
ea_item = QtGui.QStandardItem(
('%0' + get_addr_width() + 'X') % ref['ea'])
ea_item.setEditable(False)
ea_item.setData(ref['ea'], self.ADDR_ROLE)
name_item = QtGui.QStandardItem(ref['name'])
name_item.setEditable(False)
name_item.setData(ref['ea'], self.ADDR_ROLE)
apis = ', '.join(ref['tags'][tag])
api_name = QtGui.QStandardItem(apis)
api_name.setEditable(False)
api_name.setData(ref['ea'], self.ADDR_ROLE)
api_name.setToolTip(apis)
return [ea_item, name_item, api_name]
def _init_model(self):
self._model.clear()
root_node = self._model.invisibleRootItem()
self._tv_init_header(self._model)
for tag, refs in self._data.items():
item_tag_list = self._tv_make_tag_item(tag)
item_tag = item_tag_list[0]
root_node.appendRow(item_tag_list)
for ref in refs:
ref_item_list = self._tv_make_ref_item(tag, ref)
item_tag.appendRow(ref_item_list)
def on_navigate_to_method_requested(self, index):
addr = index.data(role=self.ADDR_ROLE)
if addr is not None:
idaapi.jumpto(addr)
if result:
if index:
tv.setCurrentIndex(index)
return
tv.clearSelection()
app = QApplication([])
model = treeModel()
dialog = QDialog()
dialog.setMinimumSize(300, 150)
layout = QVBoxLayout(dialog)
tv = QTreeView(dialog)
tv.setModel(model)
tv.setAlternatingRowColors(True)
layout.addWidget(tv)
label = QLabel("Search for the following person")
layout.addWidget(label)
buts = []
frame = QFrame(dialog)
layout2 = QHBoxLayout(frame)
for person in model.people:
but = QPushButton(person.fname, frame)
buts.append(but)
layout2.addWidget(but)
QObject.connect(but, SIGNAL("clicked()"), but_clicked)
class MainWindow(QMainWindow):
def __init__(self, parent=None):
super(MainWindow, self).__init__(parent)
self.dirty = False
# Models
# - Filesystem Model
self.fileSystemModel = QFileSystemModel()
# rootPath = QDesktopServices.storageLocation(
# QDesktopServices.HomeLocation)
# fileSystemRoot = self.fileSystemModel.setRootPath(rootPath)
fileSystemRoot = self.fileSystemModel.setRootPath(
"/home/ryan/Programming/Python/projects/r3tagger/r3tagger/tests"
)
# Views
# - Filesystem View
self.fileSystemView = QTreeView()
self.fileSystemView.setModel(self.fileSystemModel)
self.fileSystemView.setRootIndex(fileSystemRoot)
self.fileSystemView.doubleClicked.connect(self.updateAlbumModel)
self.fileSystemView.expanded.connect(self.fixFileSystemColumns)
self.fileSystemView.collapsed.connect(self.fixFileSystemColumns)
# - Album View
self.albumView = albumcollection.MusicCollectionView()
self.albumView.setSelectionMode(QAbstractItemView.MultiSelection)
self.albumView.clicked.connect(self.updateEditing)
self.albumView.expanded.connect(self.fixAlbumViewColumns)
self.albumView.collapsed.connect(self.fixAlbumViewColumns)
self.albumView.model().dataChanged.connect(self.fixAlbumViewColumns)
self.albumView.model().dataChanged.connect(self._setDirty)
model = self.albumView.model()
model.dataChanged.connect(self.updateEditing)
# Editing Group
self.editingGroup = QFormLayout()
self.lineArtist = QLineEdit()
self.lineAlbum = QLineEdit()
self.lineTitle = QLineEdit()
self.lineTrack = QLineEdit()
self.lineDate = QLineEdit()
self.lineGenre = QLineEdit()
self.editingGroup.addRow("Artist:", self.lineArtist)
self.editingGroup.addRow("Album:", self.lineAlbum)
self.editingGroup.addRow("Title:", self.lineTitle)
self.editingGroup.addRow("Track:", self.lineTrack)
self.editingGroup.addRow("Date:", self.lineDate)
self.editingGroup.addRow("Genre:", self.lineGenre)
self.tagsToAttribs = {
"artist": self.lineArtist,
"album": self.lineAlbum,
"title": self.lineTitle,
"tracknumber": self.lineTrack,
"date": self.lineDate,
"genre": self.lineGenre,
}
# Confirm / Cancel / Clear Group
self.buttonGroup = QHBoxLayout()
self.buttonGroup.addStretch()
confirm = QPushButton("Confirm")
confirm.clicked.connect(self.confirmChanges)
self.buttonGroup.addWidget(confirm)
cancel = QPushButton("Cancel")
cancel.clicked.connect(self.cancelChanges)
self.buttonGroup.addWidget(cancel)
clear = QPushButton("Clear")
clear.clicked.connect(self.clearAlbumView)
self.buttonGroup.addWidget(clear)
self.buttonGroup.addStretch()
# Statusbar
# status = self.statusBar()
# status.setSizeGripEnabled(False)
# status.showMessage("Ready", 5000)
# Docks
dockAllowed = Qt.LeftDockWidgetArea | Qt.RightDockWidgetArea
# - Filesystem Dock
fileSystemDock = QDockWidget("Navigate", self)
fileSystemDock.setObjectName("fileSystemDock")
fileSystemDock.setAllowedAreas(dockAllowed)
fileSystemDock.setWidget(self.fileSystemView)
self.addDockWidget(Qt.LeftDockWidgetArea, fileSystemDock)
# - Editing Dock
editingWidget = QWidget()
editingWidget.setLayout(self.editingGroup)
editingDock = QDockWidget("Editing", self)
editingDock.setObjectName("editingDock")
editingDock.setAllowedAreas(dockAllowed)
editingDock.setWidget(editingWidget)
self.addDockWidget(Qt.RightDockWidgetArea, editingDock)
# Actions
fileAddSongAction = self._createAction(
text="Add &Songs",
slot=self.fileAddSong,
shortcut=QKeySequence.Open,
icon="fileOpen",
tip="Add files (songs)",
)
fileAddAlbumAction = self._createAction(
text="Add &Album",
slot=self.fileAddAlbum,
shortcut=QKeySequence(Qt.CTRL + Qt.SHIFT + Qt.Key_O),
icon="fileOpen",
tip="Add directory (album)",
)
fileSaveAction = self._createAction(
text="&Save Changes",
slot=self.confirmChanges,
shortcut=QKeySequence.Save,
icon="fileSave",
tip="Save Changes",
)
fileQuitAction = self._createAction(
text="&Quit", slot=self.close, shortcut=QKeySequence.Quit, icon="fileQuit", tip="Quit Program"
)
editRecognizeAction = self._createAction(
text="&Recognize",
slot=self.editRecognize,
shortcut=QKeySequence(Qt.CTRL + Qt.Key_R),
icon="editRecognize",
tip="Recognize music",
)
editReorganizeAction = self._createAction(
text="Reorganize",
slot=self.editReorganize,
shortcut=QKeySequence(Qt.CTRL + Qt.SHIFT + Qt.Key_R),
icon="editReorganize",
tip="Reorganize music",
)
editSettingsAction = self._createAction(
text="Settings",
slot=self.editSettings,
shortcut=QKeySequence.Preferences,
icon="editSettings",
tip="Edit settings",
)
helpDocsAction = self._createAction(
text="Documentation",
slot=self.helpDocs,
shortcut=QKeySequence.HelpContents,
icon="helpDocs",
tip="Documentation",
)
helpAboutAction = self._createAction(
text="About",
slot=self.helpAbout,
shortcut=QKeySequence(Qt.CTRL + Qt.SHIFT + Qt.Key_F1),
icon="helpAbout",
tip="About",
)
toggleEditing = editingDock.toggleViewAction()
toggleEditing.setIcon(QIcon(":/toggleEditing.png"))
toggleFileNav = fileSystemDock.toggleViewAction()
toggleFileNav.setIcon(QIcon(":/toggleFileNav.png"))
# Menus
fileMenu = self.menuBar().addMenu("&File")
self._addActions(fileMenu, (fileAddSongAction, fileAddAlbumAction, fileSaveAction, fileQuitAction))
editMenu = self.menuBar().addMenu("&Edit")
self._addActions(editMenu, (editReorganizeAction, editRecognizeAction, editSettingsAction))
helpMenu = self.menuBar().addMenu("&Help")
self._addActions(helpMenu, (helpDocsAction, helpAboutAction))
# Toolbars
editToolbar = self.addToolBar("EditToolbar")
editToolbar.setObjectName("editToolbar")
self._addActions(editToolbar, (editRecognizeAction, editReorganizeAction))
toggleToolbar = self.addToolBar("ToggleToolbar")
toggleToolbar.setObjectName("toggleToolbar")
self._addActions(toggleToolbar, (toggleFileNav, toggleEditing))
# Settings
settings = QSettings()
if settings.contains("MainWindow/Geometry"):
self.restoreGeometry(settings.value("MainWindow/Geometry"))
if settings.contains("MainWindow/State"):
self.restoreState(settings.value("MainWindow/State"))
self.setWindowTitle("r3tagger")
# Final Layout
centralWidget = QWidget()
centralLayout = QVBoxLayout()
centralLayout.addWidget(self.albumView)
centralLayout.addLayout(self.buttonGroup)
centralWidget.setLayout(centralLayout)
self.setCentralWidget(centralWidget)
def _createAction(self, text, slot=None, shortcut=None, icon=None, tip=None, checkable=False, signal="triggered"):
action = QAction(text, self)
if icon is not None:
action.setIcon(QIcon(":/{0}.png".format(icon)))
if shortcut is not None:
action.setShortcut(shortcut)
if tip is not None:
action.setToolTip(tip)
action.setStatusTip(tip)
if slot is not None:
getattr(action, signal).connect(slot)
if checkable:
action.setCheckable(True)
return action
def _addActions(self, target, actions):
for action in actions:
if action is None:
target.addSeparator()
else:
target.addAction(action)
def _setDirty(self):
self.dirty = True
def _fixColumns(self, index, view, model):
for column in range(model.columnCount(index)):
view.resizeColumnToContents(column)
def closeEvent(self, event):
if self.dirty:
reply = QMessageBox.question(
self,
"r3tagger - Unsaved Changes",
"Save unsaved changes?",
QMessageBox.Yes | QMessageBox.No | QMessageBox.Cancel,
)
if reply == QMessageBox.Cancel:
event.ignore()
return None
elif reply == QMessageBox.Yes:
self.confirmChanges()
settings = QSettings()
settings.setValue("MainWindow/Geometry", self.saveGeometry())
settings.setValue("MainWindow/State", self.saveState())
def fixFileSystemColumns(self, index):
self._fixColumns(index, self.fileSystemView, self.fileSystemModel)
def fixAlbumViewColumns(self, index):
model = self.albumView.model()
self._fixColumns(index, self.albumView, model)
def updateAlbumModel(self, index):
path = self.fileSystemModel.fileInfo(index).absoluteFilePath()
self.addPath(path)
self.fixAlbumViewColumns(index)
def addPath(self, path):
if os.path.isfile(path):
track = controller.build_track(path)
containerAlbum = controller.album_from_tracks([track], u"Singles")
self.albumView.model().addAlbum(containerAlbum)
else:
for album in controller.build_albums(path, recursive=True):
self.albumView.model().addAlbum(album)
def updateEditing(self, index):
self.albumView.correctListingSelection(index)
selectedTracks = self.albumView.selectedTracks()
albumOfSingles = controller.album_from_tracks(selectedTracks)
selected = self.albumView.selectedAlbums()
selected.append(albumOfSingles)
tags = controller.find_shared_tags(*selected) if selected else {}
for tag, edit in self.tagsToAttribs.items():
if not tags:
self.clearEditing()
break
edit.setText(tags.get(tag, ""))
edit.setCursorPosition(0)
def confirmChanges(self):
tags = {}
for field, lineEdit in self.tagsToAttribs.items():
tag = lineEdit.text()
tags[field] = tag
view = self.albumView
for album in view.selectedAlbums():
controller.retag_album(album, tags)
for track in view.selectedTracks():
controller.retag_track(track, tags)
self.saveChanges()
def clearAlbumView(self):
model = self.albumView.model()
model.clear()
model.setHeaders()
def clearEditing(self):
for lineEdit in self.tagsToAttribs.values():
lineEdit.setText("")
def saveChanges(self):
for track in self.albumView.model():
if track.dirty:
track.saveChanges()
self.dirty = False
self.resetModel()
def cancelChanges(self):
for track in self.albumView.model():
track.reset()
self.resetModel()
def resetModel(self):
model = self.albumView.model()
model.beginResetModel()
self.clearEditing()
expanded = []
rowCount = model.rowCount(QModelIndex())
for row in range(rowCount):
index = model.index(row, 0, QModelIndex())
if self.albumView.isExpanded(index):
expanded.append(index)
model.endResetModel()
for expandedIndex in expanded:
self.albumView.setExpanded(expandedIndex, True)
def fileAddSong(self):
selectedFiles, selectedFilter = QFileDialog.getOpenFileNames(parent=self, caption="Add Songs")
for song in selectedFiles:
self.addPath(song)
def fileAddAlbum(self):
selectedDir = QFileDialog.getExistingDirectory(parent=self, caption="Add Album")
self.addPath(selectedDir)
def editRecognize(self):
pass
def editReorganize(self):
pass
def editSettings(self):
pass
def helpDocs(self):
pass
def helpAbout(self):
pass
