Files
maxDcb-C2TeamServer/client/GraphPanel.py
T
2024-12-02 10:11:49 -05:00

329 lines
12 KiB
Python

import sys
import os
import time
from threading import Thread, Lock
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtGui import QPixmap, QTransform
from grpcClient import *
#
# Constant
#
BeaconNodeItemType = "Beacon"
ListenerNodeItemType = "Listener"
PrimaryListenerImage = "images/firewall.svg"
WindowsSessionImage = "images/pc.svg"
WindowsHighPrivSessionImage = "images/windowshighpriv.svg"
LinuxSessionImage = "images/linux.svg"
LinuxRootSessionImage = "images/linuxhighpriv.svg"
#
# Graph Tab Implementation
#
# needed to send the message of mouseMoveEvent because QGraphicsPixmapItem doesn't herit from QObject
class Signaller(QObject):
signal = pyqtSignal()
def trigger(self):
self.signal.emit()
class NodeItem(QGraphicsPixmapItem):
# Signal to notify position changes
signaller = Signaller()
def __init__(self, type, hash, os="", privilege="", hostname="", parent=None):
if type == ListenerNodeItemType:
self.type = ListenerNodeItemType
pixmap = self.addImageNode(PrimaryListenerImage, "")
self.beaconHash = ""
self.connectedListenerHash = ""
self.listenerHash = []
self.listenerHash.append(hash)
elif type == BeaconNodeItemType:
self.type = BeaconNodeItemType
# print("NodeItem beaconHash", hash, "os", os, "privilege", privilege)
if "linux" in os.lower():
if privilege == "root":
pixmap = self.addImageNode(LinuxRootSessionImage, hostname)
else:
pixmap = self.addImageNode(LinuxSessionImage, hostname)
elif "windows" in os.lower():
if privilege == "HIGH":
pixmap = self.addImageNode(WindowsHighPrivSessionImage, hostname)
else:
pixmap = self.addImageNode(WindowsSessionImage, hostname)
else:
pixmap = QPixmap(LinuxSessionImage).scaled(64, 64, Qt.KeepAspectRatio, Qt.SmoothTransformation)
self.beaconHash=hash
self.hostname = hostname
self.connectedListenerHash = ""
self.listenerHash=[]
super().__init__(pixmap)
def print(self):
print("NodeItem", self.type, "beaconHash", self.beaconHash, "listenerHash", self.listenerHash, "connectedListenerHash", self.connectedListenerHash)
def isResponsableForListener(self, hash):
if hash in self.listenerHash:
return True
else:
return False
def mouseMoveEvent(self, event):
super().mouseMoveEvent(event)
self.signaller.trigger()
def mousePressEvent(self, event):
super().mousePressEvent(event)
self.setCursor(Qt.ClosedHandCursor)
def mouseReleaseEvent(self, event):
super().mouseReleaseEvent(event)
self.setCursor(Qt.ArrowCursor)
def addImageNode(self, image_path, legend_text, font_size=9, padding=5, text_color=Qt.white):
# Load and scale the image
pixmap = QPixmap(image_path).scaled(64, 64, Qt.KeepAspectRatio, Qt.SmoothTransformation)
# Create a new QPixmap larger than the original for the image and text
legend_height = font_size + padding * 2
legend_width = len(legend_text) * font_size + padding * 2
combined_pixmap = QPixmap(max(legend_width, pixmap.width()), pixmap.height() + legend_height)
combined_pixmap.fill(Qt.transparent) # Transparent background
# Paint the image and the legend onto the combined pixmap
painter = QPainter(combined_pixmap)
image_x = (combined_pixmap.width() - pixmap.width()) // 2
painter.drawPixmap(image_x, 0, pixmap) # Draw the image
pen = QPen()
pen.setColor(text_color) # Set the desired text color
painter.setPen(pen)
# Set font for the legend
font = QFont()
font.setPointSize(font_size)
painter.setFont(font)
# Draw the legend text centered below the image
text_rect = painter.boundingRect(
0, pixmap.height(), combined_pixmap.width(), legend_height, Qt.AlignCenter, legend_text
)
painter.drawText(text_rect, Qt.AlignCenter, legend_text)
painter.end()
return combined_pixmap
class Connector(QGraphicsLineItem):
def __init__(self, listener, beacon, pen=None):
super().__init__()
self.listener = listener
self.beacon = beacon
self.pen = pen or QPen(QColor("white"), 3)
self.setPen(self.pen)
self.update_line()
def print(self):
print("Connector", "beaconHash", self.beacon.beaconHash, "connectedListenerHash", self.beacon.connectedListenerHash, "listenerHash", self.listener.listenerHash)
def update_line(self):
# print("listener", self.listener.pos())
# print("beacon", self.beacon.pos())
center1 = self.listener.pos() + self.listener.boundingRect().center()
center2 = self.beacon.pos() + self.beacon.boundingRect().center()
self.setLine(QLineF(center1, center2))
class Graph(QWidget):
listNodeItem = []
listNodeItem = []
listConnector = []
def __init__(self, parent, ip, port, devMode):
super(QWidget, self).__init__(parent)
width = self.frameGeometry().width()
height = self.frameGeometry().height()
self.ip = ip
self.port = port
self.grpcClient = GrpcClient(ip, port, devMode)
self.scene = QGraphicsScene()
self.view = QGraphicsView(self.scene)
self.view.setRenderHint(QPainter.Antialiasing)
self.vbox = QVBoxLayout()
self.vbox.setContentsMargins(0, 0, 0, 0)
self.vbox.addWidget(self.view)
self.setLayout(self.vbox)
self.thread = QThread()
self.getGraphInfoWorker = GetGraphInfoWorker()
self.getGraphInfoWorker.moveToThread(self.thread)
self.thread.started.connect(self.getGraphInfoWorker.run)
self.getGraphInfoWorker.checkin.connect(self.updateGraph)
self.thread.start()
# self.updateScene()
def __del__(self):
self.getGraphInfoWorker.quit()
self.thread.quit()
self.thread.wait()
def updateConnectors(self):
for connector in self.listConnector:
connector.update_line()
# Update the graphe every X sec with information from the team server
def updateGraph(self):
#
# Update beacons
#
responses = self.grpcClient.getSessions()
sessions = list()
for response in responses:
sessions.append(response)
# delete beacon
for ix, nodeItem in enumerate(self.listNodeItem):
runing=False
for session in sessions:
if session.beaconHash == nodeItem.beaconHash:
runing=True
if not runing and self.listNodeItem[ix].type == BeaconNodeItemType:
for ix2, connector in enumerate(self.listConnector):
if connector.beacon.beaconHash == nodeItem.beaconHash:
print("[-] delete connector")
self.scene.removeItem(self.listConnector[ix2])
del self.listConnector[ix2]
print("[-] delete beacon", nodeItem.beaconHash)
self.scene.removeItem(self.listNodeItem[ix])
del self.listNodeItem[ix]
# add beacon
for session in sessions:
inStore=False
for ix, nodeItem in enumerate(self.listNodeItem):
if session.beaconHash == nodeItem.beaconHash:
inStore=True
if not inStore:
item = NodeItem(BeaconNodeItemType, session.beaconHash, session.os, session.privilege, session.hostname)
item.connectedListenerHash = session.listenerHash
item.signaller.signal.connect(self.updateConnectors)
self.scene.addItem(item)
self.listNodeItem.append(item)
print("[+] add beacon", session.beaconHash)
#
# Update listener
#
responses= self.grpcClient.getListeners()
listeners = list()
for listener in responses:
listeners.append(listener)
# delete listener
for ix, nodeItem in enumerate(self.listNodeItem):
runing=False
for listener in listeners:
if nodeItem.isResponsableForListener(listener.listenerHash):
runing=True
if not runing:
# primary listener
if self.listNodeItem[ix].type == ListenerNodeItemType:
for ix2, connector in enumerate(self.listConnector):
if self.listNodeItem[ix2].listenerHash in connector.listener.listenerHash:
print("[-] delete connector")
self.scene.removeItem(self.listConnector[ix2])
del self.listConnector[ix2]
print("[-] delete primary listener", nodeItem.listenerHash)
self.scene.removeItem(self.listNodeItem[ix])
del self.listNodeItem[ix]
# beacon listener
elif self.listNodeItem[ix].type == BeaconNodeItemType:
if listener.listenerHash in self.listNodeItem[ix].listenerHash:
for ix2, connector in enumerate(self.listConnector):
if self.listNodeItem[ix2].listenerHash in connector.listener.listenerHash:
print("[-] delete connector")
self.scene.removeItem(self.listConnector[ix2])
del self.listConnector[ix2]
print("[-] delete secondary listener", nodeItem.listenerHash)
self.listNodeItem[ix].listenerHash.remove(listener.listenerHash)
# add listener
for listener in listeners:
inStore=False
for ix, nodeItem in enumerate(self.listNodeItem):
if nodeItem.isResponsableForListener(listener.listenerHash):
inStore=True
if not inStore:
if not listener.beaconHash:
item = NodeItem(ListenerNodeItemType, listener.listenerHash)
item.signaller.signal.connect(self.updateConnectors)
self.scene.addItem(item)
self.listNodeItem.append(item)
print("[+] add primary listener", listener.listenerHash)
else:
for nodeItem2 in self.listNodeItem:
if nodeItem2.beaconHash == listener.beaconHash:
nodeItem2.listenerHash.append(listener.listenerHash)
print("[+] add secondary listener", listener.listenerHash)
#
# Update connectors
#
for nodeItem in self.listNodeItem:
if nodeItem.type == BeaconNodeItemType:
inStore=False
beaconHash = nodeItem.beaconHash
listenerHash = nodeItem.connectedListenerHash
for connector in self.listConnector:
if connector.listener.isResponsableForListener(listenerHash) and connector.beacon.beaconHash == beaconHash:
inStore=True
if not inStore:
for listener in self.listNodeItem:
if listener.isResponsableForListener(listenerHash)==True:
connector = Connector(listener, nodeItem)
self.scene.addItem(connector)
connector.setZValue(-1)
self.listConnector.append(connector)
print("[+] add connector listener:", listenerHash, "beacon", beaconHash)
for item in self.listNodeItem:
item.setFlag(QGraphicsItem.ItemIsMovable)
item.setFlag(QGraphicsItem.ItemIsSelectable)
class GetGraphInfoWorker(QObject):
checkin = pyqtSignal()
exit=False
def run(self):
while self.exit==False:
self.checkin.emit()
time.sleep(5)
def quit(self):
self.exit=True