diff --git a/core/dukto.py b/core/dukto.py index 37d24eb..1f11d65 100644 --- a/core/dukto.py +++ b/core/dukto.py @@ -34,13 +34,19 @@ class DuktoProtocol: self.peers: Dict[str, Peer] = {} self.is_sending = False self.is_receiving = False + self.is_awaiting_approval = False self.running = False + # Confirmation for receiving + self._transfer_decision = threading.Event() + self._transfer_approved = False + # Callbacks self.on_peer_added: Optional[Callable[[Peer], None]] = None self.on_peer_removed: Optional[Callable[[Peer], None]] = None self.on_receive_start: Optional[Callable[[str], None]] = None + self.on_receive_request: Optional[Callable[[str], None]] = None self.on_receive_complete: Optional[Callable[[List[str], int], None]] = None self.on_receive_text: Optional[Callable[[str, int], None]] = None self.on_send_complete: Optional[Callable[[List[str]], None]] = None @@ -143,6 +149,14 @@ class DuktoProtocol: threading.Thread(target=self._send_text_thread, args=(ip_dest, port, text), daemon=True).start() + def approve_transfer(self): + self._transfer_approved = True + self._transfer_decision.set() + + def reject_transfer(self): + self._transfer_approved = False + self._transfer_decision.set() + def _udp_listener(self): while self.running: try: @@ -188,16 +202,35 @@ class DuktoProtocol: while self.running: try: conn, addr = self.tcp_server.accept() - if self.is_receiving or self.is_sending: + if self.is_receiving or self.is_sending or self.is_awaiting_approval: conn.close() continue - threading.Thread(target=self._receive_files, + threading.Thread(target=self._handle_transfer_request, args=(conn, addr[0]), daemon=True).start() except Exception as e: if self.running: print(f"TCP listener error: {e}") - + + def _handle_transfer_request(self, conn: socket.socket, sender_ip: str): + try: + self.is_awaiting_approval = True + self._transfer_decision.clear() + + if self.on_receive_request: + self.on_receive_request(sender_ip) + else: + self.reject_transfer() + + self._transfer_decision.wait() + + if self._transfer_approved: + self._receive_files(conn, sender_ip) + else: + conn.close() + finally: + self.is_awaiting_approval = False + def _receive_files(self, conn: socket.socket, sender_ip: str): self.is_receiving = True diff --git a/main.py b/main.py index 6234f30..ae31d44 100644 --- a/main.py +++ b/main.py @@ -9,6 +9,7 @@ from core.web_search import MullvadLetaWrapper from core.discord_presence import presence from core.app_launcher import list_apps, launch from core.updater import update_repository, is_update_available +from core.dukto import DuktoProtocol ASSET = Path(__file__).parent / "assets" / "2ktan.png" @@ -318,8 +319,9 @@ class WebSearchResults(QtWidgets.QDialog): class MainWindow(QtWidgets.QMainWindow): show_menu_signal = QtCore.Signal() + receive_request_signal = QtCore.Signal(str) - def __init__(self, restart=False, no_quit=False, super_menu=True): + def __init__(self, dukto_handler, restart=False, no_quit=False, super_menu=True): super().__init__() flags = ( @@ -344,6 +346,9 @@ class MainWindow(QtWidgets.QMainWindow): self.setMask(mask) self.super_menu = super_menu + self.dukto_handler = dukto_handler + self.dukto_handler.on_receive_request = self.on_dukto_receive_request + self.receive_request_signal.connect(self.show_receive_confirmation) self.tray = QtWidgets.QSystemTrayIcon(self) self.tray.setIcon(QtGui.QIcon(str(ASSET))) @@ -418,6 +423,22 @@ class MainWindow(QtWidgets.QMainWindow): def toggle_visible(self): self.setVisible(not self.isVisible()) + def on_dukto_receive_request(self, sender_ip: str): + self.receive_request_signal.emit(sender_ip) + + def show_receive_confirmation(self, sender_ip: str): + reply = QtWidgets.QMessageBox.question( + self, + "Incoming Transfer", + f"You have an incoming transfer from {sender_ip}.\nDo you want to accept it?", + QtWidgets.QMessageBox.StandardButton.Yes | QtWidgets.QMessageBox.StandardButton.No, + QtWidgets.QMessageBox.StandardButton.No + ) + if reply == QtWidgets.QMessageBox.StandardButton.Yes: + self.dukto_handler.approve_transfer() + else: + self.dukto_handler.reject_transfer() + def start_app_launcher(self): self.app_launcher_dialog = AppLauncherDialog(self) self.app_launcher_dialog.move(QtGui.QCursor.pos()) @@ -519,6 +540,7 @@ class MainWindow(QtWidgets.QMainWindow): def restart_application(self): presence.end() + self.dukto_handler.shutdown() args = [sys.executable] + sys.argv @@ -534,9 +556,16 @@ def main(): restart = "--restart" in sys.argv no_quit = "--no-quit" in sys.argv super_menu = not "--no-super" in sys.argv - pet = MainWindow(restart=restart, no_quit=no_quit, super_menu=super_menu) + + dukto_handler = DuktoProtocol() + + pet = MainWindow(dukto_handler=dukto_handler, restart=restart, no_quit=no_quit, super_menu=super_menu) + presence.start() + dukto_handler.initialize() + dukto_handler.say_hello() + # bottom right corner screen_geometry = app.primaryScreen().availableGeometry() pet_geometry = pet.frameGeometry() @@ -547,6 +576,7 @@ def main(): pet.show() app.aboutToQuit.connect(presence.end) + app.aboutToQuit.connect(dukto_handler.shutdown) sys.exit(app.exec())