Merge pull request #44 from maxDcb/develop

Develop

Author: maxDcb

.gitignore

@@ -1,4 +1,11 @@
 Tests
 C2Client/build/
 .vscode
-build/
+build/
+C2Client/.cmdHistory
+C2Client/.termHistory
+C2Client/Beacon.exe
+C2Client/C2Client/Scripts/__init__.py
+C2Client/C2Client/TerminalModules/__pycache__/
+C2Client/loader.bin
+updateRelease.sh

C2Client/C2Client/GUI.py

@@ -1,16 +1,22 @@
 import argparse
 import logging
+import os
 import signal
 import sys
-from typing import Optional
+from typing import Optional, Tuple
 
 from PyQt6.QtWidgets import (
     QApplication,
+    QDialog,
+    QDialogButtonBox,
     QGridLayout,
     QHBoxLayout,
+    QLabel,
+    QLineEdit,
     QMainWindow,
     QPushButton,
     QTabWidget,
+    QVBoxLayout,
     QWidget,
 )
 
@@ -27,18 +33,75 @@
 signal.signal(signal.SIGINT, signal.SIG_DFL)
 
 
+class CredentialDialog(QDialog):
+    """Prompt for credentials when environment variables are absent."""
+
+    def __init__(self, parent: Optional[QWidget] = None, default_username: str = "") -> None:
+        super().__init__(parent)
+        self.setWindowTitle("Login")
+        self.setModal(True)
+
+        layout = QVBoxLayout(self)
+        description = QLabel("Login:")
+        description.setWordWrap(True)
+        layout.addWidget(description)
+
+        self.username_input = QLineEdit(self)
+        self.username_input.setPlaceholderText("Username")
+        if default_username:
+            self.username_input.setText(default_username)
+        layout.addWidget(self.username_input)
+
+        self.password_input = QLineEdit(self)
+        self.password_input.setPlaceholderText("Password")
+        self.password_input.setEchoMode(QLineEdit.EchoMode.Password)
+        layout.addWidget(self.password_input)
+
+        self.error_label = QLabel("Username and password are required.")
+        self.error_label.setStyleSheet("color: red;")
+        self.error_label.setVisible(False)
+        layout.addWidget(self.error_label)
+
+        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
+        buttons.accepted.connect(self._handle_accept)
+        buttons.rejected.connect(self.reject)
+        layout.addWidget(buttons)
+
+    def _handle_accept(self) -> None:
+        username = self.username_input.text().strip()
+        password = self.password_input.text()
+        if not username or not password:
+            self.error_label.setVisible(True)
+            return
+        self.accept()
+
+    def credentials(self) -> Tuple[str, str]:
+        return self.username_input.text().strip(), self.password_input.text()
+
+
 class App(QMainWindow):
     """Main application window for the C2 client."""
 
-    def __init__(self, ip: str, port: int, devMode: bool) -> None:
+    def __init__(self, ip: str, port: int, devMode: bool, credentials: Optional[Tuple[str, str]] = None) -> None:
         super().__init__()
 
         self.ip = ip
         self.port = port
         self.devMode = devMode
 
+        username: Optional[str] = None
+        password: Optional[str] = None
+        if credentials:
+            username, password = credentials
+
         try:
-            self.grpcClient = GrpcClient(self.ip, self.port, self.devMode)
+            self.grpcClient = GrpcClient(
+                self.ip,
+                self.port,
+                self.devMode,
+                username=username,
+                password=password,
+            )
         except ValueError as e:
             raise e
 
@@ -136,8 +199,20 @@ def main() -> None:
     app = QApplication(sys.argv)
     app.setStyleSheet(qdarktheme.load_stylesheet())
 
+    username = os.getenv("C2_USERNAME")
+    password = os.getenv("C2_PASSWORD")
+
+    credentials: Optional[Tuple[str, str]] = None
+    if username and password:
+        credentials = (username, password)
+    else:
+        dialog = CredentialDialog(default_username=username or "")
+        if dialog.exec() != QDialog.DialogCode.Accepted:
+            sys.exit(1)
+        credentials = dialog.credentials()
+
     try:
-        window = App(args.ip, args.port, args.dev)
+        window = App(args.ip, args.port, args.dev, credentials)
         window.show()
         sys.exit(app.exec())
     except ValueError:

C2Client/C2Client/grpcClient.py

@@ -9,7 +9,7 @@
 import os
 import sys
 import uuid
-from typing import Any, Iterable, List, Tuple
+from typing import Any, Iterable, List, Tuple, Optional
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/libGrpcMessages/build/py/')
 
@@ -34,9 +34,21 @@ class GrpcClient:
         If ``True`` the SSL hostname check is disabled.
     token:
         Bearer token used for authentication metadata.
+    username:
+        Username to authenticate with. If omitted, environment variables are used.
+    password:
+        Password to authenticate with. If omitted, environment variables are used.
     """
 
-    def __init__(self, ip: str, port: int, devMode: bool, token: str = "my-secret-token") -> None:
+    def __init__(
+        self,
+        ip: str,
+        port: int,
+        devMode: bool,
+        token: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+    ) -> None:
         env_cert_path = os.getenv('C2_CERT_PATH')
 
         if env_cert_path and os.path.isfile(env_cert_path):
@@ -91,11 +103,37 @@ def __init__(self, ip: str, port: int, devMode: bool, token: str = "my-secret-to
             raise ValueError("grpcClient: unable to connect") from exc
 
         self.stub = TeamServerApi_pb2_grpc.TeamServerApiStub(self.channel)
+
+        if token is None:
+            if username is None or password is None:
+                username, password = self._load_credentials_from_env()
+            token = self._authenticate(username, password)
+
         self.metadata: MetadataType = [
             ("authorization", f"Bearer {token}"),
             ("clientid", str(uuid.uuid4())[:16]),
         ]
 
+    def _load_credentials_from_env(self) -> Tuple[str, str]:
+        username = os.getenv("C2_USERNAME")
+        password = os.getenv("C2_PASSWORD")
+        if not username or not password:
+            raise ValueError(
+                "grpcClient: missing C2_USERNAME or C2_PASSWORD environment variables for authentication",
+            )
+        return username, password
+
+    def _authenticate(self, username: str, password: str) -> str:
+        request = TeamServerApi_pb2.AuthRequest(username=username, password=password)
+        response = self.stub.Authenticate(request)
+        if response.status != TeamServerApi_pb2.OK or not response.token:
+            message = response.message or "unknown authentication error"
+            logging.error("Authentication failed for user %s: %s", username, message)
+            raise ValueError(f"grpcClient: authentication failed: {message}")
+
+        logging.info("Authenticated against TeamServer as %s", username)
+        return response.token
+
     def getListeners(self) -> Any:
         """Return the list of listeners registered on the TeamServer."""
 

C2Client/C2Client/libGrpcMessages/build/py/TeamServerApi_pb2.py

@@ -15,6 +15,11 @@ def __init__(self, channel):
         Args:
             channel: A grpc.Channel.
         """
+        self.Authenticate = channel.unary_unary(
+                '/teamserverapi.TeamServerApi/Authenticate',
+                request_serializer=TeamServerApi__pb2.AuthRequest.SerializeToString,
+                response_deserializer=TeamServerApi__pb2.AuthResponse.FromString,
+                _registered_method=True)
         self.GetListeners = channel.unary_stream(
                 '/teamserverapi.TeamServerApi/GetListeners',
                 request_serializer=TeamServerApi__pb2.Empty.SerializeToString,
@@ -66,6 +71,12 @@ class TeamServerApiServicer(object):
     """Interface exported by the server.
     """
 
+    def Authenticate(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
     def GetListeners(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
@@ -123,6 +134,11 @@ def SendTermCmd(self, request, context):
 
 def add_TeamServerApiServicer_to_server(servicer, server):
     rpc_method_handlers = {
+            'Authenticate': grpc.unary_unary_rpc_method_handler(
+                    servicer.Authenticate,
+                    request_deserializer=TeamServerApi__pb2.AuthRequest.FromString,
+                    response_serializer=TeamServerApi__pb2.AuthResponse.SerializeToString,
+            ),
             'GetListeners': grpc.unary_stream_rpc_method_handler(
                     servicer.GetListeners,
                     request_deserializer=TeamServerApi__pb2.Empty.FromString,
@@ -180,6 +196,33 @@ class TeamServerApi(object):
     """Interface exported by the server.
     """
 
+    @staticmethod
+    def Authenticate(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/teamserverapi.TeamServerApi/Authenticate',
+            TeamServerApi__pb2.AuthRequest.SerializeToString,
+            TeamServerApi__pb2.AuthResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
     @staticmethod
     def GetListeners(request,
             target,

C2Client/C2Client/libGrpcMessages/build/py/TeamServerApi_pb2_grpc.py

@@ -17,7 +17,8 @@ dependencies = [
     "pwn==1.0",
     "pefile==2024.8.26",
     "openai==1.102.0",
-    "donut-shellcode"
+    "donut-shellcode",
+    "markdown"
 ]
 
 [project.optional-dependencies]
@@ -41,4 +42,4 @@ C2Client = [
 ]
 
 [project.scripts]
-c2client = "C2Client.GUI:main"  # Entry point for CLI tool
+c2client = "C2Client.GUI:main"  # Entry point for CLI tool

C2Client/pyproject.toml

@@ -11,3 +11,4 @@ openai==1.102.0
 pytest==8.4.1
 pytest-qt==4.5.0
 donut-shellcode
+markdown