Merge pull request #30 from maxDcb/codex/assess-and-improve-c2client-code-quality

Refactor C2 client imports and add tests

Author: maxDcb

C2Client/C2Client/AssistantPanel.py

@@ -5,11 +5,19 @@
 from datetime import datetime
 
 from threading import Thread, Lock, Semaphore
-from PyQt5.QtWidgets import *
-from PyQt5.QtGui import *
-from PyQt5.QtCore import *
 
-from grpcClient import *
+from PyQt5.QtCore import Qt, QEvent, QTimer, pyqtSignal
+from PyQt5.QtGui import QFont, QTextCursor, QStandardItem, QStandardItemModel
+from PyQt5.QtWidgets import (
+    QCompleter,
+    QLineEdit,
+    QPlainTextEdit,
+    QShortcut,
+    QVBoxLayout,
+    QWidget,
+)
+
+from .grpcClient import TeamServerApi_pb2
 
 import openai
 from openai import OpenAI

C2Client/C2Client/ConsolePanel.py

@@ -4,18 +4,27 @@
 import re, html
 from datetime import datetime
 from threading import Thread, Lock
-from PyQt5.QtWidgets import *
-from PyQt5.QtGui import *
-from PyQt5.QtCore import *
 
-from grpcClient import *
-
-from TerminalPanel import *
-from ScriptPanel import *
-from AssistantPanel import *
-
-sys.path.insert(1, './Credentials')
-import credentials
+from PyQt5.QtCore import QObject, Qt, QEvent, QThread, QTimer, pyqtSignal, pyqtSlot
+from PyQt5.QtGui import QFont, QStandardItem, QStandardItemModel, QTextCursor
+from PyQt5.QtWidgets import (
+    QWidget,
+    QTabWidget,
+    QVBoxLayout,
+    QHBoxLayout,
+    QTextEdit,
+    QLineEdit,
+    QShortcut,
+    QCompleter,
+    QTableWidget,
+    QTableWidgetItem,
+)
+
+from .grpcClient import GrpcClient, TeamServerApi_pb2
+from .TerminalPanel import Terminal
+from .ScriptPanel import Script
+from .AssistantPanel import Assistant
+from .TerminalModules.Credentials import credentials
 
 
 #
@@ -570,15 +579,13 @@ def runCommand(self):
             self.printInTerminal("", "", "")
 
         else:
-            cmdHistoryFile = open(CmdHistoryFileName, 'a')
-            cmdHistoryFile.write(commandLine)
-            cmdHistoryFile.write('\n')
-            cmdHistoryFile.close()
+            with open(CmdHistoryFileName, 'a') as cmdHistoryFile:
+                cmdHistoryFile.write(commandLine)
+                cmdHistoryFile.write('\n')
 
-            logFile = open(logsDir+"/"+self.logFileName, 'a')
-            logFile.write('[+] send: \"' + commandLine + '\"')
-            logFile.write('\n')
-            logFile.close()
+            with open(os.path.join(logsDir, self.logFileName), 'a') as logFile:
+                logFile.write('[+] send: \"' + commandLine + '\"')
+                logFile.write('\n')
 
             self.commandEditor.setCmdHistory()
             instructions = commandLine.split()
@@ -612,11 +619,10 @@ def displayResponse(self):
             self.printInTerminal("", response.instruction + " " + response.cmd, response.response.decode('utf-8', 'replace'))
             self.setCursorEditorAtEnd()
 
-            logFile = open(logsDir+"/"+self.logFileName, 'a')
-            logFile.write('[+] result: \"' + response.instruction + " " + response.cmd + '\"')
-            logFile.write('\n' + response.response.decode('utf-8', 'replace')  + '\n')
-            logFile.write('\n')
-            logFile.close()
+            with open(os.path.join(logsDir, self.logFileName), 'a') as logFile:
+                logFile.write('[+] result: \"' + response.instruction + " " + response.cmd + '\"')
+                logFile.write('\n' + response.response.decode('utf-8', 'replace')  + '\n')
+                logFile.write('\n')
 
     def setCursorEditorAtEnd(self):
         cursor = self.editorOutput.textCursor()
@@ -625,32 +631,36 @@ def setCursorEditorAtEnd(self):
 
 
 class GetSessionResponse(QObject):
+    """Background worker querying session responses."""
+
     checkin = pyqtSignal()
 
-    exit=False
+    def __init__(self) -> None:
+        super().__init__()
+        self.exit = False
 
-    def run(self):
-        while self.exit==False:
+    def run(self) -> None:
+        while not self.exit:
             self.checkin.emit()
             time.sleep(1)
 
-    def quit(self):
-        self.exit=True
+    def quit(self) -> None:
+        self.exit = True
 
 
 class CommandEditor(QLineEdit):
     tabPressed = pyqtSignal()
-    cmdHistory = []
-    idx = 0
 
-    def __init__(self, parent=None):
+    def __init__(self, parent: QWidget | None = None) -> None:
         super().__init__(parent)
 
-        if(os.path.isfile(CmdHistoryFileName)):
-            cmdHistoryFile = open(CmdHistoryFileName)
-            self.cmdHistory = cmdHistoryFile.readlines()
-            self.idx=len(self.cmdHistory)-1
-            cmdHistoryFile.close()
+        self.cmdHistory: list[str] = []
+        self.idx: int = 0
+
+        if os.path.isfile(CmdHistoryFileName):
+            with open(CmdHistoryFileName) as cmdHistoryFile:
+                self.cmdHistory = cmdHistoryFile.readlines()
+            self.idx = len(self.cmdHistory) - 1
 
         QShortcut(Qt.Key_Up, self, self.historyUp)
         QShortcut(Qt.Key_Down, self, self.historyDown)
@@ -686,11 +696,10 @@ def historyDown(self):
             cmd = self.cmdHistory[self.idx%len(self.cmdHistory)]
             self.setText(cmd.strip())
 
-    def setCmdHistory(self):
-        cmdHistoryFile = open(CmdHistoryFileName)
-        self.cmdHistory = cmdHistoryFile.readlines()
-        self.idx=len(self.cmdHistory)-1
-        cmdHistoryFile.close()
+    def setCmdHistory(self) -> None:
+        with open(CmdHistoryFileName) as cmdHistoryFile:
+            self.cmdHistory = cmdHistoryFile.readlines()
+        self.idx = len(self.cmdHistory) - 1
 
     def clearLine(self):
         self.clear()

C2Client/C2Client/GUI.py

@@ -1,20 +1,24 @@
-import sys
-import os
-import signal
 import argparse
 import logging
-
-from PyQt5.QtWidgets import *
-from PyQt5.QtGui import *
-from PyQt5.QtCore import *
-
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-
-from grpcClient import *
-from ListenerPanel import *
-from SessionPanel import *
-from ConsolePanel import *
-from GraphPanel import *
+import signal
+import sys
+from typing import Optional
+
+from PyQt5.QtWidgets import (
+    QApplication,
+    QGridLayout,
+    QHBoxLayout,
+    QMainWindow,
+    QPushButton,
+    QTabWidget,
+    QWidget,
+)
+
+from .grpcClient import GrpcClient
+from .ListenerPanel import Listeners
+from .SessionPanel import Sessions
+from .ConsolePanel import ConsolesTab
+from .GraphPanel import Graph
 
 import qdarktheme
 
@@ -24,8 +28,9 @@
 
 
 class App(QMainWindow):
+    """Main application window for the C2 client."""
 
-    def __init__(self, ip, port, devMode):
+    def __init__(self, ip: str, port: int, devMode: bool) -> None:
         super().__init__()
 
         self.ip = ip
@@ -37,7 +42,7 @@ def __init__(self, ip, port, devMode):
         except ValueError as e:
             raise e
 
-        self.createPayloadWindow = None
+        self.createPayloadWindow: Optional[QWidget] = None
 
         self.title = 'Exploration C2'
         self.left = 0
@@ -68,19 +73,17 @@ def __init__(self, ip, port, devMode):
         self.sessionsWidget.interactWithSession.connect(self.consoleWidget.addConsole)
 
         self.consoleWidget.script.mainScriptMethod("start", "", "", "")
-        
-        self.show()
-
 
-    def topLayout(self):
+    def topLayout(self) -> None:
+        """Initialise the upper part of the main window."""
 
         self.topWidget = QTabWidget()
 
         self.m_main = QWidget()
 
         self.m_main.layout = QHBoxLayout(self.m_main)
         self.m_main.layout.setContentsMargins(0, 0, 0, 0)
-        
+
         self.sessionsWidget = Sessions(self, self.grpcClient)
         self.listenersWidget = Listeners(self, self.grpcClient)
 
@@ -96,24 +99,33 @@ def topLayout(self):
         self.mainLayout.addWidget(self.topWidget, 1, 1, 1, 1)
 
 
-    def botLayout(self):
+    def botLayout(self) -> None:
+        """Initialise the bottom console area."""
 
         self.consoleWidget = ConsolesTab(self, self.grpcClient)
         self.mainLayout.addWidget(self.consoleWidget, 2, 0, 1, 2)
 
 
-    def __del__(self):
+    def __del__(self) -> None:
+        """Ensure scripts are stopped when the window is destroyed."""
         if hasattr(self, 'consoleWidget'):
             self.consoleWidget.script.mainScriptMethod("stop", "", "", "")
 
 
-    def payloadForm(self):
+    def payloadForm(self) -> None:
+        """Display the payload creation window."""
         if self.createPayloadWindow is None:
+            try:
+                from .ScriptPanel import CreatePayload  # type: ignore
+            except Exception:
+                CreatePayload = QWidget  # fallback to simple widget
             self.createPayloadWindow = CreatePayload()
         self.createPayloadWindow.show()
 
 
-def main():
+def main() -> None:
+    """Entry point used by the project script."""
+
     parser = argparse.ArgumentParser(description='TeamServer IP and port.')
     parser.add_argument('--ip', default='127.0.0.1', help='IP address (default: 127.0.0.1)')
     parser.add_argument('--port', type=int, default=50051, help='Port number (default: 50051)')
@@ -125,11 +137,13 @@ def main():
     app.setStyleSheet(qdarktheme.load_stylesheet())
 
     try:
-        ex = App(args.ip, args.port, args.dev)
-    except ValueError as e:
+        window = App(args.ip, args.port, args.dev)
+        window.show()
+        sys.exit(app.exec_())
+    except ValueError:
         sys.exit(1)
     sys.exit(app.exec_())
 
 
 if __name__ == "__main__":
-    main()
+    main()

C2Client/C2Client/GraphPanel.py

@@ -2,12 +2,18 @@
 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 *
+from PyQt5.QtCore import QObject, Qt, QThread, QLineF, pyqtSignal
+from PyQt5.QtGui import QColor, QFont, QPainter, QPen, QPixmap
+from PyQt5.QtWidgets import (
+    QGraphicsLineItem,
+    QGraphicsPixmapItem,
+    QGraphicsScene,
+    QGraphicsView,
+    QVBoxLayout,
+    QWidget,
+    QGraphicsItem,
+)
 
 
 #

C2Client/C2Client/ListenerPanel.py

@@ -1,11 +1,23 @@
 import time
 import logging
 
-from PyQt5.QtWidgets import *
-from PyQt5.QtGui import *
-from PyQt5.QtCore import *
-
-from grpcClient import *
+from PyQt5.QtCore import Qt, QThread, pyqtSignal, QObject
+from PyQt5.QtWidgets import (
+    QComboBox,
+    QFormLayout,
+    QGridLayout,
+    QLabel,
+    QLineEdit,
+    QMenu,
+    QPushButton,
+    QTableView,
+    QTableWidget,
+    QTableWidgetItem,
+    QWidget,
+    QHeaderView,
+)
+
+from .grpcClient import TeamServerApi_pb2
 
 
 #