maq/LinkScope_Client
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

The "make collectors work" commit.

Browse Source
This commit is contained in:
AccentuSoft
2022-05-06 04:23:00 +03:00
parent 194254ab55
commit 95b033d8bb
3 changed files with 199 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
Core/FrontendCommunicationsHandler.py
View File

@@ -50,6 +50,8 @@ class CommunicationsHandler(QtCore.QObject):
receive_question_answer = QtCore.Signal(dict)
receive_chat_message = QtCore.Signal(str)
receive_collectors_signal = QtCore.Signal(dict, dict)
receive_start_collector_signal = QtCore.Signal(str, str, str, list, dict)
receive_collector_result_signal = QtCore.Signal(str, str, str, list)
receive_resolutions_signal = QtCore.Signal(dict)
receive_completed_resolution_result_signal = QtCore.Signal(str, list)
receive_completed_resolution_string_result_signal = QtCore.Signal(str, str)
@@ -86,6 +88,8 @@ class CommunicationsHandler(QtCore.QObject):
self.receive_question_answer.connect(self.mainWindow.questionAnswerListener)
self.receive_chat_message.connect(self.mainWindow.receiveChatMessage)
self.receive_collectors_signal.connect(self.mainWindow.addCollectorsFromServerListener)
self.receive_start_collector_signal.connect(self.mainWindow.startNewCollectorListener)
self.receive_collector_result_signal.connect(self.mainWindow.receiveCollectorResultListener)
self.receive_resolutions_signal.connect(self.mainWindow.addResolutionsFromServerListener)
self.receive_completed_resolution_result_signal.connect(self.mainWindow.resolutionSignalListener)
self.receive_completed_resolution_string_result_signal.connect(self.mainWindow.resolutionSignalListener)
@@ -317,17 +321,21 @@ class CommunicationsHandler(QtCore.QObject):
}}
self.transmitMessage(message)
def receiveCollectors(self, server_collectors) -> None:
print('RECVD COLL', server_collectors)
self.receive_collectors_signal.emit(server_collectors)
def receiveStartCollector(self, collector_category: str, collector_name: str, collector_uid: str,
collector_entities: list, collector_parameters: dict):
self.receive_start_collector_signal.emit(collector_category, collector_name, collector_uid,
collector_entities, collector_parameters)
def receiveCollectors(self, server_collectors: dict, continuing_collectors_info: dict) -> None:
self.receive_collectors_signal.emit(server_collectors, continuing_collectors_info)
def startServerCollector(self, collector_name: str, collector_entities: list, collector_parameters: dict,
continueTimestamp: int) -> None:
resolution_entities_to_send = []
continueTimestamp: int = 0) -> None:
collector_entities_to_send = []
for entity in collector_entities:
try:
dereferenced_entity = dict(entity)
resolution_entities_to_send.append(dereferenced_entity)
collector_entities_to_send.append(dereferenced_entity)
# Icon is not necessary for any collector as of now: 2022/4/3.
# Cutting it out saves data.
dereferenced_entity['Icon'] = ''
@@ -336,24 +344,21 @@ class CommunicationsHandler(QtCore.QObject):
message = {'Operation': 'Start Server Collector',
'Arguments': {
'collector_name': collector_name,
'collector_entities': resolution_entities_to_send,
'collector_entities': collector_entities_to_send,
'collector_parameters': collector_parameters,
'continue_timestamp': continueTimestamp
'continue_time': continueTimestamp
}}
self.transmitMessage(message)
def stopServerCollector(self, collector_name: str, collector_uid: str) -> None:
def stopServerCollector(self, collector_uid: str) -> None:
message = {"Operation": "Stop Server Collector",
"Arguments": {
'collector_name': collector_name,
'collector_uid': collector_uid
}}
self.transmitMessage(message)
def receiveCollectorResult(self, collector_name: str, results: list) -> None:
print('Received results from ' + str(collector_name))
print('Results:', results)
# TODO
def receiveCollectorResult(self, collector_name: str, collector_uid: str, timestamp: str, results: list) -> None:
self.receive_collector_result_signal.emit(collector_name, collector_uid, timestamp, results)
def askServerForResolutions(self) -> None:
message = {"Operation": "Get Server Resolutions",
@@ -659,6 +664,8 @@ class CommunicationsHandler(QtCore.QObject):
self.receiveFileUploadAbort(**arguments)
elif operation == "Get Server Collectors":
self.receiveCollectors(**arguments)
elif operation == "Start Collector":
self.receiveStartCollector(**arguments)
else:
self.mainWindow.MESSAGEHANDLER.warning('Unhandled message: ' + str(message) +
' On Operation: ' + str(operation))
@@ -727,6 +734,9 @@ class CommunicationsHandler(QtCore.QObject):
file_name = message.split(': ', 1)[1]
# Remove file from uploading files list.
self.file_upload_abort_signal.emit(file_name)
elif operation == 'Stop Collector':
# No need to do anything here - stopping collectors is only done by the client.
pass
else:
self.mainWindow.MESSAGEHANDLER.warning('Unhandled status message: ' + message +
' Code: ' + str(status_code) +

174
Core/Interface/MenuBar.py
View File

@@ -907,8 +907,10 @@ class MenuBar(QtWidgets.QMenuBar):
def manageCollectors(self) -> None:
with self.parent().serverCollectorsLock:
currentCollectors = dict(self.parent().collectors)
print('CURR COLL', currentCollectors)
collectorsDialog = CollectorsDialog(currentCollectors if self.parent().FCOM.isConnected() else None)
runningCollectors = dict(self.parent().runningCollectors)
collectorsDialog = CollectorsDialog(self.parent(),
currentCollectors if self.parent().FCOM.isConnected() else None,
runningCollectors)
collectorsDialog.exec()
def selectAllNodes(self) -> None:
@@ -1754,14 +1756,17 @@ class ViewAndStopResolutionsDialogOption(QtWidgets.QPushButton):
class CollectorsDialog(QtWidgets.QDialog):
def __init__(self, collectorsDict: dict = None):
def __init__(self, mainWindow, collectorsDict: dict = None, runningCollectors: dict = None):
super(CollectorsDialog, self).__init__()
self.mainWindow = mainWindow
self.setStyleSheet(Stylesheets.MAIN_WINDOW_STYLESHEET)
self.baseLayout = QtWidgets.QVBoxLayout()
self.setLayout(self.baseLayout)
self.setModal(True)
self.runningCollectorTreeItems = {}
collectorsLabel = QtWidgets.QLabel("Collectors")
collectorsLabel.setAlignment(QtCore.Qt.AlignCenter)
self.baseLayout.addWidget(collectorsLabel)
if collectorsDict is None:
@@ -1771,27 +1776,45 @@ class CollectorsDialog(QtWidgets.QDialog):
connectedToServerFormWidget = QtWidgets.QWidget()
self.connectedToServerFormWidgetLayout = QtWidgets.QVBoxLayout()
for collector in collectorsDict:
newCollectorWidget = QtWidgets.QWidget()
newCollectorWidgetLayout = QtWidgets.QGridLayout()
newCollectorWidget.setLayout(newCollectorWidgetLayout)
for category in collectorsDict:
categoryLabel = QtWidgets.QLabel("Category: " + str(category))
self.connectedToServerFormWidgetLayout.addWidget(categoryLabel)
for collector in collectorsDict[category]:
newCollectorWidget = QtWidgets.QWidget()
newCollectorWidgetLayout = QtWidgets.QGridLayout()
newCollectorWidget.setLayout(newCollectorWidgetLayout)
newCollectorLabel = QtWidgets.QLabel(collectorsDict[collector]['name'])
newCollectorLabel.setToolTip(collectorsDict[collector]['description'])
newCollectorButton = QtWidgets.QPushButton('Create New')
newCollectorButton.clicked.connect(lambda: self.close()) # TODO
newCollectorLabel = QtWidgets.QLabel(collectorsDict[category][collector]['name'])
newCollectorLabel.setToolTip(collectorsDict[category][collector]['description'])
newCollectorButton = QtWidgets.QPushButton('Create New')
newCollectorButton.clicked.connect(lambda: self.startSelectedCollector(
collectorsDict[category][collector]))
newCollectorWidgetLayout.addWidget(newCollectorLabel, 0, 0)
newCollectorWidgetLayout.addWidget(newCollectorButton, 0, 1)
newCollectorWidgetLayout.addWidget(newCollectorLabel, 0, 0)
newCollectorWidgetLayout.addWidget(newCollectorButton, 0, 1)
newCollectorInstanceTree = QtWidgets.QTreeWidget()
newCollectorInstanceTree.setColumnCount(2)
newCollectorInstanceTree.setHeaderLabels(['UID', 'Stop Button']) # TODO
newCollectorInstanceTree = QtWidgets.QTreeWidget()
newCollectorInstanceTree.setColumnCount(2)
newCollectorInstanceTree.setHeaderLabels(['UID', 'Stop Button'])
newCollectorInstanceTree.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
newCollectorInstanceTree.header().setStretchLastSection(False)
newCollectorInstanceTree.header().setSectionResizeMode(0, QtWidgets.QHeaderView.Stretch)
newCollectorWidgetLayout.addWidget(newCollectorInstanceTree, 1, 0, 2, 2)
self.connectedToServerFormWidgetLayout.addWidget(newCollectorWidget)
if runningCollectors is not None:
for runningCollectorCategory in runningCollectors:
for runningCollector in runningCollectors[runningCollectorCategory][collector]:
newTreeItem = QtWidgets.QTreeWidgetItem(newCollectorInstanceTree)
newTreeItem.setText(0, runningCollector['uid'])
stopButton = QtWidgets.QPushButton("Stop")
stopButton.clicked.connect(lambda: self.stopSelectedCollector(runningCollector['uid']))
newCollectorInstanceTree.setItemWidget(newTreeItem, 1, stopButton)
self.runningCollectorTreeItems[runningCollector['uid']] = (newTreeItem,
newCollectorInstanceTree)
self.connectedToServerFormWidgetLayout.addWidget(newCollectorWidget)
connectedToServerFormWidget.setLayout(self.connectedToServerFormWidgetLayout)
# TODO: manage collectors.
self.baseLayout.addWidget(connectedToServerFormWidget)
@@ -1800,18 +1823,45 @@ class CollectorsDialog(QtWidgets.QDialog):
self.baseLayout.addWidget(closeButton)
def startSelectedCollector(self, collectorToStartDict: dict):
newCollector = CollectorStartDialog(self.mainWindow.LENTDB, collectorToStartDict)
if newCollector.exec_():
collector_name = collectorToStartDict['name']
try:
self.mainWindow.FCOM.startServerCollector(collector_name,
newCollector.chosenItems,
newCollector.chosenParameters)
self.mainWindow.MESSAGEHANDLER.info('Starting server collector: ' + collector_name)
except Exception as e:
self.mainWindow.MESSAGEHANDLER.error('Error starting server collector: ' + str(e))
def stopSelectedCollector(self, collectorToStop: str):
self.mainWindow.MESSAGEHANDLER.info('Stopping server collector with UID: ' + collectorToStop)
self.mainWindow.FCOM.stopServerCollector(collectorToStop)
self.runningCollectorTreeItems[collectorToStop][1].takeTopLevelItem(
self.runningCollectorTreeItems[collectorToStop][1].indexOfTopLevelItem(
self.runningCollectorTreeItems[collectorToStop][0]))
currentClientCollectors = self.mainWindow.getClientCollectors()
try:
currentClientCollectors.pop(collectorToStop)
self.mainWindow.setClientCollectors(currentClientCollectors)
except KeyError:
self.mainWindow.MESSAGEHANDLER.error('Trying to stop Collector that was never ran. Was confirmation that '
'the collector was started by the server received?', popUp=False)
class CollectorStartDialog(QtWidgets.QDialog):
def __init__(self, collectorDict: dict):
def __init__(self, entityDB, collectorDict: dict):
super(CollectorStartDialog, self).__init__()
self.entityDB = entityDB
self.setStyleSheet(Stylesheets.MAIN_WINDOW_STYLESHEET)
self.setModal(True)
self.setWindowTitle('Collector Wizard')
self.parametersList = []
# Have two separate dicts for readability.
self.chosenParameters = {}
self.parametersToRemember = {}
self.chosenItems = []
dialogLayout = QtWidgets.QGridLayout()
self.setLayout(dialogLayout)
@@ -1820,44 +1870,60 @@ class CollectorStartDialog(QtWidgets.QDialog):
dialogLayout.setRowStretch(0, 1)
dialogLayout.setColumnStretch(0, 1)
if includeEntitySelector is not None and originTypes is not None:
entitySelectTab = QtWidgets.QWidget()
entitySelectTab.setLayout(QtWidgets.QVBoxLayout())
labelText = ""
if resolutionDescription is not None:
labelText += resolutionDescription + "\n\n"
labelText += 'Select the entities to use for this resolution.\nAccepted Origin Types: ' + \
', '.join(originTypes)
entitySelectTabLabel = QtWidgets.QLabel(labelText)
entitySelectTabLabel.setWordWrap(True)
entitySelectTabLabel.setMaximumWidth(600)
entitySelectTab = QtWidgets.QWidget()
entitySelectTab.setLayout(QtWidgets.QVBoxLayout())
originTypes = collectorDict['originTypes']
labelText = collectorDict['description'] + "\n\n"
labelText += 'Select the entities to use for this collector.\nAccepted Origin Types: ' + \
', '.join(collectorDict['originTypes'])
entitySelectTabLabel = QtWidgets.QLabel(labelText)
entitySelectTabLabel.setWordWrap(True)
entitySelectTabLabel.setMaximumWidth(600)
entitySelectTabLabel.setAlignment(QtCore.Qt.AlignCenter)
entitySelectTab.layout().addWidget(entitySelectTabLabel)
entitySelectTabLabel.setAlignment(QtCore.Qt.AlignCenter)
entitySelectTab.layout().addWidget(entitySelectTabLabel)
self.entitySelector = QtWidgets.QListWidget()
self.entitySelector.addItems(includeEntitySelector)
self.entitySelector = QtWidgets.QTreeWidget()
self.entitySelector.setHeaderLabels(['Primary Field', 'Entity Type', 'Icon'])
self.entitySelector.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
self.entitySelector.setSortingEnabled(True)
# Stretch the first column, since it contains the primary field.
self.entitySelector.header().setStretchLastSection(False)
self.entitySelector.header().setSectionResizeMode(0, QtWidgets.QHeaderView.Stretch)
relevantEntityFields = [(entity['uid'], entity[list(entity)[1]], entity['Entity Type'], entity['Icon'])
for entity in entityDB.getAllEntities()
if entity['Entity Type'] in originTypes or originTypes == '*']
for eligibleEntity in relevantEntityFields:
newTreeWidgetItem = QtWidgets.QTreeWidgetItem(self.entitySelector)
newTreeWidgetItemPixmap = QtGui.QPixmap()
newTreeWidgetItemPixmap.loadFromData(eligibleEntity[3])
newTreeWidgetItem.setText(0, eligibleEntity[1])
newTreeWidgetItem.setText(1, eligibleEntity[2])
newTreeWidgetItem.setIcon(2, newTreeWidgetItemPixmap)
# Hidden, so we can pull the UID later.
newTreeWidgetItem.setText(3, eligibleEntity[0])
self.entitySelector.setSelectionMode(self.entitySelector.MultiSelection)
entitySelectTab.layout().addWidget(self.entitySelector)
self.entitySelector.setSelectionMode(self.entitySelector.MultiSelection)
entitySelectTab.layout().addWidget(self.entitySelector)
self.childWidget.addTab(entitySelectTab, 'Entities')
self.childWidget.addTab(entitySelectTab, 'Entities')
for key in properties:
parameters = collectorDict['parameters']
for key in parameters:
propertyWidget = QtWidgets.QWidget()
propertyKeyLayout = QtWidgets.QVBoxLayout()
propertyWidget.setLayout(propertyKeyLayout)
propertyLabel = QtWidgets.QLabel(properties[key].get('description'))
propertyLabel = QtWidgets.QLabel(parameters[key].get('description'))
propertyLabel.setWordWrap(True)
propertyLabel.setMaximumWidth(600)
propertyLabel.setAlignment(QtCore.Qt.AlignCenter)
propertyKeyLayout.addWidget(propertyLabel)
propertyType = properties[key].get('type')
propertyValue = properties[key].get('value')
propertyDefaultValue = properties[key].get('default')
propertyType = parameters[key].get('type')
propertyValue = parameters[key].get('value')
propertyDefaultValue = parameters[key].get('default')
if propertyType == 'String':
propertyInputField = StringPropertyInput(propertyValue, propertyDefaultValue)
@@ -1875,14 +1941,10 @@ class CollectorStartDialog(QtWidgets.QDialog):
propertyKeyLayout.addWidget(propertyInputField)
propertyInputField.setStyleSheet(Stylesheets.CHECK_BOX_STYLESHEET)
rememberChoiceCheckbox = QtWidgets.QCheckBox('Remember Choice')
rememberChoiceCheckbox.setStyleSheet(Stylesheets.CHECK_BOX_STYLESHEET)
rememberChoiceCheckbox.setChecked(False)
propertyKeyLayout.addWidget(rememberChoiceCheckbox)
propertyKeyLayout.setStretch(1, 1)
self.childWidget.addTab(propertyWidget, key)
self.parametersList.append((key, propertyInputField, rememberChoiceCheckbox))
self.parametersList.append((key, propertyInputField))
nextButton = QtWidgets.QPushButton('Next')
nextButton.setStyleSheet(Stylesheets.BUTTON_STYLESHEET_2)
@@ -1915,23 +1977,21 @@ class CollectorStartDialog(QtWidgets.QDialog):
self.childWidget.setCurrentIndex(currentIndex - 1)
def accept(self) -> None:
for resolutionParameterName, resolutionParameterInput, resolutionParameterRemember in self.parametersList:
for item in self.entitySelector.selectedItems():
self.chosenItems.append(self.entityDB.getEntity(item.text(3)))
for resolutionParameterName, resolutionParameterInput in self.parametersList:
value = resolutionParameterInput.getValue()
if value == '':
msgBox = QtWidgets.QMessageBox()
msgBox.setModal(True)
QtWidgets.QMessageBox.warning(msgBox,
"Not all parameters filled in",
"Some of the required parameters for the resolution have been left blank."
"Not all parameters were filled in",
"Some of the required parameters for the collector have been left blank."
" Please fill them in before proceeding.")
return
self.chosenParameters[resolutionParameterName] = value
if resolutionParameterRemember.isChecked():
self.parametersToRemember[resolutionParameterName] = value
super(ResolutionParametersSelector, self).accept()
super(CollectorStartDialog, self).accept()
class ImportLinksFromCSVFile(QtWidgets.QDialog):

69
LinkScope.py
View File

@@ -826,7 +826,7 @@ class MainWindow(QtWidgets.QMainWindow):
self.saveProject()
def changeGraphics(self):
def changeGraphics(self) -> None:
settingsDialog = GraphicsEditDialog(self.SETTINGS, self.RESOURCEHANDLER)
settingsConfirm = settingsDialog.exec()
@@ -1162,6 +1162,39 @@ class MainWindow(QtWidgets.QMainWindow):
if resolutionThread[0].isFinished() and resolutionThread[1] is False:
self.resolutions.remove(resolutionThread)
def getClientCollectors(self) -> dict:
with self.serverCollectorsLock:
try:
clientCollectorUIDs = literal_eval(self.SETTINGS.value('Project/Server/Collectors'))
if not isinstance(clientCollectorUIDs, dict):
raise ValueError('Collectors were not saved in the correct format.')
except Exception as e:
self.MESSAGEHANDLER.error('Unable to load Collectors from Settings file.',
popUp=False,
exc_info=False)
self.MESSAGEHANDLER.debug('Cannot eval Project/Server/Collectors setting: ' + str(e))
clientCollectorUIDs = {}
return clientCollectorUIDs
def setClientCollectors(self, newClientCollectorsDict: dict) -> None:
if not isinstance(newClientCollectorsDict, dict):
self.MESSAGEHANDLER.error('Unable to save Collectors to Settings: Invalid format.')
self.MESSAGEHANDLER.debug('Collectors argument: ' + str(newClientCollectorsDict))
self.MESSAGEHANDLER.debug('Collectors format: ' + str(type(newClientCollectorsDict)))
else:
with self.serverCollectorsLock:
self.SETTINGS.setValue('Project/Server/Collectors', str(newClientCollectorsDict))
self.SETTINGS.save()
self.MESSAGEHANDLER.info('Client Collectors state updated.')
def receiveCollectorResultListener(self, collector_name: str, collector_uid: str, timestamp: str, results: list):
# Signal ints are 4 bytes long and signed, so we use strings to communicate timestamps.
currentCollectors = self.getClientCollectors()
currentCollectors[collector_uid] = int(timestamp)
self.setClientCollectors(currentCollectors)
self.centralWidget().tabbedPane.facilitateResolution('Collector ' + str(collector_uid), results)
self.notifyUser("New entities discovered by collector: " + str(collector_name), "Collector Update")
# Server functions
def statusMessageListener(self, message: str, showPopup: bool = True) -> None:
if showPopup:
@@ -1189,15 +1222,7 @@ class MainWindow(QtWidgets.QMainWindow):
status = "Getting Collectors..."
self.MESSAGEHANDLER.info(status)
self.setStatus(status)
try:
clientCollectorUIDs = literal_eval(self.SETTINGS.value('Project/Server/Collectors'))
except Exception as e:
self.MESSAGEHANDLER.error('Unable to load Collectors from Settings file.',
popUp=False,
exc_info=False)
self.MESSAGEHANDLER.debug('Cannot eval Project/Server/Collectors setting: ' + str(e))
clientCollectorUIDs = {}
self.FCOM.askServerForCollectors(clientCollectorUIDs)
self.FCOM.askServerForCollectors(self.getClientCollectors())
status = "Getting server projects list..."
self.MESSAGEHANDLER.info(status)
self.setStatus(status)
@@ -1249,6 +1274,28 @@ class MainWindow(QtWidgets.QMainWindow):
self.collectors = server_collectors
self.runningCollectors = continuing_collectors_info
def startNewCollectorListener(self, collector_category: str, collector_name: str, collector_uid: str,
collector_entities: list, collector_parameters: dict):
with self.serverCollectorsLock:
if collector_category not in self.runningCollectors:
self.runningCollectors[collector_category] = {}
if collector_name not in self.runningCollectors[collector_category]:
self.runningCollectors[collector_category][collector_name] = []
# Re-running a collector would generate duplicate info; we don't want that.
duplicateExists = False
for collectorInstance in self.runningCollectors[collector_category][collector_name]:
if collectorInstance['uid'] == collector_uid:
duplicateExists = True
break
if not duplicateExists:
self.runningCollectors[collector_category][collector_name].append({'uid': collector_uid,
'entities': collector_entities,
'parameters': collector_parameters})
currentCollectors = self.getClientCollectors()
currentCollectors[collector_uid] = time.time_ns() // 1000
self.setClientCollectors(currentCollectors)
def receiveProjectsListListener(self, projects: list) -> None:
with self.serverProjectsLock:
self.serverProjects = projects
@@ -2565,7 +2612,7 @@ class ResolutionParametersSelector(QtWidgets.QDialog):
msgBox = QtWidgets.QMessageBox()
msgBox.setModal(True)
QtWidgets.QMessageBox.warning(msgBox,
"Not all parameters filled in",
"Not all parameters were filled in",
"Some of the required parameters for the resolution have been left blank."
" Please fill them in before proceeding.")
return