diff --git a/client/GraphPanel.py b/client/GraphPanel.py index e4f0dc0..19329a7 100644 --- a/client/GraphPanel.py +++ b/client/GraphPanel.py @@ -18,7 +18,9 @@ ListenerNodeItemType = "Listener" PrimaryListenerImage = "images/firewall.svg" WindowsSessionImage = "images/pc.svg" +WindowsHighPrivSessionImage = "images/windowshighpriv.svg" LinuxSessionImage = "images/linux.svg" +LinuxRootSessionImage = "images/linuxhighpriv.svg" # @@ -36,10 +38,10 @@ class NodeItem(QGraphicsPixmapItem): # Signal to notify position changes signaller = Signaller() - def __init__(self, type, hash, os="", privilege="", parent=None): + def __init__(self, type, hash, os="", privilege="", hostname="", parent=None): if type == ListenerNodeItemType: self.type = ListenerNodeItemType - pixmap = QPixmap(PrimaryListenerImage).scaled(64, 64, Qt.KeepAspectRatio, Qt.SmoothTransformation) + pixmap = self.addImageNode(PrimaryListenerImage, "") self.beaconHash = "" self.connectedListenerHash = "" self.listenerHash = [] @@ -48,12 +50,19 @@ class NodeItem(QGraphicsPixmapItem): self.type = BeaconNodeItemType # print("NodeItem beaconHash", hash, "os", os, "privilege", privilege) if "linux" in os.lower(): - pixmap = QPixmap(LinuxSessionImage).scaled(64, 64, Qt.KeepAspectRatio, Qt.SmoothTransformation) + if privilege == "root": + pixmap = self.addImageNode(LinuxRootSessionImage, hostname) + else: + pixmap = self.addImageNode(LinuxSessionImage, hostname) elif "windows" in os.lower(): - pixmap = QPixmap(WindowsSessionImage).scaled(64, 64, Qt.KeepAspectRatio, Qt.SmoothTransformation) + 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=[] @@ -79,6 +88,38 @@ class NodeItem(QGraphicsPixmapItem): 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): @@ -184,7 +225,7 @@ class Graph(QWidget): if session.beaconHash == nodeItem.beaconHash: inStore=True if not inStore: - item = NodeItem(BeaconNodeItemType, session.beaconHash, session.os, session.privilege) + 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) diff --git a/client/images/linuxhighpriv.svg b/client/images/linuxhighpriv.svg new file mode 100644 index 0000000..d3d1a6c --- /dev/null +++ b/client/images/linuxhighpriv.svg @@ -0,0 +1,2 @@ + + diff --git a/client/images/windowshighpriv.svg b/client/images/windowshighpriv.svg new file mode 100644 index 0000000..7d421d0 --- /dev/null +++ b/client/images/windowshighpriv.svg @@ -0,0 +1,2 @@ + +