diff --git a/C2Client/C2Client/GUI.py b/C2Client/C2Client/GUI.py index 64e9b5b..9a274d5 100644 --- a/C2Client/C2Client/GUI.py +++ b/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 @@ logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - % 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("Authentication Required") + self.setModal(True) + + layout = QVBoxLayout(self) + description = QLabel("Enter the credentials to connect to the TeamServer.") + 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: diff --git a/C2Client/C2Client/grpcClient.py b/C2Client/C2Client/grpcClient.py index bed5e15..c937f23 100644 --- a/C2Client/C2Client/grpcClient.py +++ b/C2Client/C2Client/grpcClient.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: Optional[str] = None) -> 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): @@ -93,7 +105,8 @@ class GrpcClient: self.stub = TeamServerApi_pb2_grpc.TeamServerApiStub(self.channel) if token is None: - username, password = self._load_credentials_from_env() + if username is None or password is None: + username, password = self._load_credentials_from_env() token = self._authenticate(username, password) self.metadata: MetadataType = [