class VectorManager(DataBaseManager):
TABLE = "Vectors"
def __init__(self):
super().__init__()
self.table = self.db[VectorManager.TABLE]
self.nodeManager = NodeManager()
def addVector(self, name, desc):
v = {
"name": name,
"desc": desc,
"nodes": []
}# Vector(name, desc)
self.add(v)
def vectorExists(self, name):
for vector in self.get(None):
if vector["name"] == name:
return True
return False
def getVectors(self):
vectors = [
Vector(vector["name"], vector["desc"], vector["nodes"]) for vector in self.get(None)
]
return vectors
def getVectorByName(self, name):
query = {"name":name}
results = self.get(query)
vector = None
for result in results:
vector = Vector(
result["name"],
result["desc"],
result["nodes"]
)
break
return vector
def associateLogEntry(self, logentry, vector):
self.nodeManager.addNode(logentry)
node = self.nodeManager.getNodeByLogRef(logentry.getNumber())
self.table.update({"name": vector.getName()}, {"$push": {"nodes":node.getId()}})
def updateVector(self, vector_name, name, desc):
query = {"name": vector_name}
update = {"$set": {"name": name, "desc": desc} }
self.update(query, update)
def deleteVector(self, name):
query = {"name": name}
self.delete(query)
def setCurrentVector(self, name):
self.curVec = self.getVectorByName(name)
def getCurrentVector(self):
return self.curVec
class GraphManager(object):
def __init__(self):
self.graphs = dict()
self.nodeManager = NodeManager()
def getGraph(self, vector):
if vector.getName() in self.graphs.keys():
return self.graphs[vector.getName()].getGraph()
# Generate the graph if we havent generated it for the vecotr
self.generateGraph(vector)
return self.graphs[vector.getName()].getGraph()
def generateGraph(self, vector):
gGenerator = GraphGenerator(vector.getName())
nodes = vector.getNodes()
for nodeid in nodes:
node = self.nodeManager.getNode(nodeid)
gGenerator.addNode(**{
"name": node.getName(),
"label": "",
"shape": node.getIcon()
})
gGenerator.build()
self.graphs[vector.getName()] = gGenerator
def addNode(self, vector, **kwargs):
self.graphs[vector.getName()].addNode(**kwargs)
self.graphs[vector.getName()].build()
def addEdge(self, vector, src, dest):
self.graphs[vector.getName()].addEdge(src, dest)
self.graphs[vector.getName()].build()
def __init__(self):
super().__init__()
self.table = self.db[VectorManager.TABLE]
self.nodeManager = NodeManager()
def __init__(self, parent=None):
super(QWidget, self).__init__(parent)
self.graphManager = GraphManager()
self.nodeManager = NodeManager()
self.vector = None
self.initUI()
class GraphWidget(QWidget):
def __init__(self, parent=None):
super(QWidget, self).__init__(parent)
self.graphManager = GraphManager()
self.nodeManager = NodeManager()
self.vector = None
self.initUI()
def initUI(self):
self.setLayout(QVBoxLayout())
graphControls = QHBoxLayout()
addNode = QPushButton("Add Node")
addNode.clicked.connect(self.addNode)
graphControls.addWidget(addNode)
addEdge = QPushButton("Add Edge")
addEdge.clicked.connect(self.addEdge)
graphControls.addWidget(addEdge)
self.layout().addLayout(graphControls)
# selectedVector = self.vectorManager.getCurrentVector()
# if selectedVector:
# graphGenerator.generateVectorGraph(selectedVector)
# self.graphGenerator.build()
if self.vector:
qgv = self.graphManager.getGraph(self.vector)
self.layout().addWidget(qgv)
else:
self.layout().addWidget(QWidget())
def updateGraph(self, vector):
# clear current widget
for i in reversed(range(self.layout().count())):
widget = self.layout().itemAt(i).widget()
if widget:
widget.setParent(None)
self.vector = vector
qgv = self.graphManager.getGraph(self.vector)
self.layout().addWidget(qgv)
def addNode(self):
dlg = QDialog()
dlg.ok = False
dlg.node_name = ""
dlg.node_label = ""
dlg.node_type = "None"
addNodeWidget = AddNodeWidget(dlg)
dlg.setLayout(addNodeWidget.layout())
dlg.exec()
result = dlg.result()
if result == QDialog.Accepted:
results = addNodeWidget.getResults()
if results[2] == "image":
#TODO validate path to image
self.graphManager.addNode(self.vector,
name=results[0],
label=results[1],
shape=results[3])
else:
self.graphManager.addNode(self.vector,
name=results[0],
label=results[1],
shape=results[2])
else:
# Just putting this here in case we need to handel the rejected case
pass
def addEdge(self):
# nodes = self.graphGenerator.getNodes()
dlg = QDialog()
nodes = []
for id in self.vector.getNodes():
node = self.nodeManager.getNode(id)
nodes.append(node)
addEdgeWidget = AddEdgeWidget([node.getName() for node in nodes], dlg)
dlg.setLayout(addEdgeWidget.layout())
dlg.exec()
result = dlg.result()
if result == QDialog.Accepted:
results = addEdgeWidget.getResults()
self.graphManager.addEdge(self.vector, results[0], results[1])
def __init__(self, parent=None):
super(QWidget, self).__init__(parent)
self.vectorManager = VectorManager()
self.nodeManager = NodeManager()
self.logentryManager = LogEntryManager.get_instance()
self.initUI()
class AnalysisView(QWidget):
def __init__(self, parent=None):
super(QWidget, self).__init__(parent)
self.vectorManager = VectorManager()
self.nodeManager = NodeManager()
self.logentryManager = LogEntryManager.get_instance()
self.initUI()
def initUI(self):
self.setupMainLayout()
self.setLayout(self.mainLayout)
def setupMainLayout(self):
#Log Entries table
self.logEntriesTbl = QTableView()
self.logEntryModel = QStandardItemModel()
self.logEntryModel.setHorizontalHeaderLabels([
'Host', 'Timestamp', 'Content', 'Source', 'Source Type',
'Associated Vectors'
])
self.logEntriesTbl.setModel(self.logEntryModel)
self.logEntriesTbl.horizontalHeader().setSectionResizeMode(
QHeaderView.Stretch)
self.setupVectorTab()
self.tabWidget = QTabWidget()
self.tabWidget.addTab(self.logEntriesTbl, "Log Entries")
self.tabWidget.addTab(self.vectorTab, "Vector View")
# Label for our Vector List
vectorLbl = QListWidgetItem()
vectorLbl.setText("Vector Databases")
vectorLbl.setTextAlignment(Qt.AlignCenter)
vectorLbl.setFlags(Qt.NoItemFlags)
self.workspace = QHBoxLayout()
self.workspace.addWidget(self.tabWidget, 90)
# Filtering and search
hSpacer = QSpacerItem(0, 0, QSizePolicy.Expanding, QSizePolicy.Fixed)
self.search = QLineEdit()
self.search.setFixedWidth(250)
self.filterBox = QComboBox()
# Container for tab/table controls
self.controls = QHBoxLayout()
self.controls.addItem(hSpacer)
self.controls.addWidget(self.search)
self.controls.addWidget(self.filterBox)
# Container for all workspace
self.mainLayout = QVBoxLayout()
self.mainLayout.addLayout(self.controls)
self.mainLayout.addLayout(self.workspace)
def setupVectorTab(self):
self.graph = GraphWidget()
self.nodes = QTableView()
nodeModel = QStandardItemModel()
nodeModel.setHorizontalHeaderLabels([
"Name", "Timestamp", "Description", "Log Entry Refrence",
"Log Creator", "Icon", "Source", "Visible"
])
self.nodes.setModel(nodeModel)
# self.nodes.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
self.vectorViews = QHBoxLayout()
self.vectorViews.addWidget(self.nodes, 30)
self.vectorViews.addWidget(self.graph, 70)
hSpacer = QSpacerItem(0, 0, QSizePolicy.Expanding, QSizePolicy.Fixed)
vectorsLable = QLabel("Vectors")
self.vectorsCB = QComboBox()
self.vectorsCB.currentIndexChanged.connect(self.vectorHandle)
self.vectorAdded()
# self.unitsCb = QComboBox()
# self.interval = QLineEdit()
# Graph controls such as orientation, interval units, interval
self.graphContols = QHBoxLayout()
self.graphContols.addItem(hSpacer)
self.graphContols.addWidget(vectorsLable)
self.graphContols.addWidget(self.vectorsCB)
# self.graphContols.addWidget(self.unitsCb)
# self.graphContols.addWidget(self.interval)
self.container = QVBoxLayout()
self.container.addLayout(self.graphContols)
self.container.addLayout(self.vectorViews)
self.vectorTab = QWidget()
self.vectorTab.setLayout(self.container)
def vectorAdded(self):
self.vectorsCB.clear()
self.vectorsCB.addItem("")
vectors = self.vectorManager.getVectors()
for vector in vectors:
self.vectorsCB.addItem(vector.getName())
def vectorHandle(self, item):
vectorname = self.vectorsCB.currentText()
vector = self.vectorManager.getVectorByName(vectorname)
if vector == None:
return
# Update node tableview
nodeModel = self.nodes.model()
nodeModel.removeRows(0, nodeModel.rowCount())
for nodeId in vector.getNodes():
node = self.nodeManager.getNode(nodeId)
# "Name", "Timestamp", "Description", "Log Entry Refrence", "Log Creator", "Icon", "Source", "Visible"
nodeModel.appendRow([
QStandardItem(node.getName()),
QStandardItem(node.getTimeStamp()),
QStandardItem(node.getDesc()),
QStandardItem(node.getLogEntryRef()),
QStandardItem(node.getLogCreator()),
QStandardItem(node.getIcon()),
QStandardItem(node.getSource()),
QStandardItem(node.getVisible())
])
self.nodes.setModel(nodeModel)
# Todo add nodes to graph
self.graph.updateGraph(vector)
def addLogEntry(self, logentry):
host = QStandardItem(logentry.getHost())
timestamp = QStandardItem(logentry.getTimestamp())
content = QStandardItem(logentry.getContent())
source = QStandardItem(logentry.getSource())
sourcetype = QStandardItem(logentry.getSourceType())
testWidget = QtWidgets.QWidget()
combobox = CheckableComboBox(testWidget)
combobox.itemcheck_callback.connect(self.handelLogEntryChange)
vectors = self.vectorManager.getVectors()
for vector in vectors:
combobox.addItem(vector.name)
self.logEntryModel.appendRow([
host,
timestamp,
content,
source,
sourcetype,
])
row = self.logEntryModel.rowCount() - 1
col = self.logEntryModel.columnCount() - 1
a = self.logEntryModel.index(row, col)
self.logEntriesTbl.setIndexWidget(a, combobox)
def handelLogEntryChange(self, item):
if item.checkState() == Qt.Checked:
vectorName = item.text()
row = item.row()
logEntryContent = self.logEntryModel.item(row, 2).text()
logentry = self.logentryManager.getEntryByContent(logEntryContent)
vector = self.vectorManager.getVectorByName(vectorName)
self.vectorManager.associateLogEntry(logentry, vector)
else:
# TODO: Remove association
print("uncheck")
def __init__(self):
self.graphs = dict()
self.nodeManager = NodeManager()