class MainGUI(QMainWindow):
def __init__(self, logman, comment_mgr, val):
logging.debug("MainGUI(): Instantiated")
super(MainGUI, self).__init__()
self.setWindowTitle('Traffic Annotation Workflow')
self.setFixedSize(670,565)
self.logman = logman
self.comment_mgr = comment_mgr
self.val = val
self.project_sessions = ProjectSessions()
self.cm = ConfigurationManager.get_instance()
#shared data between widgets
self.existingconfignames = {}
self.logEnabled = ''
self.closeConfirmed = ''
self.newProject_pressed = False
self.newPro = None
#get project folder
self.project_data_folder = self.cm.read_config_value("PROJECTS", "PROJECTS_BASE_PATH")
self.createRequiredSubDirectories()
self.at_start = True
self.mainWidget = QWidget()
self.setCentralWidget(self.mainWidget)
mainlayout = QVBoxLayout()
self.baseWidget = QWidget() #BaseWidget()
self.annotateWidget = QWidget()
self.resultsWidget = QWidget()
self.projectTree = QtWidgets.QTreeWidget()
self.baseWidgets = {}
self.blankTreeContextMenu = {}
quit = QAction("Quit", self)
quit.triggered.connect(self.closeEvent)
#Add tab widget - RES
tabWidget = QtWidgets.QTabWidget()
tabWidget.setGeometry(QtCore.QRect(0, 15, 668, 565))
tabWidget.setObjectName("tabWidget")
#BaseWidget
self.baseWidget.setWindowTitle("BaseWidget")
self.baseWidget.setObjectName("BaseWidget")
baseLayoutWidget = QtWidgets.QWidget()
baseLayoutWidget.setObjectName("layoutWidget")
self.baseOuterVertBox = QtWidgets.QVBoxLayout()
self.baseOuterVertBox.setObjectName("outerVertBox")
baseLayoutWidget.setLayout(self.baseOuterVertBox)
self.baseWidget.setLayout(self.baseOuterVertBox)
#Configuration window - RES
## windowBoxHLayout contains:
###projectTree (Left)
###basedataStackedWidget (Right)
windowWidget = QtWidgets.QWidget()
windowWidget.setObjectName("windowWidget")
windowBoxHLayout = QtWidgets.QHBoxLayout()
windowBoxHLayout.setObjectName("windowBoxHLayout")
windowWidget.setLayout(windowBoxHLayout)
self.projectTree.itemSelectionChanged.connect(self.onItemSelected)
self.projectTree.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
self.projectTree.customContextMenuRequested.connect(self.showContextMenu)
self.projectTree.setEnabled(True)
self.projectTree.setMaximumSize(200,521)
self.projectTree.setObjectName("projectTree")
self.projectTree.headerItem().setText(0, "Projects")
self.projectTree.setSortingEnabled(False)
windowBoxHLayout.addWidget(self.projectTree)
self.basedataStackedWidget = QStackedWidget()
self.basedataStackedWidget.setObjectName("basedataStackedWidget")
windowBoxHLayout.addWidget(self.basedataStackedWidget)
tabWidget.addTab(windowWidget, "Configuration")
#ADD TAB WIDGET - RES
self.initMenu()
mainlayout = QVBoxLayout()
mainlayout.addWidget(self.mainMenu)
mainlayout.addWidget(tabWidget)
self.mainWidget.setLayout(mainlayout)
#load any saved projects
self.load_saved()
self.load_sessions()
self.at_start = False
logging.debug("MainWindow(): Complete")
def createRequiredSubDirectories(self):
logging.debug("MainApp:createRequiredSubDirectories() instantiated")
if os.path.exists(self.project_data_folder) == False:
try:
os.makedirs(self.project_data_folder)
except:
logging.error("MainApp:createRequiredSubDirectories(): An error occured when trying to create project directories: ")
exc_type, exc_value, exc_traceback = sys.exc_info()
traceback.print_exception(exc_type, exc_value, exc_traceback)
QMessageBox.error(self,
"Create Error",
"Could not create project subdirectories, quitting...",
QMessageBox.Ok)
exit()
#RES Method
def onItemSelected(self):
logging.debug("MainApp:onItemSelected instantiated")
# Get the selected item
self.selectedItem = self.projectTree.currentItem()
if self.selectedItem == None:
logging.debug("MainApp:onItemSelected no configurations left")
self.statusBar.showMessage("No configuration items selected or available.")
return
#Check if it's the case that an project name was selected
parentSelectedItem = self.selectedItem.parent()
if(parentSelectedItem == None):
#A base widget was selected
#logging.debug"PROJECT_WIDGET: " + str((self.baseWidgets[self.selectedItem.text(0)]["ProjectWidget"])))
self.basedataStackedWidget.setCurrentWidget(self.baseWidgets[self.selectedItem.text(0)]["ProjectWidget"])
else:
#for children
parentOfParent = parentSelectedItem.parent()
#Check if it's the case that a Session Name was selected
if(self.selectedItem.text(0)[0] == "S"):
#logging.debug"SESSION_WIDGET: " + str(self.baseWidgets[parentSelectedItem.text(0)][self.selectedItem.text(0)]["SessionWidget"]))
logging.debug("Setting right widget: " + str(self.baseWidgets[parentSelectedItem.text(0)][self.selectedItem.text(0)]["SessionWidget"]))
self.basedataStackedWidget.setCurrentWidget(self.baseWidgets[parentSelectedItem.text(0)][self.selectedItem.text(0)]["SessionWidget"])
#Check if it's the case that a Annotate was selected
elif(self.selectedItem.text(0)[0] == "A"):
#logging.debug"ANNOTATE " + str(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["AnnotateWidget"]))
logging.debug("Setting right widget: " + str(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["AnnotateWidget"]))
self.basedataStackedWidget.setCurrentWidget(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["AnnotateWidget"])
#Check if it's the case that a Rules was selected
elif(self.selectedItem.text(0)[0] == "R"):
#logging.debug"RULES " + str(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["RulesWidget"]))
logging.debug("Setting right widget: " + str(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["RulesWidget"]))
self.basedataStackedWidget.setCurrentWidget(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["RulesWidget"])
#Check if it's the case that a Results was selected
elif(self.selectedItem.text(0)[0] == "X"):
#logging.debug"RESULTS " + str(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["ResultsWidget"]))
logging.debug("Setting right widget: " + str(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["ResultsWidget"]))
self.basedataStackedWidget.setCurrentWidget(self.baseWidgets[parentOfParent.text(0)][parentSelectedItem.text(0)]["ResultsWidget"])
def on_add_curation_clicked(self):
logging.debug("on_add_curation_clicked(): Instantiated")
selectedItem = self.projectTree.currentItem()
self.nsd = NewSessionDialog(selectedItem.text(0))
projectSessions = []
for pot_session in self.baseWidgets[selectedItem.text(0)]:
if pot_session.startswith("S:"):
projectSessions.append(pot_session)
if len(projectSessions) > 0:
self.nsd.templateNameComboBox.addItems(projectSessions)
self.nsd.templateNameComboBox.setEnabled(True)
self.nsd.created.connect(self.session_created)
self.nsd.setWindowModality(QtCore.Qt.ApplicationModal)
self.nsd.show()
def addSession(self, projectItem, sessionName):
logging.debug("addSession(): Instantiated")
projectName = projectItem.text(0)
add_session_result = self.add_session_list(projectName, sessionName)
if add_session_result == False:
logging.error("An error occured when trying to add session to GUI: " + str(sessionName))
return
else:
sessionLabel = "S: " + sessionName
#create tree widget item
sessionItem = QtWidgets.QTreeWidgetItem(projectItem)
sessionItem.setText(0,sessionLabel)
sessionWidget = SessionWidget(sessionName)
self.baseWidgets[projectName][sessionLabel] = {} #project name (Parent of Parent) + session name (parent of children)
self.baseWidgets[projectName][sessionLabel]["SessionWidget"] = sessionWidget
self.basedataStackedWidget.addWidget(sessionWidget)
#create other widget items
##ANNOTATE
annItem = QtWidgets.QTreeWidgetItem()
annLabel = "A: " + "Annotate"
annItem.setText(0, annLabel)
sessionItem.addChild(annItem)
self.annotateWidget = AnnotateWidget(self.project_data_folder, projectName, os.path.basename(self.existingconfignames[projectName]), sessionName, self.comment_mgr) #send project name for the corresponding directory
self.baseWidgets[projectName][sessionLabel]["AnnotateWidget"] = self.annotateWidget #child
self.basedataStackedWidget.addWidget(self.annotateWidget)
##RULES
rulesItem = QtWidgets.QTreeWidgetItem()
rulesLabel = "R: " + "Rules"
rulesItem.setText(0, rulesLabel)
sessionItem.addChild(rulesItem)
#add the corresponding directory -- if it is already created, skip
rulesDir = os.path.join(self.project_data_folder, projectName)
rulesDir = os.path.join(rulesDir, ConfigurationManager.STRUCTURE_RULES_GEN_PATH)
if os.path.exists(rulesDir) == False:
os.mkdir(rulesDir)
rulesWidget = RulesWidget(self.project_data_folder, projectName, os.path.basename(self.existingconfignames[projectName]), sessionName, rulesDir, self.comment_mgr, self.val)
self.baseWidgets[projectName][sessionLabel]["RulesWidget"] = rulesWidget
self.basedataStackedWidget.addWidget(rulesWidget)
##RESULTS
resultsItem = QtWidgets.QTreeWidgetItem()
resultsLabel = "X: " + "Results"
resultsItem.setText(0, resultsLabel)
sessionItem.addChild(resultsItem)
#add the corresponding directory -- if it is already created, skip
resultsDir = os.path.join(self.project_data_folder, projectName)
resultsDir = os.path.join(resultsDir, "IDS-ALERTS")
if os.path.exists(resultsDir) == False:
os.mkdir(resultsDir)
self.resultsWidget = ResultsWidget(self.project_data_folder, projectName, sessionName, resultsDir, self.val)
self.baseWidgets[projectName][sessionLabel]["ResultsWidget"] = self.resultsWidget
self.basedataStackedWidget.addWidget(self.resultsWidget)
logging.debug("on_add_curation_clicked(): Completed")
def on_export_clicked(self):
#get project dir
selectedItem = self.projectTree.currentItem()
selectedItemName = selectedItem.text(0)
project_folder = os.path.join(self.project_data_folder, selectedItemName)
project_folder = os.path.abspath(project_folder)
self.exportPro = ExportDialog(self, selectedItemName, project_folder).exec_()
#self.exportPro.setWindowModality(QtCore.Qt.ApplicationModal)
#self.exportPro.show()
#RES METHOD
def showContextMenu(self, position):
logging.debug("MainApp:showContextMenu() instantiated: " + str(position))
if(self.projectTree.itemAt(position) == None):
self.blankTreeContextMenu.popup(self.projectTree.mapToGlobal(position))
elif(self.projectTree.itemAt(position).parent() == None):
self.projectContextMenu.popup(self.projectTree.mapToGlobal(position))
#RES METHOD
def initMenu(self):
logging.debug("MainApp:initMenu() instantiated")
self.mainMenu = QMenuBar()
self.fileMenu = self.mainMenu.addMenu("File")
self.blankTreeContextMenu = QtWidgets.QMenu()
self.newProjectContextMenu = QtWidgets.QMenu("New Project")
self.blankTreeContextMenu.addMenu(self.newProjectContextMenu)
self.fromCapture = self.newProjectContextMenu.addAction("Create Capture")
self.fromCapture.triggered.connect(self.newFromCapture)
self.fromPCAP = self.newProjectContextMenu.addAction("From PCAP")
self.fromPCAP.triggered.connect(self.newFromPCAP)
# Experiment context menu
self.importContextMenu = QtWidgets.QMenu("Import Project")
self.blankTreeContextMenu.addMenu(self.importContextMenu)
self.fromFolderContextSubMenu = self.importContextMenu.addAction("From Folder")
self.fromFolderContextSubMenu.triggered.connect(self.importFromFolder)
self.fromZipContextSubMenu = self.importContextMenu.addAction("From Zip")
self.fromZipContextSubMenu.triggered.connect(self.importFromZip)
#Context menu project
self.projectContextMenu = QtWidgets.QMenu()
self.addCuration = self.projectContextMenu.addAction("Create curation session")
self.addCuration.triggered.connect(self.on_add_curation_clicked)
self.exportProject = self.projectContextMenu.addAction("Export project")
self.exportProject.triggered.connect(self.on_export_clicked)
self.quitAppMenuButton = QAction(QIcon(), "Quit", self)
self.quitAppMenuButton.setShortcut("Ctrl+Q")
self.quitAppMenuButton.setStatusTip("Quit App")
self.quitAppMenuButton.triggered.connect(self.closeEvent)
self.fileMenu.addAction(self.quitAppMenuButton)
#Used to create a new project, this is where the prompt to write a name for the project is taken.
def newFromCapture(self):
logging.debug("MainApp:newFromCapture() instantiated")
#Creating a custom widget to display what is needed for creating a new project:
self.newPro = NewFromCollectDataDialog(self.logman, self.existingconfignames)
#slots to receive data from the custom widget
self.newPro.logEnabled.connect(self.log_enabled)
self.newPro.created.connect(self.project_created)
self.newPro.closeConfirmed.connect(self.close_confirmed)
self.newProject_pressed = True
self.newPro.setWindowModality(QtCore.Qt.ApplicationModal)
self.newPro.show()
def newFromPCAP(self):
logging.debug("MainApp:newFromPCAP() instantiated")
#Creating a custom widget to display what is needed for creating a new project:
self.newCap = NewFromPCAPDialog(self.existingconfignames)
#slot to receive data from the custom widget
self.newCap.created.connect(self.project_created)
self.newProject_pressed = True
self.newCap.setWindowModality(QtCore.Qt.ApplicationModal)
self.newCap.show()
#Slot for when the user created the new project, path and configname
@QtCore.pyqtSlot(str, dict, str)
def project_created(self, configname, existingconfignames, path):
#update project info with new info selected from widget
self.existingconfignames = existingconfignames
#create the new project with the updated information
self.addProject(configname, self.existingconfignames[configname], path)
@QtCore.pyqtSlot(str)
def session_created(self, newsessionname):
#call the function to add session to gui
selectedProject = self.projectTree.currentItem()
self.addSession(selectedProject, newsessionname)
#Slot to let us know if the logging has started
@QtCore.pyqtSlot(str)
def log_enabled(self, status):
self.logEnabled = status
#Slot to let us know if the close has been confirmed or canceled
@QtCore.pyqtSlot(str)
def close_confirmed(self, status):
self.closeConfirmed = status
#Used to create a new project, and this is where the project will actually be populated
def addProject(self, configname, projectPCAP, path):
if configname in self.baseWidgets:
logging.debug("Project already in tree: " + str(configname + "; skipping"))
return
#create the folders and files for new project if it's not already there:
self.projectWidget = ProjectWidget(configname, projectPCAP, path)
configTreeWidgetItem = QtWidgets.QTreeWidgetItem(self.projectTree)
configTreeWidgetItem.setText(0,configname)
self.projectWidget.addProjectItem(configname)
#Add base info
self.baseWidgets[configname] = {"BaseWidget": {}, "ProjectWidget": {}, "SessionWidget": {}, "AnnotateWidget": {}, "RulesWidget": {}, "ResultsWidget": {}}
self.baseWidgets[configname]["BaseWidget"] = self.baseWidget
self.basedataStackedWidget.addWidget(self.baseWidget)
self.baseWidgets[configname]["ProjectWidget"] = self.projectWidget
self.basedataStackedWidget.addWidget(self.projectWidget)
self.basedataStackedWidget.addWidget(self.baseWidget)
#add to list
self.project_sessions.add_project(configname)
def importFromZip(self):
logging.debug("MainApp:importFromZip() instantiated")
zip_file = QFileDialog()
zip_file.setWindowTitle("Select File")
zip_file.setFileMode(QtWidgets.QFileDialog.ExistingFile)
zip_file.setNameFilter("Zip Files (*.zip)")
filenames = zip_file.getOpenFileName()
filename = filenames[0]
if filename == "":
logging.debug("File choose cancelled")
return
else:
configname = os.path.basename(filename)
configname = os.path.splitext(configname)[0]
if configname in self.existingconfignames:
QMessageBox.warning(self,
"Name Exists",
"A project with the same name already exists.",
QMessageBox.Ok)
return None
else:
#instance of package manage
pack_mgr = PackageManager()
self.populate_import(pack_mgr, configname, os.path.abspath(filename))
def importFromFolder(self, configname):
logging.debug("MainApp:importFromFolder() instantiated")
folder_chosen = str(QFileDialog.getExistingDirectory(self, "Select Directory to Store Data"))
if folder_chosen == "":
logging.debug("File choose cancelled")
return
if len(folder_chosen) > 0:
baseNoExt = os.path.basename(folder_chosen)
baseNoExt = os.path.splitext(baseNoExt)[0]
configname = ''.join(e for e in baseNoExt if e.isalnum)
if configname in self.existingconfignames:
QMessageBox.warning(self,
"Name Exists",
"A project with the same name already exists.",
QMessageBox.Ok)
return None
else:
self.populate_import("dir", configname, folder_chosen)
def populate_import(self, function, configname, from_file):
importedProjectPath = os.path.join(self.project_data_folder, configname)
#copy selected dir to new dir
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.completion_signal.connect(self.copy_dir_complete)
if function == "dir":
self.batch_thread.add_function(self.copy_dir, from_file, importedProjectPath)
else:
self.batch_thread.add_function(function.unzip, from_file, configname, self.project_data_folder)
self.progress_dialog_overall = ProgressBarDialog(self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
def load_project_widgets(self):
logging.debug("load_project_widgets(): instantiated")
for name in self.existingconfignames:
path = os.path.join(self.project_data_folder, name)
projectPCAP = os.path.join(path, self.existingconfignames[name])
configname = name
self.addProject(configname, projectPCAP, path)
logging.debug("load_project_widgets(): Complete")
def copy_dir(self, from_dir, to_dir):
logging.debug("copy_dir(): copying selected directory")
copy_tree(from_dir, to_dir)
logging.debug("copy_dir(): copying complete")
def copy_dir_complete(self):
logging.debug("copy_dir_complete(): Instantiated")
self.progress_dialog_overall.update_progress()
self.progress_dialog_overall.hide()
#Need to load projects (add to existing...)
##GET Project PCAP Name Here and also add to existingprojectnames
self.load_saved()
self.load_sessions()
logging.debug("copy_dir_complete(): Complete")
def load_saved(self):
i = 0
#for each subdir, import the saved projects
for (dirName, subdirlist, filelist) in os.walk(self.project_data_folder):
for projectFolder in subdirlist:
pcapSubdir = os.path.join(self.project_data_folder, projectFolder, ConfigurationManager.STRUCTURE_PCAP_SUBDIR)
abspcapSubdir = os.path.abspath(pcapSubdir)
if os.path.exists(abspcapSubdir):
#go one level down
#if this is the pcap directory, then check if a pcap exists
#if so, this is a good project; create an entry in existingconfignames
for (subdirName, subsubdirlist, subfilelist) in os.walk(abspcapSubdir):
if len(subfilelist) == 1:
filename = subfilelist[0]
spt = os.path.splitext(filename)
if spt[1] == ".pcap" or spt[1] == ".pcapng":
self.existingconfignames[projectFolder] = filename
break
#once everything has been added, populate widget
self.load_project_widgets()
def load_sessions(self):
for name in self.existingconfignames:
#for already saved project
project_path = os.path.join(self.project_data_folder, name)
project_pcap_session = os.path.join(project_path, ConfigurationManager.STRUCTURE_PCAP_SUBDIR)
if os.path.exists(project_pcap_session):
paths, dirs, files = next(os.walk(project_pcap_session))
if len(dirs) > 0:
self.traverse_sessions(name, project_pcap_session)
def traverse_sessions(self, project_name, path):
#if RULES dir exists in project folder, then sessions exists
i = 0
for (dirName, subdirlist, filelist) in os.walk(path):
for sessionName in subdirlist:
if os.path.isfile(sessionName):
#skip
break
elif self.add_session_list(project_name, sessionName) == True:
self.add_session_widgets(project_name, self.existingconfignames[project_name], sessionName)
def add_session_widgets(self, project_name, project_pcap_filename, sessionName):
sessionLabel = "S: " + sessionName
#create tree widget item
selectedItem = self.projectTree.findItems(project_name, Qt.Qt.MatchContains)
sessionItem = QtWidgets.QTreeWidgetItem(selectedItem[0])
sessionItem.setText(0,sessionLabel)
sessionWidget = SessionWidget(sessionName)
self.baseWidgets[project_name][sessionLabel] = {} #project name (Parent of Parent) + session name (parent of children)
self.baseWidgets[project_name][sessionLabel]["SessionWidget"] = sessionWidget
self.basedataStackedWidget.addWidget(sessionWidget)
#create other widget items
##ANNOTATE
annItem = QtWidgets.QTreeWidgetItem()
annLabel = "A: " + "Annotate"
annItem.setText(0, annLabel)
sessionItem.addChild(annItem)
self.annotateWidget = AnnotateWidget(self.project_data_folder, project_name, project_pcap_filename, sessionName, self.comment_mgr) #send project name for the corresponding directory
self.baseWidgets[project_name][sessionLabel]["AnnotateWidget"] = self.annotateWidget #child
self.basedataStackedWidget.addWidget(self.annotateWidget)
##RULES
rulesItem = QtWidgets.QTreeWidgetItem()
rulesLabel = "R: " + "Rules"
rulesItem.setText(0, rulesLabel)
sessionItem.addChild(rulesItem)
#add the corresponding directory -- if it is already created, skip
rulesDir = os.path.join(self.project_data_folder, project_name)
rulesDir = os.path.join(rulesDir, ConfigurationManager.STRUCTURE_RULES_GEN_PATH)
if os.path.exists(rulesDir) == False:
os.mkdir(rulesDir)
self.rulesWidget = RulesWidget(self.project_data_folder, project_name, project_pcap_filename, sessionName, rulesDir, self.comment_mgr, self.val)
self.baseWidgets[project_name][sessionLabel]["RulesWidget"] = self.rulesWidget
self.basedataStackedWidget.addWidget(self.rulesWidget)
##RESULTS
resultsItem = QtWidgets.QTreeWidgetItem()
resultsLabel = "X: " + "Results"
resultsItem.setText(0, resultsLabel)
sessionItem.addChild(resultsItem)
#add the corresponding directory -- if it is already created, skip
resultsDir = os.path.join(self.project_data_folder, project_name)
resultsDir = os.path.join(resultsDir, "IDS-ALERTS")
if os.path.exists(resultsDir) == False:
os.mkdir(resultsDir)
self.resultsWidget = ResultsWidget(self.project_data_folder, project_name, sessionName, resultsDir, self.val)
self.baseWidgets[project_name][sessionLabel]["ResultsWidget"] = self.resultsWidget
self.basedataStackedWidget.addWidget(self.resultsWidget)
def add_session_list(self, project_name, project_session):
#method returns true if session was successfully added
success = self.project_sessions.add_project_session(project_name, project_session)
if success == False:
#self.project_sessions.print_d()
return False
else:
#self.project_sessions.print_d()
return True
def update_progress_bar(self):
logging.debug('update_progress_bar(): Instantiated')
self.progress_dialog_overall.update_progress()
logging.debug('update_progress_bar(): Complete')
def closeEvent(self, event):
logging.debug("closeEvent(): instantiated")
self.quit_event = event
if self.logEnabled == "TRUE":
#This means that the new project widget is still running so call the close event
#for that widget first to stop logger
self.newPro.closeEvent(event)
#Check if the close was confirmed or not
if self.closeConfirmed == "TRUE":
#after that's done, make sure to quit the app
self.quit_event.accept()
#self.close()
qApp.quit()
else:
return
else:
close = QMessageBox.question(self,
"QUIT",
"Are you sure you want to quit?",
QMessageBox.Yes | QMessageBox.No)
if close == QMessageBox.Yes:
#call the delete data function from new project, just to make sure
#everything has been cleared out
if self.newProject_pressed == True:
if self.newPro != None:
self.newPro.delete_data()
qApp.quit()
return
elif close == QMessageBox.No and not type(self.quit_event) == bool:
self.quit_event.ignore()
pass
return
def quit_app(self):
self.quit_event.accept()
qApp.quit()
return
class RulesWidget(QtWidgets.QWidget):
def __init__(self, projectfolder, projectName, projectPCAPName,
sessionLabel, rulesDir, comment_mgr, val):
QtWidgets.QWidget.__init__(self, parent=None)
self.rulesDir = rulesDir
self.comment_mgr = comment_mgr
self.val = val
self.outerVertBox = QtWidgets.QVBoxLayout()
self.outerVertBox.setObjectName("outerVertBoxAnnot")
self.setWindowTitle("RulesWidget")
self.setObjectName("RulesWidget")
#Label - Rules Title
self.labelVerBoxSess = QtWidgets.QVBoxLayout()
self.labelVerBoxSess.setObjectName("labeVerBoxPro")
self.rulesLabel = QtWidgets.QLabel("RULES")
labelFont = QtGui.QFont()
labelFont.setBold(True)
self.rulesLabel.setFont(labelFont)
self.rulesLabel.setAlignment(Qt.AlignCenter)
self.labelVerBoxSess.addWidget(self.rulesLabel)
#get project pcap path
projectpath = os.path.join(projectfolder, projectName)
projectPCAPFolder = os.path.join(
projectpath, ConfigurationManager.STRUCTURE_PCAP_SUBDIR)
projectpcap = os.path.join(projectPCAPFolder, projectPCAPName)
sessionPCAPFolder = os.path.join(projectPCAPFolder, sessionLabel)
self.sessionPCAP = os.path.join(sessionPCAPFolder,
"SessionGenerated.pcapng")
#create a directory for session rules
self.sessionRulesDir = os.path.join(rulesDir, sessionLabel)
#if corresponding session dir doesnt exist, create dir
if os.path.exists(self.sessionRulesDir) == False:
os.mkdir(self.sessionRulesDir)
#Project PCAP
self.pcapHorBox = QtWidgets.QHBoxLayout()
self.pcapHorBox.setObjectName("pcapHorBox")
self.pcapLabel = QtWidgets.QLabel()
self.pcapLabel.setObjectName("pcapLabel")
self.pcapLabel.setText("Project PCAP: ")
self.pcapHorBox.addWidget(self.pcapLabel)
self.pcapLineEdit = QtWidgets.QLineEdit()
self.pcapLineEdit.setAcceptDrops(False)
self.pcapLineEdit.setReadOnly(True)
self.pcapLineEdit.setObjectName("pcapLineEdit")
self.pcapLineEdit.setText(projectpcap)
self.pcapLineEdit.setAlignment(Qt.AlignLeft)
self.pcapHorBox.addWidget(self.pcapLineEdit)
self.pcapPathViewButton = QPushButton("View")
self.pcapPathViewButton.clicked.connect(
lambda x: self.on_view_button_clicked(x, projectPCAPFolder))
self.pcapHorBox.addWidget(self.pcapPathViewButton)
#Annotated PCAP
self.pcapHorBox2 = QtWidgets.QHBoxLayout()
self.pcapHorBox2.setObjectName("pcapHorBox2")
self.pcapLabel2 = QtWidgets.QLabel()
self.pcapLabel2.setObjectName("pcapLabel2")
self.pcapLabel2.setText("Annotated PCAP: ")
self.pcapHorBox2.addWidget(self.pcapLabel2)
self.pcapLineEdit2 = QtWidgets.QLineEdit()
self.pcapLineEdit2.setAcceptDrops(False)
self.pcapLineEdit2.setReadOnly(True)
self.pcapLineEdit2.setObjectName("pcapLineEdit2")
self.pcapLineEdit2.setText(self.sessionPCAP)
self.pcapLineEdit2.setAlignment(Qt.AlignLeft)
self.pcapHorBox2.addWidget(self.pcapLineEdit2)
self.pcapPathViewButton2 = QPushButton("View")
self.pcapPathViewButton2.clicked.connect(
lambda x: self.on_view_button_clicked(x, sessionPCAPFolder))
self.pcapHorBox2.addWidget(self.pcapPathViewButton2)
#Generate Rules Button
self.ruleButtonHorBox = QtWidgets.QHBoxLayout()
self.ruleButtonHorBox.setObjectName("ruleButtonHorBox")
self.rulesButton = QPushButton("Generate Rules")
self.rulesButton.clicked.connect(self.on_rules_button_clicked)
self.ruleButtonHorBox.setAlignment(Qt.AlignRight)
self.ruleButtonHorBox.addWidget(self.rulesButton)
#Path to where the rules are going to be outputted
self.ruleHorBox = QtWidgets.QHBoxLayout()
self.ruleHorBox.setObjectName("ruleHorBox")
self.rulePathLabel = QtWidgets.QLabel()
self.rulePathLabel.setObjectName("rulePathLabel")
self.rulePathLabel.setText("Rule Output Path: ")
self.ruleHorBox.addWidget(self.rulePathLabel)
self.ruleLineEdit = QtWidgets.QLineEdit()
self.ruleLineEdit.setAcceptDrops(False)
self.ruleLineEdit.setReadOnly(True)
self.ruleLineEdit.setObjectName("ruleLineEdit")
self.ruleLineEdit.setText(self.sessionRulesDir)
self.ruleLineEdit.setAlignment(Qt.AlignLeft)
self.ruleHorBox.addWidget(self.ruleLineEdit)
self.rulePathViewButton = QPushButton("View")
self.rulePathViewButton.clicked.connect(
lambda x: self.on_view_button_clicked(x, self.sessionRulesDir))
self.ruleHorBox.addWidget(self.rulePathViewButton)
#put everything together
self.outerVertBox.addLayout(self.labelVerBoxSess)
self.outerVertBox.addLayout(self.pcapHorBox)
self.outerVertBox.addLayout(self.pcapHorBox2)
self.outerVertBox.addLayout(self.ruleHorBox)
self.outerVertBox.addLayout(self.ruleButtonHorBox)
self.outerVertBox.addStretch()
self.setLayout(self.outerVertBox)
def on_view_button_clicked(self, x, folder_path=None):
if isinstance(folder_path, QTextEdit):
folder_path = folder_path.toPlainText()
self.file_explore_thread = FileExplorerRunner(
folder_location=folder_path)
self.file_explore_thread.start()
def on_rules_button_clicked(self):
logging.debug('on_genrules_out_start_button_clicked(): Instantiated')
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.completion_signal.connect(
self.genrules_button_batch_completed)
self.batch_thread.add_function(self.comment_mgr.extract_json,
self.sessionPCAP)
comment_filename = os.path.join(
self.sessionRulesDir, ConfigurationManager.STRUCTURE_JSON_COMMENTS)
self.batch_thread.add_function(
self.comment_mgr.write_comment_json_to_file, comment_filename)
self.batch_thread.add_function(self.val.extract_rules,
comment_filename)
rules_filename = os.path.join(
self.sessionRulesDir,
ConfigurationManager.STRUCTURE_RULES_GEN_FILE)
self.batch_thread.add_function(self.val.write_rules_to_file,
rules_filename)
self.progress_dialog_overall = ProgressBarDialog(
self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
logging.debug('on_genrules_out_start_button_clicked(): Complete')
def genrules_button_batch_completed(self):
logging.debug('thread_finish(): Instantiated')
self.progress_dialog_overall.update_progress()
self.progress_dialog_overall.hide()
logging.debug('thread_finish(): Completed')
def update_progress_bar(self):
logging.debug('update_progress_bar(): Instantiated')
self.progress_dialog_overall.update_progress()
logging.debug('update_progress_bar(): Complete')
class NewFromCollectDataDialog(QtWidgets.QWidget):
#Signal for when the user is done creating the new project
created = QtCore.pyqtSignal(str, dict, str)
logEnabled = QtCore.pyqtSignal(str)
closeConfirmed = QtCore.pyqtSignal(str)
def __init__(self, logman, existingProjects):
QtWidgets.QWidget.__init__(self, parent=None)
self.logger_started_once = False
self.existingconfignames = existingProjects
self.projectPath = ""
self.projectName = ""
self.cancel_pressed = False
self.saved_pressed = False
self.cm = ConfigurationManager.get_instance()
self.project_data_folder = self.cm.read_config_value(
"PROJECTS", "PROJECTS_BASE_PATH")
quit = QAction("Quit", self)
quit.triggered.connect(self.closeEvent)
#Title of window
self.outerVertBoxPro = QtWidgets.QVBoxLayout()
self.outerVertBoxPro.setObjectName("outerVertBox")
self.setWindowTitle("Collect Data")
self.setObjectName("NewFromCollectDataDialog")
#Label - New Project Title
self.labelVerBoxPro = QtWidgets.QVBoxLayout()
self.labelVerBoxPro.setObjectName("labeVerBoxPro")
self.newProjectLabel = QLabel("Collect Data for New Project")
labelFont = QtGui.QFont()
labelFont.setBold(True)
self.newProjectLabel.setFont(labelFont)
self.newProjectLabel.setAlignment(Qt.AlignCenter)
self.nameVerBoxPro = QtWidgets.QHBoxLayout()
self.nameVerBoxPro.setObjectName("nameVerBoxPro")
self.nameLabel = QtWidgets.QLabel()
self.nameLabel.setObjectName("nameLabel")
self.nameLabel.setText("Project Name:")
self.nameVerBoxPro.addWidget(self.nameLabel)
self.configname = QLineEdit()
self.configname.returnPressed.connect(self.on_log_start_button_clicked)
###### Fixed Height for project name text box
self.configname.setFixedHeight(27)
#Create buttons for creating new file
self.logOutStartButton = QPushButton("Start Logging")
self.logOutStopButton = QPushButton("Stop Logging")
self.logOutSaveButton = QPushButton("Save")
self.logOutCancelButton = QPushButton("Cancel")
#Add on click event
self.logOutStartButton.clicked.connect(
self.on_log_start_button_clicked)
self.logOutStartButton.setEnabled(True)
self.logOutStopButton.clicked.connect(self.on_log_stop_button_clicked)
self.logOutStopButton.setEnabled(False)
self.logOutSaveButton.clicked.connect(self.on_log_save_button_clicked)
self.logOutSaveButton.setEnabled(False)
self.logOutCancelButton.clicked.connect(self.on_cancel_button_clicked)
#Set the button layouts
self.bottomButtons_layout = QtWidgets.QHBoxLayout()
#Put all the components together
self.labelVerBoxPro.addWidget(self.newProjectLabel)
self.nameVerBoxPro.addWidget(self.configname)
self.bottomButtons_layout.addWidget(self.logOutStartButton)
self.bottomButtons_layout.addWidget(self.logOutStopButton)
self.bottomButtons_layout.addWidget(self.logOutSaveButton)
self.bottomButtons_layout.addWidget(self.logOutCancelButton,
alignment=QtCore.Qt.AlignRight)
self.outerVertBoxPro.addLayout(self.labelVerBoxPro)
self.outerVertBoxPro.addLayout(self.nameVerBoxPro)
self.outerVertBoxPro.addLayout(self.bottomButtons_layout)
#Auto Adjust Size
self.setFixedSize(self.labelVerBoxPro.sizeHint())
self.setFixedSize(self.nameVerBoxPro.sizeHint())
self.setFixedSize(self.bottomButtons_layout.sizeHint())
self.outerVertBoxPro.addStretch()
self.setLayout(self.outerVertBoxPro)
self.logman = logman
def on_log_start_button_clicked(self):
logging.debug('on_log_start_button_clicked(): Instantiated')
#Remove any special characters or spaces:
self.projectName = self.configname.text()
self.projectName = re.sub('\W+', '', self.projectName)
#check if name has been filed out in order to create a project folder
#with the name that was chosen:
if self.projectName != '':
self.projectPath = os.path.join(self.project_data_folder,
self.projectName)
#show the project name edited
self.configname.setText(self.projectName)
if self.logger_started_once == True and os.path.exists(
self.projectPath) == True:
buttonReply = QMessageBox.question(
self, 'Confirmation',
"Restarting the Logger will Remove any Previous Data. \r\n Continue?",
QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
shutil.rmtree(self.projectPath)
if buttonReply != QMessageBox.Yes:
logging.debug('on_log_start_button_clicked(): Cancelled')
return
if os.path.exists(self.projectPath) == True:
QMessageBox.warning(
self, "Name Exists",
"The project name specified and directory already exists",
QMessageBox.Ok)
return None
self.logger_started_once = True
self.logman.remove_data_all()
self.logman.start_collectors()
self.logEnabled.emit("TRUE")
self.logOutStartButton.setEnabled(False)
self.logOutStopButton.setEnabled(True)
self.configname.setReadOnly(True)
else:
QMessageBox.warning(self, "Name is Empty",
"Project Name is Empty!", QMessageBox.Ok)
logging.debug('on_log_start_button_clicked(): Complete')
def on_log_stop_button_clicked(self):
logging.debug('on_log_stop_button_clicked(): Instantiated')
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.completion_signal.connect(
self.stop_button_batch_completed)
self.batch_thread.add_function(self.logman.stop_collectors)
self.batch_thread.add_function(self.logman.parse_data_all)
self.batch_thread.add_function(self.logman.export_data,
os.path.abspath(self.projectPath))
parsedLogs = os.path.join(self.projectPath,
ConfigurationManager.STRUCTURE_PARSED_PATH)
latest_captured_PCAP = os.path.join(
self.projectPath,
ConfigurationManager.STRUCTURE_DEFAULT_RAW_PCAP_FILE)
click_out_path = os.path.join(
self.projectPath, ConfigurationManager.STRUCTURE_CLICKS_PATH)
timed_out_path = os.path.join(
self.projectPath, ConfigurationManager.STRUCTURE_TIMED_PATH)
self.batch_thread.add_function(self.logman.copy_latest_data,
self.projectPath, parsedLogs,
latest_captured_PCAP, click_out_path,
timed_out_path)
dissectorsPath = os.path.join(
self.projectPath,
ConfigurationManager.STRUCTURE_GEN_DISSECTORS_PATH)
self.batch_thread.add_function(self.logman.generate_dissectors,
parsedLogs, dissectorsPath, None)
self.progress_dialog_overall = ProgressBarDialog(
self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
logging.debug('on_log_stop_button_clicked(): Complete')
def update_progress_bar(self):
logging.debug('update_progress_bar(): Instantiated')
self.progress_dialog_overall.update_progress()
logging.debug('update_progress_bar(): Complete')
def stop_button_batch_completed(self):
logging.debug('thread_finish(): Instantiated')
self.progress_dialog_overall.update_progress()
self.progress_dialog_overall.hide()
output_dissected = "Saved Logs. \r\n\r\nCreated:\r\n"
for dissected in self.logman.get_generated_dissector_filenames():
output_dissected += str(os.path.basename(dissected)) + "\r\n"
if output_dissected == "":
QMessageBox.about(self, "Completed", "No files processed")
else:
QMessageBox.about(self, "Completed", output_dissected)
self.logOutStartButton.setEnabled(True)
self.logOutStopButton.setEnabled(False)
self.logOutSaveButton.setEnabled(True)
self.pcapFilename = os.path.join(
self.projectPath,
ConfigurationManager.STRUCTURE_DEFAULT_RAW_PCAP_FILE)
self.logEnabled.emit("FALSE")
logging.debug('thread_finish(): Completed')
def on_log_save_button_clicked(self):
logging.debug('on_log_save_button_clicked(): Instantiated')
if self.projectName != '':
if self.projectName in self.existingconfignames and os.path.exists(
self.projectPath) == True:
QMessageBox.warning(
self, "Name Exists",
"The project name specified and directory already exists",
QMessageBox.Ok)
return None
else:
#if all good, add to existing file names list
self.existingconfignames[self.projectName] = self.pcapFilename
self.saved_pressed = True
saveComplete = QMessageBox.warning(
self, "Creation Successful", "Success! Project Created.",
QMessageBox.Ok)
#Once save is hit, it should close the new project pop up and return to the main window
if saveComplete == QMessageBox.Ok:
#let main window know everything is ready:
#Send signal to slot
config = self.projectName
self.created.emit(config, self.existingconfignames,
self.projectPath)
self.close()
else:
QMessageBox.warning(self, "Name is Empty",
"Project Name is Empty!", QMessageBox.Ok)
logging.debug('on_log_save_button_clicked(): Complete')
def on_cancel_button_clicked(self, event):
logging.debug('on_cancel_button_clicked(): Instantiated')
self.quit_event = event
cancel = QMessageBox.question(
self, "Close New Project",
"Are you sure you want to quit? Any Collected data work will be discarded.",
QMessageBox.Close | QMessageBox.Cancel)
if cancel == QMessageBox.Close:
#call closing event
self.cancel_pressed = True #confrm that cancel was pressed
self.closeEvent(event)
elif cancel == QMessageBox.Cancel:
pass
logging.debug('on_cancel_button_clicked(): Complete')
def closeEvent(self, event):
logging.debug("closeEvent(): instantiated")
self.quit_event = event
if self.logOutStartButton.isEnabled() == True:
if self.cancel_pressed == True:
self.delete_data()
self.close()
elif self.configname.text() != '' and self.saved_pressed == False:
delete_temp = QMessageBox.question(
self, "Delete Temp Data",
"Closing... Discard data for this unsaved project?",
QMessageBox.Yes | QMessageBox.No)
if delete_temp == QMessageBox.Yes:
self.delete_data()
self.close()
#if none of these conditions are met, then close
self.close()
if self.logOutStartButton.isEnabled() == False:
logging.debug("closeEvent(): Creating Quit Command Load")
close = QMessageBox.question(self, "CLOSE",
"Logger is running. Stop and Close?",
QMessageBox.Yes | QMessageBox.No)
if close == QMessageBox.Yes:
self.closeConfirmed.emit("TRUE")
delete_temp = QMessageBox.question(
self, "Delete Temp Data",
"Closing... Discard data for this unsaved project?",
QMessageBox.Yes | QMessageBox.No)
if delete_temp == QMessageBox.Yes:
self.delete_data()
#run stop process:
self.stop_logger()
return
elif close == QMessageBox.No and not type(self.quit_event) == bool:
self.closeConfirmed.emit("FALSE")
self.quit_event.ignore()
self.closeConfirmed.emit("FALSE")
pass
logging.debug("closeEvent(): returning ignore")
return
def quit_app(self):
self.destroy()
self.progress_dialog_overall.hide()
return
def stop_logger(self):
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.add_function(self.logman.stop_collectors)
self.progress_dialog_overall = ProgressBarDialog(
self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
self.batch_thread.completion_signal.connect(self.quit_app)
self.logEnabled.emit("FALSE")
self.closeConfirmed.emit("TRUE")
return
def delete_data(self):
#make sure to not delete existing projects:
if self.projectName not in self.existingconfignames:
if self.logger_started_once == True:
#Delete Temp Data
self.logman.remove_data_all()
#If project directory has already been created, make sure to delete it
if os.path.exists(self.projectPath):
shutil.rmtree(self.projectPath)
class NewSessionDialog(QtWidgets.QWidget):
#Signal for when the user is done creating the new session
created = QtCore.pyqtSignal(str)
def __init__(self, projectName):
QtWidgets.QWidget.__init__(self, parent=None)
self.logger_started_once = False
self.projectPath = ""
self.projectName = projectName
self.cancel_pressed = False
self.saved_pressed = False
self.cm = ConfigurationManager.get_instance()
self.project_sessions_folder = self.cm.read_config_value(
"PROJECTS", "PROJECTS_BASE_PATH")
self.project_sessions_folder = os.path.join(
self.project_sessions_folder, self.projectName)
quit = QAction("Quit", self)
quit.triggered.connect(self.closeEvent)
#Title of window
self.outerVertBoxPro = QtWidgets.QVBoxLayout()
self.outerVertBoxPro.setObjectName("outerVertBox")
self.setWindowTitle("Create Session")
self.setObjectName("New Project Session")
#Label - New Session Title
self.labelVerBoxPro = QtWidgets.QVBoxLayout()
self.labelVerBoxPro.setObjectName("labeVerBoxPro")
self.newSessionLabel = QLabel("Create Session")
labelFont = QtGui.QFont()
labelFont.setBold(True)
self.newSessionLabel.setFont(labelFont)
self.newSessionLabel.setAlignment(Qt.AlignCenter)
self.sessionNameHorBoxPro = QtWidgets.QHBoxLayout()
self.sessionNameHorBoxPro.setObjectName("sessionNameHorBoxPro")
self.sessionNameLabel = QtWidgets.QLabel()
self.sessionNameLabel.setObjectName("sessionNameLabel")
self.sessionNameLabel.setText("Session Name:")
self.sessionNameHorBoxPro.addWidget(self.sessionNameLabel)
self.newsessionname = QLineEdit()
self.newsessionname.textChanged.connect(self.on_newsessionname_changed)
self.newsessionname.setFixedHeight(27)
self.sessionNameHorBoxPro.addWidget(self.newsessionname)
self.templateNameHorBoxPro = QtWidgets.QHBoxLayout()
self.templateNameHorBoxPro.setObjectName("templateNameHorBoxPro")
self.templateNameLabel = QtWidgets.QLabel()
self.templateNameLabel.setObjectName("templateNameLabel")
self.templateNameLabel.setText("From Session:")
self.templateNameHorBoxPro.addWidget(self.templateNameLabel)
self.templateNameComboBox = QComboBox()
self.templateNameComboBox.addItem("None")
self.templateNameComboBox.setEnabled(False)
self.templateNameComboBox.setFixedHeight(27)
self.templateNameHorBoxPro.addWidget(self.templateNameComboBox)
#Create buttons for OK/Cancel
self.okButton = QPushButton("OK")
self.okButton.setEnabled(False)
self.okButton.clicked.connect(self.on_ok_button_clicked)
self.cancelButton = QPushButton("Cancel")
self.cancelButton.clicked.connect(self.on_cancel_button_clicked)
#Set the button layouts
self.bottomButtons_layout = QtWidgets.QHBoxLayout()
#Put all the components together
self.labelVerBoxPro.addWidget(self.newSessionLabel)
self.bottomButtons_layout.addWidget(self.okButton)
self.bottomButtons_layout.addWidget(self.cancelButton,
alignment=QtCore.Qt.AlignRight)
self.outerVertBoxPro.addLayout(self.labelVerBoxPro)
self.outerVertBoxPro.addLayout(self.sessionNameHorBoxPro)
self.outerVertBoxPro.addLayout(self.templateNameHorBoxPro)
self.outerVertBoxPro.addLayout(self.bottomButtons_layout)
self.outerVertBoxPro.addStretch()
self.setLayout(self.outerVertBoxPro)
def on_newsessionname_changed(self):
#logging.debug('on_configname_changed(): Instantiated')
if self.newsessionname.text() != "":
self.okButton.setEnabled(True)
else:
self.okButton.setEnabled(False)
def on_ok_button_clicked(self):
logging.debug('on_ok_button_clicked(): Instantiated')
#Remove any special characters or spaces:
self.sessionName = self.newsessionname.text()
self.sessionName = re.sub('\W+', '', self.sessionName)
#check if name has been filed out in order to create a session folder
#with the name that was chosen:
if self.sessionName != '':
self.sessionPath = os.path.join(self.project_sessions_folder,
self.sessionName)
#show the project name edited
self.newsessionname.setText(self.sessionName)
if os.path.exists(self.sessionPath) == True:
QMessageBox.warning(
self, "Name Exists",
"The session name specified and directory already exists",
QMessageBox.Ok)
return None
else:
QMessageBox.warning(self, "Name is Empty",
"Session Name is Empty!", QMessageBox.Ok)
return None
#copy over the pcap (with a progress bar)
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.completion_signal.connect(
self.create_session_completed)
#TODO: check if template chosen and then add the function call
if self.templateNameComboBox.currentText() == "None":
self.batch_thread.add_function(self.create_session,
self.newsessionname.text())
else:
self.batch_thread.add_function(
self.create_session, self.newsessionname.text(),
self.templateNameComboBox.currentText())
#create subdirectories (empty) for this project
self.progress_dialog_overall = ProgressBarDialog(
self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
logging.debug('on_ok_button_clicked(): Complete')
def update_progress_bar(self):
logging.debug('update_progress_bar(): Instantiated')
self.progress_dialog_overall.update_progress()
logging.debug('update_progress_bar(): Complete')
def create_session_completed(self):
logging.debug('copy_completed(): Instantiated')
self.progress_dialog_overall.update_progress()
self.progress_dialog_overall.hide()
newsessionname = self.newsessionname.text()
QMessageBox.about(self, "Completed", "Finished Setting Up Session")
self.created.emit(newsessionname)
self.close()
logging.debug('copy_completed(): Completed')
def on_cancel_button_clicked(self, event):
logging.debug('on_cancel_button_clicked(): Instantiated')
self.quit_event = event
self.closeEvent(event)
logging.debug('on_cancel_button_clicked(): Complete')
def closeEvent(self, event):
logging.debug("closeEvent(): instantiated")
self.quit_event = event
self.progress_dialog_overall.close()
self.close()
return
def create_session(self, session_name, template_name=None):
logging.debug("create_session(): instantiated")
try:
if template_name != None:
#first copy required
logging.debug("Creating session " + str(session_name) +
" from template: " + str(template_name))
if template_name.startswith("S: "):
logging.debug(
"Removing Preceeding characters from template name: " +
str(template_name))
template_name = template_name.split("S: ", 1)[1]
logging.debug("New template name: " + str(template_name))
templatePathPCAP = os.path.join(
self.project_sessions_folder,
ConfigurationManager.STRUCTURE_PCAP_SUBDIR, template_name)
dstSessionPCAP = os.path.join(
self.project_sessions_folder,
ConfigurationManager.STRUCTURE_PCAP_SUBDIR, session_name)
if os.path.exists(templatePathPCAP):
logging.debug("Template PCAP Path exists: " +
str(templatePathPCAP))
if os.path.exists(dstSessionPCAP):
logging.debug(
"Session PCAP Path already exists, removing: " +
str(dstSessionPCAP))
shutil.rmtree(dstSessionPCAP)
logging.debug("Copying Template PCAP " +
str(templatePathPCAP) + " to New Session: " +
str(dstSessionPCAP))
res_dir = shutil.copytree(templatePathPCAP, dstSessionPCAP)
else:
logging.error(
"PCAP for Template Session does not exist: " +
str(template_name))
return
templatePathRULES = os.path.join(
self.project_sessions_folder,
ConfigurationManager.STRUCTURE_RULES_GEN_PATH,
template_name)
dstSessionPathRULES = os.path.join(
self.project_sessions_folder,
ConfigurationManager.STRUCTURE_RULES_GEN_PATH,
session_name)
if os.path.exists(templatePathRULES):
if os.path.exists(dstSessionPathRULES):
shutil.rmtree(dstSessionPathRULES)
res_dir = shutil.copytree(templatePathRULES,
dstSessionPathRULES)
else:
logging.debug(
"RULES PATH for Template Session does not exist: " +
str(template_name))
templatePathIDS_ALERTS = os.path.join(
self.project_sessions_folder,
ConfigurationManager.STRUCTURE_ALERT_GEN_PATH,
template_name)
dstPathIDS_ALERTS = os.path.join(
self.project_sessions_folder,
ConfigurationManager.STRUCTURE_ALERT_GEN_PATH,
session_name)
if os.path.exists(templatePathIDS_ALERTS):
if os.path.exists(dstPathIDS_ALERTS):
shutil.rmtree(dstPathIDS_ALERTS)
res_dir = shutil.copytree(templatePathIDS_ALERTS,
dstPathIDS_ALERTS)
else:
logging.debug(
"RULES PATH for Template Session does not exist: " +
str(template_name))
except:
logging.error(
"create_session(): An error occured when trying to copy session folders"
)
exc_type, exc_value, exc_traceback = sys.exc_info()
traceback.print_exception(exc_type, exc_value, exc_traceback)
def copy_Session(self, src_folder, dst_folder):
logging.debug('copy_Session(): Instantiated')
try:
dst_folder = os.path.dirname(dst_folder)
if os.path.exists(dst_folder) == False:
os.makedirs(dst_folder)
shutil.copy2(src_folder, dst_folder)
except:
logging.error(
"copy_Session(): An error occured when trying to copy log files"
)
exc_type, exc_value, exc_traceback = sys.exc_info()
traceback.print_exception(exc_type, exc_value, exc_traceback)
class ResultsWidget(QtWidgets.QWidget):
def __init__(self, projectfolder, projectName, sessionLabel, resultsDir, val):
QtWidgets.QWidget.__init__(self, parent=None)
self.resultsDir = resultsDir
self.val = val
self.outerVertBox = QtWidgets.QVBoxLayout()
self.outerVertBox.setObjectName("outerVertBoxAnnot")
self.setWindowTitle("ResultsWidget")
self.setObjectName("ResultsWidget")
#Label - Results Title
self.labelVerBoxSess = QtWidgets.QVBoxLayout()
self.labelVerBoxSess.setObjectName("labeVerBoxPro")
self.resultLabel = QtWidgets.QLabel("RESULTS")
labelFont = QtGui.QFont()
labelFont.setBold(True)
self.resultLabel.setFont(labelFont)
self.resultLabel.setAlignment(Qt.AlignCenter)
self.labelVerBoxSess.addWidget(self.resultLabel)
#get project-session rules path
projectpath = os.path.join(projectfolder, projectName)
projectPCAPFolder = os.path.join(projectpath, ConfigurationManager.STRUCTURE_PCAP_SUBDIR)
sessionPCAPFolder = os.path.join(projectPCAPFolder, sessionLabel)
self.sessionPCAP = os.path.join(sessionPCAPFolder, "SessionGenerated.pcapng")
self.suspectPCAP = os.path.join(sessionPCAPFolder, "SuspectPCAP.pcapng")
#Create dir for session alerts
self.sessionAlertsDir = os.path.join(resultsDir, sessionLabel)
try:
if os.path.exists(self.sessionAlertsDir):
shutil.rmtree(self.sessionAlertsDir, ignore_errors=True)
os.makedirs(self.sessionAlertsDir)
except Exception as e:
exc_type, exc_value, exc_traceback = sys.exc_info()
logging.error("comment_to_json(): An error occured ")
traceback.print_exception(exc_type, exc_value, exc_traceback)
#get session rules
self.sessionRulesDir = os.path.join(projectpath, ConfigurationManager.STRUCTURE_RULES_GEN_PATH)
self.sessionRulesDir = os.path.join(self.sessionRulesDir, sessionLabel)
self.rules_filename = os.path.join(self.sessionRulesDir, ConfigurationManager.STRUCTURE_RULES_GEN_FILE)
#Path to where the rules stored
self.ruleHorBox = QtWidgets.QHBoxLayout()
self.ruleHorBox.setObjectName("ruleHorBox")
self.rulePathLabel = QtWidgets.QLabel()
self.rulePathLabel.setObjectName("rulePathLabel")
self.rulePathLabel.setText("Rule Path: ")
self.ruleHorBox.addWidget(self.rulePathLabel)
self.ruleLineEdit = QtWidgets.QLineEdit()
self.ruleLineEdit.setAcceptDrops(False)
self.ruleLineEdit.setReadOnly(True)
self.ruleLineEdit.setObjectName("ruleLineEdit")
self.ruleLineEdit.setText(self.sessionRulesDir)
self.ruleLineEdit.setAlignment(Qt.AlignLeft)
self.ruleHorBox.addWidget(self.ruleLineEdit)
self.rulePathViewButton = QPushButton("View")
self.rulePathViewButton.clicked.connect(lambda x: self.on_view_button_clicked(x, self.sessionRulesDir))
self.ruleHorBox.addWidget(self.rulePathViewButton)
#Annotated/Suspect PCAP
self.pcapHorBox2 = QtWidgets.QHBoxLayout()
self.pcapHorBox2.setObjectName("pcapHorBox2")
self.pcapLabel2 = QtWidgets.QLabel()
self.pcapLabel2.setObjectName("pcapLabel2")
self.pcapLabel2.setText("Apply Rules to PCAP: ")
self.pcapHorBox2.addWidget(self.pcapLabel2)
self.pcapLineEdit2 = QtWidgets.QLineEdit()
self.pcapLineEdit2.setFixedWidth(150)
#self.pcapLineEdit2.setFixedHeight(25)
if os.path.exists(self.suspectPCAP):
self.pcapLineEdit2.setText(self.suspectPCAP)
self.pcapLineEdit2.setAcceptDrops(False)
self.pcapLineEdit2.setReadOnly(True)
self.pcapLineEdit2.setObjectName("pcapLineEdit2")
self.pcapLineEdit2.setAlignment(Qt.AlignLeft)
self.pcapHorBox2.addWidget(self.pcapLineEdit2)
#view and ... button
self.suspViewButton = QPushButton("View")
self.suspPathButton = QPushButton("...")
self.pcapHorBox2.addWidget(self.suspPathButton)
self.pcapHorBox2.addWidget(self.suspViewButton)
#on buttons clicked
self.suspViewButton.clicked.connect(lambda x: self.on_view_button_clicked(x, self.pcapLineEdit2))
self.suspPathButton.clicked.connect(self.on_path_button_clicked)
#Show path to the alert folder
self.alertHorBox = QtWidgets.QHBoxLayout()
self.alertHorBox.setObjectName("alertHorBox")
self.alertPathLabel = QtWidgets.QLabel()
self.alertPathLabel.setObjectName("alertPathLabel")
self.alertPathLabel.setText("Alert Output Path: ")
self.alertHorBox.addWidget(self.alertPathLabel)
self.alertLineEdit = QtWidgets.QLineEdit()
self.alertLineEdit.setAcceptDrops(False)
self.alertLineEdit.setReadOnly(True)
self.alertLineEdit.setObjectName("alertLineEdit")
self.alertLineEdit.setText(self.sessionAlertsDir)
self.alertLineEdit.setAlignment(Qt.AlignLeft)
self.alertHorBox.addWidget(self.alertLineEdit)
self.alertPathViewButton = QPushButton("View")
self.alertPathViewButton.clicked.connect(lambda x: self.on_view_button_clicked(x, self.sessionAlertsDir))
self.alertHorBox.addWidget(self.alertPathViewButton)
#Generate Alerts Button
self.alertButtonHorBox = QtWidgets.QHBoxLayout()
self.alertButtonHorBox.setObjectName("alertButtonHorBox")
self.alertButton = QPushButton("Generate Alerts")
self.alertButton.clicked.connect(self.on_alert_button_clicked)
self.alertButtonHorBox.setAlignment(Qt.AlignRight)
self.alertButtonHorBox.addWidget(self.alertButton)
if self.pcapLineEdit2 == "":
self.alertButton.setEnabled(False)
self.outerVertBox.addLayout(self.labelVerBoxSess)
self.outerVertBox.addLayout(self.ruleHorBox)
self.outerVertBox.addLayout(self.pcapHorBox2)
self.outerVertBox.addLayout(self.alertHorBox)
self.outerVertBox.addLayout(self.alertButtonHorBox)
self.outerVertBox.addStretch()
self.setLayout(self.outerVertBox)
def on_view_button_clicked(self, x, folder_path=None):
if isinstance(folder_path, QTextEdit):
folder_path = folder_path.toPlainText()
elif isinstance(folder_path, QtWidgets.QLineEdit):
folder_path = folder_path.text()
if folder_path == "":
QMessageBox.warning(self,
"No path selected",
"There is no path selected",
QMessageBox.Ok)
return None
self.file_explore_thread = FileExplorerRunner(folder_location=folder_path)
self.file_explore_thread.start()
def on_alert_button_clicked(self):
logging.debug('on_analyze_out_start_button_clicked(): Instantiated')
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.completion_signal.connect(self.analyze_button_batch_completed)
alertOutPath = os.path.join(self.sessionAlertsDir)
try:
if os.path.exists(alertOutPath):
shutil.rmtree(self.sessionAlertsDir, ignore_errors=True)
os.makedirs(alertOutPath)
except Exception as e:
exc_type, exc_value, exc_traceback = sys.exc_info()
logging.error("comment_to_json(): An error occured ")
traceback.print_exception(exc_type, exc_value, exc_traceback)
suricata_config_filename = ConfigurationManager.get_instance().read_config_abspath("VALIDATOR", "SURICATA_CONFIG_FILENAME")
self.batch_thread.add_function( self.val.run_suricata_with_rules, None, suricata_config_filename, alertOutPath, self.rules_filename, self.pcapLineEdit2.text())
logging.debug("SUSPECT PCAP: " + self.pcapLineEdit2.text())
self.progress_dialog_overall = ProgressBarDialog(self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
logging.debug('on_analyze_out_start_button_clicked(): Complete')
def analyze_button_batch_completed(self):
logging.debug('thread_finish(): Instantiated')
self.progress_dialog_overall.update_progress()
self.progress_dialog_overall.hide()
alertOutFile = os.path.join(self.sessionAlertsDir, ConfigurationManager.STRUCTURE_ALERT_GEN_FILE)
if os.path.exists(alertOutFile) == False or os.stat(alertOutFile).st_size < 10:
QMessageBox.about(self, "IDS Alerts", "No alerts generated")
else:
res = QMessageBox.question(self,
"Alerts written.\r\n",
"Open File?",
QMessageBox.Yes | QMessageBox.No)
if res == QMessageBox.Yes:
logging.debug("analyze_button_batch_completed(): Opening Alerts File")
self.on_view_button_clicked(None, alertOutFile)
def update_progress_bar(self):
logging.debug('update_progress_bar(): Instantiated')
self.progress_dialog_overall.update_progress()
logging.debug('update_progress_bar(): Complete')
def on_path_button_clicked(self):
logging.debug('on_path_button_clicked(): Instantiated')
choose_file = QFileDialog()
filenames, _ = QFileDialog.getOpenFileNames(choose_file, "Select Suspect File")
if len(filenames) < 0:
logging.debug("File choose cancelled")
return
if len(filenames) > 0:
suspect_pcap_chosen = filenames[0]
session_pcaps = os.path.dirname(os.path.abspath(self.sessionPCAP))
self.suspect_pcap_path = os.path.join(session_pcaps, "SuspectPCAP.pcapng")
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.completion_signal.connect(self.copy_suspect_complete)
self.batch_thread.add_function(self.copy_file, suspect_pcap_chosen, self.suspect_pcap_path)
#shutil.copy(suspect_pcap_chosen, suspect_pcap_path)
self.progress_dialog_overall = ProgressBarDialog(self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
logging.debug('on_path_button_clicked(): Complete')
def copy_file(self, file, dst):
shutil.copy(file, dst)
def copy_suspect_complete(self):
self.pcapLineEdit2.setText(self.suspect_pcap_path)
self.progress_dialog_overall.update_progress()
self.progress_dialog_overall.hide()
self.alertButton.setEnabled(True)
class NewFromPCAPDialog(QtWidgets.QWidget):
#Signal for when the user is done creating the new project
created = QtCore.pyqtSignal(str, dict, str)
def __init__(self, existingProjects):
QtWidgets.QWidget.__init__(self, parent=None)
self.logger_started_once = False
self.existingconfignames = existingProjects
self.projectPath = ""
self.projectName = ""
self.cancel_pressed = False
self.saved_pressed = False
self.cm = ConfigurationManager.get_instance()
self.project_data_folder = self.cm.read_config_value(
"PROJECTS", "PROJECTS_BASE_PATH")
quit = QAction("Quit", self)
quit.triggered.connect(self.closeEvent)
#Title of window
self.outerVertBoxPro = QtWidgets.QVBoxLayout()
self.outerVertBoxPro.setObjectName("outerVertBox")
self.setWindowTitle("Project from PCAP")
self.setObjectName("NewFromPCAP")
#Label - New Project Title
self.labelVerBoxPro = QtWidgets.QVBoxLayout()
self.labelVerBoxPro.setObjectName("labeVerBoxPro")
self.newProjectLabel = QLabel("Create Project from PCAP")
labelFont = QtGui.QFont()
labelFont.setBold(True)
self.newProjectLabel.setFont(labelFont)
self.newProjectLabel.setAlignment(Qt.AlignCenter)
self.projectNameHorBoxPro = QtWidgets.QHBoxLayout()
self.projectNameHorBoxPro.setObjectName("projectNameHorBoxPro")
self.projectNameLabel = QtWidgets.QLabel()
self.projectNameLabel.setObjectName("projectNameLabel")
self.projectNameLabel.setText("Project Name:")
self.projectNameHorBoxPro.addWidget(self.projectNameLabel)
self.configname = QLineEdit()
self.configname.textChanged.connect(self.on_configname_changed)
self.configname.setFixedHeight(27)
self.projectNameHorBoxPro.addWidget(self.configname)
self.pcapNameHorBoxPro = QtWidgets.QHBoxLayout()
self.pcapNameHorBoxPro.setObjectName("pcapNameHorBoxPro")
self.pcapNameLabel = QtWidgets.QLabel()
self.pcapNameLabel.setObjectName("pcapNameLabel")
self.pcapNameLabel.setText("Path to PCAP:")
self.pcapNameHorBoxPro.addWidget(self.pcapNameLabel)
self.pcapNameLineEdit = QLineEdit()
self.pcapNameLineEdit.setEnabled(False)
self.pcapNameLineEdit.setFixedHeight(27)
self.pcapNameHorBoxPro.addWidget(self.pcapNameLineEdit)
self.selectPCAPButton = QPushButton("...")
self.selectPCAPButton.clicked.connect(
self.on_select_PCAP_button_clicked)
self.pcapNameHorBoxPro.addWidget(self.selectPCAPButton)
#Create buttons for OK/Cancel
self.okButton = QPushButton("OK")
self.okButton.setEnabled(False)
self.okButton.clicked.connect(self.on_ok_button_clicked)
self.cancelButton = QPushButton("Cancel")
self.cancelButton.clicked.connect(self.on_cancel_button_clicked)
#Set the button layouts
self.bottomButtons_layout = QtWidgets.QHBoxLayout()
#Put all the components together
self.labelVerBoxPro.addWidget(self.newProjectLabel)
self.bottomButtons_layout.addWidget(self.okButton)
self.bottomButtons_layout.addWidget(self.cancelButton,
alignment=QtCore.Qt.AlignRight)
self.outerVertBoxPro.addLayout(self.labelVerBoxPro)
self.outerVertBoxPro.addLayout(self.projectNameHorBoxPro)
self.outerVertBoxPro.addLayout(self.pcapNameHorBoxPro)
self.outerVertBoxPro.addLayout(self.bottomButtons_layout)
self.outerVertBoxPro.addStretch()
self.setLayout(self.outerVertBoxPro)
def on_configname_changed(self):
#logging.debug('on_configname_changed(): Instantiated')
if self.pcapNameLineEdit.text() != "" and self.configname.text() != "":
self.okButton.setEnabled(True)
else:
self.okButton.setEnabled(False)
def on_select_PCAP_button_clicked(self):
logging.debug('select_pcap_button_clicked(): Instantiated')
pcap_file = QFileDialog()
pcap_file.setWindowTitle("Select File")
pcap_file.setFileMode(QtWidgets.QFileDialog.ExistingFile)
pcap_file.setNameFilter("Zip Files (*.zip)")
filenames = pcap_file.getOpenFileName()
filename = filenames[0]
if filename == "":
logging.debug("File choose cancelled")
return
else:
# Very similar to populate_import
self.pcapNameLineEdit.setText(filename)
if self.configname.text() != "":
self.okButton.setEnabled(True)
def on_ok_button_clicked(self):
logging.debug('on_ok_button_clicked(): Instantiated')
#Remove any special characters or spaces:
self.projectName = self.configname.text()
self.projectName = re.sub('\W+', '', self.projectName)
#check if name has been filed out in order to create a project folder
#with the name that was chosen:
if self.projectName != '':
self.projectPath = os.path.join(self.project_data_folder,
self.projectName)
#show the project name edited
self.configname.setText(self.projectName)
if os.path.exists(self.projectPath) == True:
QMessageBox.warning(
self, "Name Exists",
"The project name specified and directory already exists",
QMessageBox.Ok)
return None
else:
QMessageBox.warning(self, "Name is Empty",
"Project Name is Empty!", QMessageBox.Ok)
return None
#copy over the pcap (with a progress bar)
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.completion_signal.connect(self.copy_completed)
src_abs_filename = self.pcapNameLineEdit.text()
src_filename = os.path.basename(src_abs_filename)
dst_filename = os.path.join(self.projectPath,
ConfigurationManager.STRUCTURE_PCAP_SUBDIR,
str(src_filename))
self.batch_thread.add_function(self.copy_PCAP, src_abs_filename,
os.path.abspath(dst_filename))
#create subdirectories (empty) for this project
self.progress_dialog_overall = ProgressBarDialog(
self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
logging.debug('on_ok_button_clicked(): Complete')
def update_progress_bar(self):
logging.debug('update_progress_bar(): Instantiated')
self.progress_dialog_overall.update_progress()
logging.debug('update_progress_bar(): Complete')
def copy_completed(self):
logging.debug('copy_completed(): Instantiated')
self.progress_dialog_overall.update_progress()
self.progress_dialog_overall.hide()
config = self.configname.text()
self.existingconfignames[config] = self.pcapNameLineEdit.text()
QMessageBox.about(self, "Completed", "Finished Setting Up Project")
self.created.emit(config, self.existingconfignames, self.projectPath)
self.close()
logging.debug('copy_completed(): Completed')
def on_cancel_button_clicked(self, event):
logging.debug('on_cancel_button_clicked(): Instantiated')
self.quit_event = event
self.closeEvent(event)
logging.debug('on_cancel_button_clicked(): Complete')
def closeEvent(self, event):
logging.debug("closeEvent(): instantiated")
self.quit_event = event
self.close()
return
def copy_PCAP(self, src_filename, dst_filename):
logging.debug('copy_PCAP(): Instantiated')
try:
dst_folder = os.path.dirname(dst_filename)
if os.path.exists(dst_folder) == False:
os.makedirs(dst_folder)
shutil.copy2(src_filename, dst_folder)
except:
logging.error(
"copy_PCAP(): An error occured when trying to copy log files")
exc_type, exc_value, exc_traceback = sys.exc_info()
traceback.print_exception(exc_type, exc_value, exc_traceback)
class ExportDialog(QtWidgets.QDialog):
def __init__(self, parent, project_path, project_data_path):
logging.debug("ExportDialog(): instantiated")
super(ExportDialog, self).__init__(parent)
self.parent = parent
# quit = QAction("Quit", self)
# quit.triggered.connect(self.closeEvent)
# self.cancel_pressed = False
self.project_path = project_path
self.project_data_path = project_data_path
#Title of window
self.outerVertBoxPro = QtWidgets.QVBoxLayout()
self.outerVertBoxPro.setObjectName("outerVertBox")
self.setWindowTitle("Export Project")
self.setObjectName("ExportProjectDialog")
#Label - New Project Title
self.labelVerBoxPro = QtWidgets.QVBoxLayout()
self.labelVerBoxPro.setObjectName("labeVerBoxPro")
self.newProjectLabel = QLabel("Exporting Project Settings")
labelFont = QtGui.QFont()
labelFont.setBold(True)
self.newProjectLabel.setFont(labelFont)
self.newProjectLabel.setAlignment(Qt.AlignCenter)
self.labelVerBoxPro.addWidget(self.newProjectLabel)
self.nameHorBox = QtWidgets.QHBoxLayout()
self.nameHorBox.setObjectName("nameVerBoxPro")
self.nameLabel = QtWidgets.QLabel()
self.nameLabel.setObjectName("nameLabel")
self.nameLabel.setText("Output Path:")
self.nameHorBox.addWidget(self.nameLabel)
self.exportOutputPath = QtWidgets.QLineEdit()
self.exportOutputPath.setFixedWidth(200)
self.exportOutputPath.setAcceptDrops(False)
self.exportOutputPath.setReadOnly(True)
self.exportOutputPath.setObjectName("exportOutputPath")
self.nameHorBox.addWidget(self.exportOutputPath)
self.exportPathButton = QPushButton("...")
self.exportPathButton.clicked.connect(self.on_path_button_clicked)
self.nameHorBox.addWidget(self.exportPathButton)
self.exportPathViewButton = QPushButton("View")
self.exportPathViewButton.clicked.connect(
lambda x: self.on_view_button_clicked(x, self.exportOutputPath))
self.nameHorBox.addWidget(self.exportPathViewButton)
self.buttonsLayout = QtWidgets.QHBoxLayout()
self.exportButton = QPushButton("Export")
self.exportButton.setFixedWidth(60)
self.exportButton.setEnabled(False)
self.exportButton.clicked.connect(self.on_export_clicked)
self.buttonsLayout.addWidget(self.exportButton)
self.cancelButton = QPushButton("Cancel")
self.cancelButton.setFixedWidth(60)
self.cancelButton.clicked.connect(self.on_cancel_button_clicked)
self.buttonsLayout.addWidget(self.cancelButton)
self.buttonsLayout.setAlignment(QtCore.Qt.AlignBottom
| QtCore.Qt.AlignRight)
self.outerVertBoxPro.addLayout(self.labelVerBoxPro)
self.outerVertBoxPro.addLayout(self.nameHorBox)
#self.outerVertBoxPro.addLayout(self.spacer)
self.outerVertBoxPro.addLayout(self.buttonsLayout)
self.setFixedHeight(90)
self.setFixedWidth(500)
self.setLayout(self.outerVertBoxPro)
def on_view_button_clicked(self, x, folder_path=None):
if isinstance(folder_path, QTextEdit):
folder_path = folder_path.toPlainText()
elif isinstance(folder_path, QtWidgets.QLineEdit):
folder_path = folder_path.text()
if folder_path == "":
QMessageBox.warning(self, "No path selected",
"There is no path selected", QMessageBox.Ok)
return None
self.file_explore_thread = FileExplorerRunner(
folder_location=folder_path)
self.file_explore_thread.start()
def on_path_button_clicked(self):
logging.debug('on_log_out_path_button_clicked(): Instantiated')
folder_chosen = str(
QFileDialog.getExistingDirectory(self,
"Select Directory to Store Data"))
if folder_chosen == "":
logging.debug("File choose cancelled")
return
self.exportOutputPath.setText(folder_chosen)
self.exportButton.setEnabled(True)
def on_export_clicked(self):
out_path = self.exportOutputPath.text()
#initialize package manager without any values in args
package_mgr = PackageManager()
zip_function = package_mgr.zip
self.batch_thread = BatchThread()
self.batch_thread.progress_signal.connect(self.update_progress_bar)
self.batch_thread.completion_signal.connect(self.export_complete)
self.batch_thread.add_function(zip_function, out_path,
self.project_path,
self.project_data_path)
self.progress_dialog_overall = ProgressBarDialog(
self, self.batch_thread.get_load_count())
self.batch_thread.start()
self.progress_dialog_overall.show()
def update_progress_bar(self):
logging.debug('update_progress_bar(): Instantiated')
self.progress_dialog_overall.update_progress()
logging.debug('update_progress_bar(): Complete')
def export_complete(self):
logging.debug("export_complete(): Instantiated")
self.progress_dialog_overall.update_progress()
QMessageBox.information(self, "Export Complete!",
"Success! Project Exported", QMessageBox.Ok)
self.progress_dialog_overall.hide()
self.hide()
logging.debug("export_complete(): Complete")
def on_cancel_button_clicked(self, event):
logging.debug('on_cancel_button_clicked(): Instantiated')
self.hide()
logging.debug('on_cancel_button_clicked(): Complete')
def exec_(self):
logging.debug("ExportDialog(): exec_() instantiated")
result = super(ExportDialog, self).exec_()
if str(result) == str(1):
logging.debug("dialog_response(): OK was pressed")
self.hide()
return (QMessageBox.Ok)
return (QMessageBox.Cancel)