def setup_file_explorer(tree, path=""):
"""Creates the file explorer rooted at a particular path.
Args:
tree (QTreeView): Tree to be populated tied to the file explorer.
path (str, optional): Path to the root of the project.
Returns:
tuple[QFileSystemModel, QTreeView]: File system UI element and tied populated
file tree.
"""
tree.setEnabled(os.path.isdir(path))
model = QFileSystemModel()
model.setRootPath(path)
tree.setModel(model)
tree.setRootIndex(model.index(path))
# Resize column 0 (name) to content length, and stretch the rest to the size of the widget
tree.header().setSectionResizeMode(0, QHeaderView.ResizeToContents)
for i in range(1, model.columnCount()):
tree.header().setSectionResizeMode(i, QHeaderView.Stretch)
tree.setSortingEnabled(True)
tree.sortByColumn(0, Qt.AscendingOrder)
return model, tree
class FileChooser(QWidget):
fileOpened = pyqtSignal(str)
def __init__(self, parent=None):
super().__init__(parent)
self.folderBox = QComboBox(self)
self.explorerTree = FileTreeView(self)
self.explorerTree.doubleClickCallback = self._fileOpened
self.explorerModel = QFileSystemModel(self)
self.explorerModel.setFilter(
QDir.AllDirs | QDir.Files | QDir.NoDotAndDotDot)
self.explorerModel.setNameFilters(["*.py"])
self.explorerModel.setNameFilterDisables(False)
self.explorerTree.setModel(self.explorerModel)
for index in range(1, self.explorerModel.columnCount()):
self.explorerTree.hideColumn(index)
self.setCurrentFolder()
self.folderBox.currentIndexChanged[int].connect(
self.updateCurrentFolder)
layout = QVBoxLayout(self)
layout.addWidget(self.folderBox)
layout.addWidget(self.explorerTree)
layout.setContentsMargins(5, 5, 0, 0)
def _fileOpened(self, modelIndex):
path = self.explorerModel.filePath(modelIndex)
if os.path.isfile(path):
self.fileOpened.emit(path)
def currentFolder(self):
return self.explorerModel.rootPath()
def setCurrentFolder(self, path=None):
if path is None:
app = QApplication.instance()
path = app.getScriptsDirectory()
else:
assert os.path.isdir(path)
self.explorerModel.setRootPath(path)
self.explorerTree.setRootIndex(self.explorerModel.index(path))
self.folderBox.blockSignals(True)
self.folderBox.clear()
style = self.style()
dirIcon = style.standardIcon(style.SP_DirIcon)
self.folderBox.addItem(dirIcon, os.path.basename(path))
self.folderBox.insertSeparator(1)
self.folderBox.addItem(self.tr("Browse…"))
self.folderBox.setCurrentIndex(0)
self.folderBox.blockSignals(False)
def updateCurrentFolder(self, index):
if index < self.folderBox.count() - 1:
return
path = QFileDialog.getExistingDirectory(
self, self.tr("Choose Directory"), self.currentFolder(),
QFileDialog.ShowDirsOnly)
if path:
QSettings().setValue("scripting/path", path)
self.setCurrentFolder(path)
def initRootDir(self,indexDir):
if not indexDir:
return
self.indexDir = indexDir
model = QFileSystemModel()
model.setRootPath('')
self.setModel(model)
self.setAnimated(False)
self.setIndentation(20)
self.setSortingEnabled(True)
self.setRootIndex(model.index(self.indexDir))
for i in range(1,model.columnCount()):
self.hideColumn(i)
class DirTreeView(QTreeView):
"""
A widget class used to display the contents of an Informatic project source
directory.
"""
newSelection = pyqtSignal([list])
def __init__(self, parent=None, rootdir=None):
"""
The rootdir keyword argument is a filepath for a directory initialized
as the root directory whose contents are displayed by the widget.
"""
# Invoke the QTreeView constructor
super().__init__(parent)
self.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Preferred)
# The widget's contents are based on a file system model
self.dirTree = QFileSystemModel()
# Display only files with filename extensions commonly used for Inform 6
# source files
self.dirTree.setNameFilterDisables(False)
self.dirTree.setNameFilters(['*.inf', '*.i6', '*.h'])
self.dirTree.setRootPath(rootdir)
self.setModel(self.dirTree)
self.cd(rootdir)
# Hide all but the first column, which holds the filename
for column in range(1, self.dirTree.columnCount()):
self.hideColumn(column)
def cd(self, path):
"""
Takes one argument, path, a directory filepath, and changes the root
directory displayed by the widget to the directory at that filepath.
"""
self.setRootIndex(self.dirTree.index(path))
def selectionChanged(self, selected, deselected):
"""
Emits the newSelection signal with a list of selected items whenever the
selection of items in the file tree is changed.
"""
self.newSelection.emit(selected.indexes())
def __init__(self): QTreeView.__init__(self) model = QFileSystemModel() model.setRootPath(QDir.homePath()) self.setModel(model) self.setRootIndex(model.index(QDir.homePath())) model.setReadOnly(False) for i in range(1, model.columnCount()): self.hideColumn(i) self.setHeaderHidden(True) self.setSelectionMode(self.SingleSelection) self.setDragDropMode(QAbstractItemView.InternalMove) self.setDragEnabled(True) self.setAcceptDrops(True) self.setDropIndicatorShown(True)
class MainWindow(QMainWindow):
def __init__(self, parent=None):
super(MainWindow, self).__init__(parent)
self.ui = Ui_MainWindow()
self.ui.setupUi(self)
path = QDir.currentPath()
self.browser_model = QFileSystemModel()
self.browser_model.setRootPath(path)
self.browser_model.setFilter(QDir.NoDotAndDotDot | QDir.AllDirs)
self.ui.browser.setModel(self.browser_model)
self.details_model = QFileSystemModel()
self.details_model.setRootPath(path)
self.details_model.setFilter(QDir.NoDotAndDotDot | QDir.AllEntries)
self.ui.details.setModel(self.details_model)
column_count = self.browser_model.columnCount()
for i in range(1, column_count):
self.ui.browser.hideColumn(i)
self.setupUi()
def setupUi(self):
self.ui.browser.clicked.connect(self.browser_clicked)
def browser_clicked(self, index):
file_info = self.browser_model.fileInfo(index)
path = file_info.absoluteFilePath()
self.ui.details.setRootIndex(self.details_model.setRootPath(path))
class FileManager:
""" Create a popup window to select a folder / file
Global Varables
---------------
item_index:
PyQt5 index element to have a global reference to currently clicked element
callback_func:
Global reference to set the callback function passed into create_tree_view_popup
popup_window:
Global reference to the QWidget window which serves as a popup container for the TreeView
"""
# pylint: disable=no-self-use
def __init__(self, width=0):
""" Initializes required variables by the class """
self.width = width
self.item_index = None
self.callback_func = None
self.popup_window = None
self.model = None
def create_tree_view(self, start_path="/home", callback=None):
""" Display the tree view to choose a folder """
file_manager = QTreeView()
self.callback_func = callback
self.model = QFileSystemModel()
self.model.setRootPath(start_path)
# for ViewType in (QColumnView, QTreeView):
file_manager.setModel(self.model)
# hide all columns except the filename
for i in range(1, self.model.columnCount()):
file_manager.hideColumn(i)
file_manager.setRootIndex(self.model.index(start_path))
file_manager.setFixedWidth(self.width)
file_manager.setWordWrap(True)
file_manager.clicked.connect(self.item_clicked)
return {"file_manager": file_manager, "file_model": self.model}
def create_file_dialog(self, start_path="/home"):
""" Create a file dialog window """
file_dialog = QFileDialog()
file_dialog.setFileMode(QFileDialog.AnyFile)
file_dialog.setDirectory(start_path)
return file_dialog
def create_tree_view_popup(self,
start_path="/home",
callback=None,
directory=False):
""" Function which creates a TreeView in a external Window"""
# self.popup_window = QWidget()
actual_path = None
if not os.path.isdir(start_path) and os.path.isfile(start_path):
path_arr = start_path.split('/')
path_arr.pop(len(path_arr) - 1)
actual_path = '/'.join(e for e in path_arr)
else:
actual_path = "/home"
if not directory:
filename = QFileDialog.getOpenFileName(None, "Open File",
actual_path)
if filename[0]:
callback({"file_path": filename[0]})
else:
directory = QFileDialog.getExistingDirectory(
None, 'Select a Folder:', actual_path,
QFileDialog.ShowDirsOnly)
callback({"dir_path": directory})
def item_clicked(self, index):
""" Helper function to keep track of the selected item in the TreeView """
self.item_index = index
self.select_item()
def select_item(self):
""" Passes the current selected item to the callback function and closes the window """
# get the index of the currently chosen item
index_item = self.model.index(self.item_index.row(), 0,
self.item_index.parent())
# get the file path of this item
file_path = self.model.filePath(index_item)
# return it with a callback
self.callback_func({"file_path": file_path})
class PBTKGUI(QApplication):
def __init__(self):
super().__init__(argv)
signal(SIGINT, SIG_DFL)
if '__file__' in globals():
views = dirname(realpath(__file__)) + '/views/'
else:
views = dirname(realpath(executable)) + '/views/'
self.welcome = loadUi(views + 'welcome.ui')
self.choose_extractor = loadUi(views + 'choose_extractor.ui')
self.choose_proto = loadUi(views + 'choose_proto.ui')
self.create_endpoint = loadUi(views + 'create_endpoint.ui')
self.choose_endpoint = loadUi(views + 'choose_endpoint.ui')
self.fuzzer = loadUi(views + 'fuzzer.ui')
self.welcome.step1.clicked.connect(self.load_extractors)
self.choose_extractor.rejected.connect(
partial(self.set_view, self.welcome))
self.choose_extractor.extractors.itemClicked.connect(
self.prompt_extractor)
self.welcome.step2.clicked.connect(self.load_protos)
self.proto_fs = QFileSystemModel()
self.choose_proto.protos.setModel(self.proto_fs)
self.proto_fs.directoryLoaded.connect(
self.choose_proto.protos.expandAll)
for i in range(1, self.proto_fs.columnCount()):
self.choose_proto.protos.hideColumn(i)
self.choose_proto.protos.setRootIndex(
self.proto_fs.index(str(BASE_PATH / 'protos')))
self.choose_proto.rejected.connect(partial(self.set_view,
self.welcome))
self.choose_proto.protos.clicked.connect(self.new_endpoint)
self.create_endpoint.transports.itemClicked.connect(
self.pick_transport)
self.create_endpoint.loadRespPbBtn.clicked.connect(
self.load_another_pb)
self.create_endpoint.rejected.connect(
partial(self.set_view, self.choose_proto))
self.create_endpoint.buttonBox.accepted.connect(self.write_endpoint)
self.welcome.step3.clicked.connect(self.load_endpoints)
self.choose_endpoint.rejected.connect(
partial(self.set_view, self.welcome))
self.choose_endpoint.endpoints.itemClicked.connect(self.launch_fuzzer)
self.fuzzer.rejected.connect(
partial(self.set_view, self.choose_endpoint))
self.fuzzer.fuzzFields.clicked.connect(self.fuzz_endpoint)
self.fuzzer.deleteThis.clicked.connect(self.delete_endpoint)
self.fuzzer.comboBox.activated.connect(self.launch_fuzzer)
self.fuzzer.urlField.setWordWrapMode(QTextOption.WrapAnywhere)
for tree in (self.fuzzer.pbTree, self.fuzzer.getTree):
tree.itemEntered.connect(lambda item, _: item.edit()
if hasattr(item, 'edit') else None)
tree.itemClicked.connect(
lambda item, col: item.updateCheck(col=col))
tree.header().setSectionResizeMode(QHeaderView.ResizeToContents)
self.welcome.mydirLabel.setText(self.welcome.mydirLabel.text() %
BASE_PATH)
self.welcome.mydirBtn.clicked.connect(
partial(QDesktopServices.openUrl,
QUrl.fromLocalFile(str(BASE_PATH))))
self.set_view(self.welcome)
self.exec_()
"""
Step 1 - Extract .proto structures from apps
"""
def load_extractors(self):
self.choose_extractor.extractors.clear()
for name, meta in extractors.items():
item = QListWidgetItem(meta['desc'],
self.choose_extractor.extractors)
item.setData(Qt.UserRole, name)
self.set_view(self.choose_extractor)
def prompt_extractor(self, item):
extractor = extractors[item.data(Qt.UserRole)]
inputs = []
if not assert_installed(self.view, **extractor.get('depends', {})):
return
if not extractor.get('pick_url', False):
files, mime = QFileDialog.getOpenFileNames()
for path in files:
inputs.append((path, Path(path).stem))
else:
text, good = QInputDialog.getText(self.view, ' ', 'Input an URL:')
if text:
url = urlparse(text)
inputs.append((url.geturl(), url.netloc))
if inputs:
wait = QProgressDialog('Extracting .proto structures...', None, 0,
0)
wait.setWindowTitle(' ')
self.set_view(wait)
self.worker = Worker(inputs, extractor)
self.worker.progress.connect(self.extraction_progress)
self.worker.finished.connect(self.extraction_done)
self.worker.signal_proxy.connect(self.signal_proxy)
self.worker.start()
def extraction_progress(self, info, progress):
self.view.setLabelText(info)
if progress is not None:
self.view.setRange(0, 100)
self.view.setValue(progress * 100)
else:
self.view.setRange(0, 0)
def extraction_done(self, outputs):
nb_written_all, wrote_endpoints = 0, False
for folder, output in outputs.items():
nb_written, wrote_endpoints = extractor_save(
BASE_PATH, folder, output)
nb_written_all += nb_written
if wrote_endpoints:
self.set_view(self.welcome)
QMessageBox.information(
self.view, ' ',
'%d endpoints and their <i>.proto</i> structures have been extracted! You can now reuse the <i>.proto</i>s or fuzz the endpoints.'
% nb_written_all)
elif nb_written_all:
self.set_view(self.welcome)
QMessageBox.information(
self.view, ' ',
'%d <i>.proto</i> structures have been extracted! You can now reuse the <i>.protos</i> or define endpoints for them to fuzz.'
% nb_written_all)
else:
self.set_view(self.choose_extractor)
QMessageBox.warning(
self.view, ' ',
'This extractor did not find Protobuf structures in the corresponding format for specified files.'
)
"""
Step 2 - Link .protos to endpoints
"""
# Don't load .protos from the filesystem until asked to, in order
# not to slow down startup.
def load_protos(self):
self.proto_fs.setRootPath(str(BASE_PATH / 'protos'))
self.set_view(self.choose_proto)
def new_endpoint(self, path):
if not self.proto_fs.isDir(path):
path = self.proto_fs.filePath(path)
if assert_installed(self.choose_proto, binaries=['protoc']):
if not getattr(self, 'only_resp_combo', False):
self.create_endpoint.pbRequestCombo.clear()
self.create_endpoint.pbRespCombo.clear()
has_msgs = False
for name, cls in load_proto_msgs(path):
has_msgs = True
if not getattr(self, 'only_resp_combo', False):
self.create_endpoint.pbRequestCombo.addItem(
name, (path, name))
self.create_endpoint.pbRespCombo.addItem(
name, (path, name))
if not has_msgs:
QMessageBox.warning(
self.view, ' ',
'There is no message defined in this .proto.')
return
self.create_endpoint.reqDataSubform.hide()
if not getattr(self, 'only_resp_combo', False):
self.create_endpoint.endpointUrl.clear()
self.create_endpoint.transports.clear()
self.create_endpoint.sampleData.clear()
self.create_endpoint.pbParamKey.clear()
self.create_endpoint.parsePbCheckbox.setChecked(False)
for name, meta in transports.items():
item = QListWidgetItem(meta['desc'],
self.create_endpoint.transports)
item.setData(Qt.UserRole,
(name, meta.get('ui_data_form')))
elif getattr(self, 'saved_transport_choice'):
self.create_endpoint.transports.setCurrentItem(
self.saved_transport_choice)
self.pick_transport(self.saved_transport_choice)
self.saved_transport_choice = None
self.only_resp_combo = False
self.set_view(self.create_endpoint)
def pick_transport(self, item):
name, desc = item.data(Qt.UserRole)
self.has_pb_param = desc and 'regular' in desc
self.create_endpoint.reqDataSubform.show()
if self.has_pb_param:
self.create_endpoint.pbParamSubform.show()
else:
self.create_endpoint.pbParamSubform.hide()
self.create_endpoint.sampleDataLabel.setText(
'Sample request data, one per line (in the form of %s):' % desc)
def load_another_pb(self):
self.only_resp_combo = True
self.saved_transport_choice = self.create_endpoint.transports.currentItem(
)
self.set_view(self.choose_proto)
def write_endpoint(self):
request_pb = self.create_endpoint.pbRequestCombo.itemData(
self.create_endpoint.pbRequestCombo.currentIndex())
url = self.create_endpoint.endpointUrl.text()
transport = self.create_endpoint.transports.currentItem()
sample_data = self.create_endpoint.sampleData.toPlainText()
pb_param = self.create_endpoint.pbParamKey.text()
has_resp_pb = self.create_endpoint.parsePbCheckbox.isChecked()
resp_pb = self.create_endpoint.pbRespCombo.itemData(
self.create_endpoint.pbRespCombo.currentIndex())
if not (request_pb and urlparse(url).netloc and transport and
(not self.has_pb_param or pb_param) and
(not has_resp_pb or resp_pb)):
QMessageBox.warning(
self.view, ' ', 'Please fill all relevant information fields.')
else:
json = {
'request': {
'transport':
transport.data(Qt.UserRole)[0],
'proto_path':
request_pb[0].replace(str(BASE_PATH / 'protos'),
'').strip('/\\'),
'proto_msg':
request_pb[1],
'url':
url
}
}
if self.has_pb_param:
json['request']['pb_param'] = pb_param
sample_data = list(filter(None, sample_data.split('\n')))
if sample_data:
transport_obj = transports[transport.data(Qt.UserRole)[0]]
transport_obj = transport_obj['func'](pb_param, url)
for sample_id, sample in enumerate(sample_data):
try:
sample = transport_obj.serialize_sample(sample)
except Exception:
return QMessageBox.warning(
self.view, ' ',
'Some of your sample data is not in the specified format.'
)
if not sample:
return QMessageBox.warning(
self.view, ' ',
"Some of your sample data didn't contain the Protobuf parameter key you specified."
)
sample_data[sample_id] = sample
json['request']['samples'] = sample_data
if has_resp_pb:
json['response'] = {
'format':
'raw_pb',
'proto_path':
resp_pb[0].replace(str(BASE_PATH / 'protos'),
'').strip('/\\'),
'proto_msg':
resp_pb[1]
}
insert_endpoint(BASE_PATH / 'endpoints', json)
QMessageBox.information(self.view, ' ',
'Endpoint created successfully.')
self.set_view(self.welcome)
def load_endpoints(self):
self.choose_endpoint.endpoints.clear()
for name in listdir(str(BASE_PATH / 'endpoints')):
if name.endswith('.json'):
item = QListWidgetItem(
name.split('.json')[0], self.choose_endpoint.endpoints)
item.setFlags(item.flags() & ~Qt.ItemIsEnabled)
pb_msg_to_endpoints = defaultdict(list)
with open(str(BASE_PATH / 'endpoints' / name)) as fd:
for endpoint in load(fd):
pb_msg_to_endpoints[endpoint['request']['proto_msg'].
split('.')[-1]].append(endpoint)
for pb_msg, endpoints in pb_msg_to_endpoints.items():
item = QListWidgetItem(' ' * 4 + pb_msg,
self.choose_endpoint.endpoints)
item.setFlags(item.flags() & ~Qt.ItemIsEnabled)
for endpoint in endpoints:
item = QListWidgetItem(
' ' * 8 +
(urlparse(endpoint['request']['url']).path or '/'),
self.choose_endpoint.endpoints)
item.setData(Qt.UserRole, endpoint)
self.set_view(self.choose_endpoint)
"""
Step 3: Fuzz and test endpoints live
"""
def launch_fuzzer(self, item):
if type(item) == int:
data, sample_id = self.fuzzer.comboBox.itemData(item)
else:
data, sample_id = item.data(Qt.UserRole), 0
if data and assert_installed(self.view, binaries=['protoc']):
self.current_req_proto = BASE_PATH / 'protos' / data['request'][
'proto_path']
self.pb_request = load_proto_msgs(self.current_req_proto)
self.pb_request = dict(
self.pb_request)[data['request']['proto_msg']]()
if data.get('response') and data['response']['format'] == 'raw_pb':
self.pb_resp = load_proto_msgs(BASE_PATH / 'protos' /
data['response']['proto_path'])
self.pb_resp = dict(
self.pb_resp)[data['response']['proto_msg']]
self.pb_param = data['request'].get('pb_param')
self.base_url = data['request']['url']
self.endpoint = data
self.transport_meta = transports[data['request']['transport']]
self.transport = self.transport_meta['func'](self.pb_param,
self.base_url)
sample = ''
if data['request'].get('samples'):
sample = data['request']['samples'][sample_id]
self.get_params = self.transport.load_sample(
sample, self.pb_request)
# Get initial data into the Protobuf tree view
self.fuzzer.pbTree.clear()
self.parse_desc(self.pb_request.DESCRIPTOR, self.fuzzer.pbTree)
self.parse_fields(self.pb_request)
# Do the same for transport-specific data
self.fuzzer.getTree.clear()
if self.transport_meta.get('ui_tab'):
self.fuzzer.tabs.setTabText(1, self.transport_meta['ui_tab'])
if self.get_params:
for key, val in self.get_params.items():
ProtocolDataItem(self.fuzzer.getTree, key, val, self)
else:
self.fuzzer.tabs.setTabText(1, '(disabled)')
# how to hide it ?
# Fill the request samples combo box if we're loading a new
# endpoint.
if type(item) != int:
if len(data['request'].get('samples', [])) > 1:
self.fuzzer.comboBox.clear()
for sample_id, sample in enumerate(
data['request']['samples']):
self.fuzzer.comboBox.addItem(
sample[self.pb_param] if self.pb_param else
str(sample), (data, sample_id))
self.fuzzer.comboBoxLabel.show()
self.fuzzer.comboBox.show()
else:
self.fuzzer.comboBoxLabel.hide()
self.fuzzer.comboBox.hide()
self.set_view(self.fuzzer)
self.fuzzer.frame.setUrl(QUrl("about:blank"))
self.update_fuzzer()
"""
Parsing and rendering the Protobuf message to a tree view:
Every Protobuf field is fed to ProtobufItem (a class inheriting
QTreeWidgetItem), and the created object is saved in the _items
property of the corresponding descriptor.
"""
# First, parse the descriptor (structure) of the Protobuf message.
def parse_desc(self, msg, item, path=[]):
for ds in msg.fields:
new_item = ProtobufItem(item, ds, self, path)
if ds.type == ds.TYPE_MESSAGE and ds.full_name not in path:
self.parse_desc(ds.message_type, new_item,
path + [ds.full_name])
# Then, parse the fields (contents) of the Protobuf message.
def parse_fields(self, msg, path=[]):
for ds, val in msg.ListFields():
if ds.label == ds.LABEL_REPEATED:
for val_index, val_value in enumerate(val):
if ds.type == ds.TYPE_MESSAGE:
ds._items[tuple(path)].setExpanded(True)
ds._items[tuple(path)].setDefault(parent=msg,
msg=val,
index=val_index)
self.parse_fields(val_value, path + [ds.full_name])
else:
ds._items[tuple(path)].setDefault(val_value,
parent=msg,
msg=val,
index=val_index)
ds._items[tuple(path)].duplicate(True)
else:
if ds.type == ds.TYPE_MESSAGE:
ds._items[tuple(path)].setExpanded(True)
ds._items[tuple(path)].setDefault(parent=msg, msg=val)
self.parse_fields(val, path + [ds.full_name])
else:
ds._items[tuple(path)].setDefault(val, parent=msg, msg=val)
def update_fuzzer(self):
resp = self.transport.perform_request(self.pb_request, self.get_params)
data, text, url, mime = resp.content, resp.text, resp.url, resp.headers[
'Content-Type'].split(';')[0]
meta = '%s %d %08x\n%s' % (mime, len(data), crc32(data) & 0xffffffff,
resp.url)
self.fuzzer.urlField.setText(meta)
self.fuzzer.frame.update_frame(data, text, url, mime,
getattr(self, 'pb_resp', None))
def fuzz_endpoint(self):
QMessageBox.information(self.view, ' ',
'Automatic fuzzing is not implemented yet.')
def delete_endpoint(self):
if QMessageBox.question(self.view, ' ',
'Delete this endpoint?') == QMessageBox.Yes:
path = str(BASE_PATH / 'endpoints' /
(urlparse(self.base_url).netloc + '.json'))
with open(path) as fd:
json = load(fd)
json.remove(self.endpoint)
with open(path, 'w') as fd:
dump(json, fd, ensure_ascii=False, indent=4)
if not json:
remove(path)
self.load_endpoints()
"""
Utility methods follow
"""
def set_view(self, view):
if hasattr(self, 'view'):
self.view.hide()
view.show()
self.view = view
resolution = QDesktopWidget().screenGeometry()
view.move((resolution.width() / 2) - (view.frameSize().width() / 2),
(resolution.height() / 2) - (view.frameSize().height() / 2))
"""
signal() can't be called from inside a thread, and some
extractors need it in order not to have their GUI child process
interrupt signal catched by our main thread, so here is an ugly
way to reach signal() through a slot.
"""
def signal_proxy(self, *args):
signal(*args)
class FileChooser(QWidget):
fileOpened = pyqtSignal(str)
def __init__(self, parent=None):
super().__init__(parent)
# TODO: migrate to FolderComboBox?
self.folderBox = QComboBox(self)
self.explorerTree = FileTreeView(self)
self.explorerTree.doubleClickCallback = self._fileOpened
self.explorerModel = QFileSystemModel(self)
self.explorerModel.setFilter(QDir.AllDirs | QDir.Files
| QDir.NoDotAndDotDot)
self.explorerModel.setNameFilters(["*.py"])
self.explorerModel.setNameFilterDisables(False)
self.explorerTree.setModel(self.explorerModel)
for index in range(1, self.explorerModel.columnCount()):
self.explorerTree.hideColumn(index)
self.setCurrentFolder()
self.folderBox.currentIndexChanged[int].connect(
self.updateCurrentFolder)
layout = QVBoxLayout(self)
layout.addWidget(self.folderBox)
layout.addWidget(self.explorerTree)
layout.setContentsMargins(5, 5, 0, 0)
def _fileOpened(self, modelIndex):
path = self.explorerModel.filePath(modelIndex)
if os.path.isfile(path):
self.fileOpened.emit(path)
def currentFolder(self):
return self.explorerModel.rootPath()
def setCurrentFolder(self, path=None):
if path is None:
app = QApplication.instance()
path = app.getScriptsDirectory()
else:
assert os.path.isdir(path)
self.explorerModel.setRootPath(path)
self.explorerTree.setRootIndex(self.explorerModel.index(path))
self.folderBox.blockSignals(True)
self.folderBox.clear()
style = self.style()
dirIcon = style.standardIcon(style.SP_DirIcon)
self.folderBox.addItem(dirIcon, os.path.basename(path))
self.folderBox.insertSeparator(1)
self.folderBox.addItem(self.tr("Browse…"))
self.folderBox.setCurrentIndex(0)
self.folderBox.blockSignals(False)
def updateCurrentFolder(self, index):
if index < self.folderBox.count() - 1:
return
path = QFileDialog.getExistingDirectory(self,
self.tr("Choose Directory"),
self.currentFolder(),
QFileDialog.ShowDirsOnly)
if path:
QSettings().setValue("scripting/path", path)
self.setCurrentFolder(path)
class ProjectTreeView(QTreeView):
def __init__(self, parent: QWidget):
super().__init__(parent)
self.doubleClicked.connect(self.file_or_directory_double_clicked)
signals().folder_opened_signal.connect(self.folder_opened)
self.file_system_model = QFileSystemModel(self)
# noinspection PyUnresolvedReferences
self.file_system_model.directoryLoaded.connect(self.directory_loaded)
self.file_system_model.setFilter(QDir.NoDotAndDotDot | QDir.AllEntries)
self.setModel(self.file_system_model)
self.setHeaderHidden(True)
# Hide all but the name column of the tree view.
for i in range(1, self.file_system_model.columnCount()):
self.hideColumn(i)
self.set_path(user_settings().get(Key.last_folder_opened,
Default.last_folder_opened))
self.loading_previous_folder = True
def set_path(self, path: str):
""" Change the path of the project window. """
self.loading_previous_folder = False # Do not try to reload the previous user's settings on this new folder.
self.file_system_model.setRootPath(path)
# This shows the selected path's contents rather than the path itself.
self.setRootIndex(self.file_system_model.index(path))
self.clearSelection()
# noinspection PyUnusedLocal
@pyqtSlot(str)
def directory_loaded(self, path):
# We have to restore the tree view here to ensure it is expanded
# only when the directory is ready to be searched.
if self.loading_previous_folder:
user_settings().restore_widget(self, Key.project_tree)
self.loading_previous_folder = False
@pyqtSlot(QModelIndex)
def file_or_directory_double_clicked(self, index: QModelIndex):
""" Send a signal so we can react to the user's double click on a file in the project view. """
path = self.file_system_model.filePath(index)
if path and os.path.isfile(path):
signals().file_selected_signal.emit(path)
@pyqtSlot(str)
def folder_opened(self, folder: str):
self.set_path(folder)
user_settings().save(Key.last_folder_opened, folder)
# noinspection PyPep8Naming
def saveState(self):
is_expanded: List[str] = []
for index in self.file_system_model.persistentIndexList():
if self.isExpanded(index):
is_expanded.append(self.file_system_model.filePath(index))
return pickle.dumps(is_expanded)
# noinspection PyPep8Naming
def restoreState(self, is_expanded: bytes):
is_expanded_list = pickle.loads(is_expanded)
self.setUpdatesEnabled(False)
for path in is_expanded_list:
if os.path.isdir(path):
index = self.file_system_model.index(path)
self.setExpanded(index, True)
self.setUpdatesEnabled(True)
def closeEvent(self, event: QCloseEvent):
user_settings().save_widget(self, Key.project_tree)
super().closeEvent(event)
class MenuLeft(QTreeView):
"""
displays interactive directory on
left side of text editor
"""
def __init__(self, layout_props, document, file_manager, abs_path=None):
"""
creates the directory display
:param file_manager: instance of FileManager class - manages all file communication
:param abs_path: default path to file being displayed
:return: returns nothing
"""
super().__init__()
logging.debug("Creating Directory Viewer")
self.layout_props = layout_props
self.document = document
self.fileManager = file_manager
if abs_path is None:
abs_path = QDir.currentPath()
self.model = QFileSystemModel()
self.initUI()
self.updateDirectory(abs_path)
self.updateAppearance()
def initUI(self):
"""
Initializes the Directory Viewer properties, signals, and model
:return: returns nothing
"""
logging.debug("Initializing Directory Viewer Props")
self.setModel(self.model)
# Default hide all columns except Name
for i in range(1, self.model.columnCount()):
name = self.model.headerData(i, Qt.Horizontal)
if name not in self.layout_props.getDefaultLeftMenuCols():
self.hideColumn(i)
self.setAnimated(True)
self.setIndentation(10)
self.setSortingEnabled(True)
self.setSelectionMode(QAbstractItemView.SingleSelection)
self.setExpandsOnDoubleClick(True)
# Expand or collapse directory on click
self.doubleClicked.connect(self.onClickItem)
# Shortcut for pressing enter on directory
shortcut = QShortcut(Qt.Key_Return, self)
shortcut.activated.connect(self.onClickItem)
def updateDirectory(self, abs_path: str):
"""
Updates the path of the model to the given path, sorts, and looks for the encryption key.
:param abs_path: default path to file being displayed
:return: returns nothing
"""
logging.info(abs_path)
self.model.setRootPath(abs_path)
self.setRootIndex(self.model.index(abs_path))
self.sortByColumn(0, Qt.AscendingOrder)
self.fileManager.encryptor = None
# Check for encryption key in Workspace
path_key = os.path.join(abs_path, ".leafCryptoKey")
if os.path.exists(path_key):
logging.debug("Encryption key found! %s", path_key)
with open(path_key, 'r') as f:
key = f.read()
self.fileManager.encryptor = Encryptor(key)
else:
logging.info("Workspace not encrypted")
def onClickItem(self, index: QModelIndex = None):
"""
functionality of double click on directory
:param index: location of filePath
:return: returns nothing
"""
# If coming from Enter Pressed, resolve index
if index is None:
index = self.selectionModel().currentIndex()
path = self.model.filePath(index)
# Toggle expand/collapse directory
if not self.model.isDir(index):
logging.info("Opening document")
self.fileManager.openDocument(self.document, path)
def toggleHeaderColByName(self, name: str):
"""
Shows or Hides a column based on it's name
"""
logging.debug("Toggling header column %s", name)
for i in range(1, self.model.columnCount()):
col_name = self.model.headerData(i, Qt.Horizontal)
if col_name == name:
self.setColumnHidden(i, not self.isColumnHidden(i))
def resizeColumnsToContent(self):
"""
Resizes all columns to fit content
"""
logging.debug("Resizing columns to contents")
for i in range(1, self.model.columnCount()):
self.resizeColumnToContents(i)
def selectItemFromPath(self, path: str):
"""
"Programmatically selects a path in the File System model
"""
logging.debug("Selecting path %s", path)
self.selectionModel().blockSignals(True)
self.setCurrentIndex(self.model.index(path))
self.selectionModel().blockSignals(False)
def updateAppearance(self):
"""
Updates the layout appearance based on properties
"""
logging.debug("Setting up appearance")
prop_header_margin = str(
self.layout_props.getDefaultLeftMenuHeaderMargin())
prop_header_color = self.layout_props.getDefaultHeaderColor()
prop_item_height = str(self.layout_props.getDefaultItemHeight())
prop_item_select_color = self.layout_props.getDefaultSelectColor()
prop_item_hover_color = self.layout_props.getDefaultHoverColor()
style = "QTreeView::item { height: " + prop_item_height + "px; }" + \
"QTreeView::item:selected {" \
"background-color: " + prop_item_select_color + "; }" + \
"QTreeView::item:hover:!selected { " + \
"background-color: " + prop_item_hover_color + "; }" + \
"QHeaderView::section { margin: " + prop_header_margin + ";" + \
" margin-left: 0; " + \
" background-color: " + prop_header_color + ";" + \
" color: white; }"
self.setStyleSheet(style)
