1818
1919PrimaryListenerImage = "images/firewall.svg"
2020WindowsSessionImage = "images/pc.svg"
21+ WindowsHighPrivSessionImage = "images/windowshighpriv.svg"
2122LinuxSessionImage = "images/linux.svg"
23+ LinuxRootSessionImage = "images/linuxhighpriv.svg"
2224
2325
2426#
@@ -36,10 +38,10 @@ class NodeItem(QGraphicsPixmapItem):
3638 # Signal to notify position changes
3739 signaller = Signaller ()
3840
39- def __init__ (self , type , hash , os = "" , privilege = "" , parent = None ):
41+ def __init__ (self , type , hash , os = "" , privilege = "" , hostname = "" , parent = None ):
4042 if type == ListenerNodeItemType :
4143 self .type = ListenerNodeItemType
42- pixmap = QPixmap ( PrimaryListenerImage ). scaled ( 64 , 64 , Qt . KeepAspectRatio , Qt . SmoothTransformation )
44+ pixmap = self . addImageNode ( PrimaryListenerImage , "" )
4345 self .beaconHash = ""
4446 self .connectedListenerHash = ""
4547 self .listenerHash = []
@@ -48,12 +50,19 @@ def __init__(self, type, hash, os="", privilege="", parent=None):
4850 self .type = BeaconNodeItemType
4951 # print("NodeItem beaconHash", hash, "os", os, "privilege", privilege)
5052 if "linux" in os .lower ():
51- pixmap = QPixmap (LinuxSessionImage ).scaled (64 , 64 , Qt .KeepAspectRatio , Qt .SmoothTransformation )
53+ if privilege == "root" :
54+ pixmap = self .addImageNode (LinuxRootSessionImage , hostname )
55+ else :
56+ pixmap = self .addImageNode (LinuxSessionImage , hostname )
5257 elif "windows" in os .lower ():
53- pixmap = QPixmap (WindowsSessionImage ).scaled (64 , 64 , Qt .KeepAspectRatio , Qt .SmoothTransformation )
58+ if privilege == "HIGH" :
59+ pixmap = self .addImageNode (WindowsHighPrivSessionImage , hostname )
60+ else :
61+ pixmap = self .addImageNode (WindowsSessionImage , hostname )
5462 else :
5563 pixmap = QPixmap (LinuxSessionImage ).scaled (64 , 64 , Qt .KeepAspectRatio , Qt .SmoothTransformation )
5664 self .beaconHash = hash
65+ self .hostname = hostname
5766 self .connectedListenerHash = ""
5867 self .listenerHash = []
5968
@@ -79,6 +88,38 @@ def mousePressEvent(self, event):
7988 def mouseReleaseEvent (self , event ):
8089 super ().mouseReleaseEvent (event )
8190 self .setCursor (Qt .ArrowCursor )
91+
92+ def addImageNode (self , image_path , legend_text , font_size = 9 , padding = 5 , text_color = Qt .white ):
93+ # Load and scale the image
94+ pixmap = QPixmap (image_path ).scaled (64 , 64 , Qt .KeepAspectRatio , Qt .SmoothTransformation )
95+
96+ # Create a new QPixmap larger than the original for the image and text
97+ legend_height = font_size + padding * 2
98+ legend_width = len (legend_text ) * font_size + padding * 2
99+ combined_pixmap = QPixmap (max (legend_width , pixmap .width ()), pixmap .height () + legend_height )
100+ combined_pixmap .fill (Qt .transparent ) # Transparent background
101+
102+ # Paint the image and the legend onto the combined pixmap
103+ painter = QPainter (combined_pixmap )
104+ image_x = (combined_pixmap .width () - pixmap .width ()) // 2
105+ painter .drawPixmap (image_x , 0 , pixmap ) # Draw the image
106+
107+ pen = QPen ()
108+ pen .setColor (text_color ) # Set the desired text color
109+ painter .setPen (pen )
110+ # Set font for the legend
111+ font = QFont ()
112+ font .setPointSize (font_size )
113+ painter .setFont (font )
114+
115+ # Draw the legend text centered below the image
116+ text_rect = painter .boundingRect (
117+ 0 , pixmap .height (), combined_pixmap .width (), legend_height , Qt .AlignCenter , legend_text
118+ )
119+ painter .drawText (text_rect , Qt .AlignCenter , legend_text )
120+
121+ painter .end ()
122+ return combined_pixmap
82123
83124
84125class Connector (QGraphicsLineItem ):
@@ -184,7 +225,7 @@ def updateGraph(self):
184225 if session .beaconHash == nodeItem .beaconHash :
185226 inStore = True
186227 if not inStore :
187- item = NodeItem (BeaconNodeItemType , session .beaconHash , session .os , session .privilege )
228+ item = NodeItem (BeaconNodeItemType , session .beaconHash , session .os , session .privilege , session . hostname )
188229 item .connectedListenerHash = session .listenerHash
189230 item .signaller .signal .connect (self .updateConnectors )
190231 self .scene .addItem (item )
0 commit comments