feat: sent and received waveform in drybox ui

protocol_prototype/DryBox/UI/client_state.py

@@ -1,6 +1,7 @@
 # client_state.py
 from queue import Queue
 from session import NoiseXKSession
+import time
 
 class ClientState:
     def __init__(self, client_id):

protocol_prototype/DryBox/UI/main.py

@@ -1,4 +1,3 @@
-# main.py
 import sys
 from PyQt5.QtWidgets import (
     QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
@@ -8,13 +7,42 @@ from PyQt5.QtCore import Qt, QSize
 from PyQt5.QtGui import QFont
 from phone_manager import PhoneManager
 from waveform_widget import WaveformWidget
+from phone_state import PhoneState
 
 class PhoneUI(QMainWindow):
     def __init__(self):
         super().__init__()
         self.setWindowTitle("Enhanced Dual Phone Interface")
         self.setGeometry(100, 100, 900, 750)
-        self.setStyleSheet("...")
+        self.setStyleSheet("""
+            QMainWindow { background-color: #333333; }
+            QLabel { color: #E0E0E0; font-size: 14px; }
+            QPushButton {
+                background-color: #0078D4; color: white; border: none;
+                padding: 10px 15px; border-radius: 5px; font-size: 14px;
+                min-height: 30px;
+            }
+            QPushButton:hover { background-color: #005A9E; }
+            QPushButton:pressed { background-color: #003C6B; }
+            QPushButton#settingsButton { background-color: #555555; }
+            QPushButton#settingsButton:hover { background-color: #777777; }
+            QFrame#phoneDisplay {
+                background-color: #1E1E1E; border: 2px solid #0078D4;
+                border-radius: 10px;
+            }
+            QLabel#phoneTitleLabel {
+                font-size: 18px; font-weight: bold; padding-bottom: 5px;
+                color: #FFFFFF;
+            }
+            QLabel#mainTitleLabel {
+                font-size: 24px; font-weight: bold; color: #00A2E8;
+                padding: 15px;
+            }
+            QWidget#phoneWidget {
+                border: 1px solid #4A4A4A; border-radius: 8px;
+                padding: 10px; background-color: #3A3A3A;
+            }
+        """)
 
         self.manager = PhoneManager()
         self.manager.initialize_phones()
@@ -42,11 +70,12 @@ class PhoneUI(QMainWindow):
 
         # Setup UI for phones
         for phone in self.manager.phones:
-            phone_container_widget, phone_button, waveform_widget, phone_status_label = self._create_phone_ui(
+            phone_container_widget, phone_display_frame, phone_button, waveform_widget, sent_waveform_widget, phone_status_label = self._create_phone_ui(
                 f"Phone {phone['id']+1}", lambda checked, pid=phone['id']: self.manager.phone_action(pid, self)
             )
             phone['button'] = phone_button
             phone['waveform'] = waveform_widget
+            phone['sent_waveform'] = sent_waveform_widget
             phone['status_label'] = phone_status_label
             phone_controls_layout.addWidget(phone_container_widget)
             phone['client'].data_received.connect(lambda data, cid=phone['id']: self.manager.update_waveform(cid, data))
@@ -74,7 +103,6 @@ class PhoneUI(QMainWindow):
             self.update_phone_ui(phone['id'])
 
     def _create_phone_ui(self, title, action_slot):
-        # Same as existing _create_phone_ui
         phone_container_widget = QWidget()
         phone_container_widget.setObjectName("phoneWidget")
         phone_layout = QVBoxLayout()
@@ -106,17 +134,25 @@ class PhoneUI(QMainWindow):
         phone_button.clicked.connect(action_slot)
         phone_layout.addWidget(phone_button, alignment=Qt.AlignCenter)
 
-        waveform_label = QLabel(f"{title} Audio")
+        # Received waveform
+        waveform_label = QLabel(f"{title} Received Audio")
         waveform_label.setAlignment(Qt.AlignCenter)
         waveform_label.setStyleSheet("font-size: 14px; color: #E0E0E0;")
         phone_layout.addWidget(waveform_label)
         waveform_widget = WaveformWidget(dynamic=False)
         phone_layout.addWidget(waveform_widget, alignment=Qt.AlignCenter)
 
-        return phone_container_widget, phone_display_frame, phone_button, waveform_widget, phone_status_label
+        # Sent waveform
+        sent_waveform_label = QLabel(f"{title} Sent Audio")
+        sent_waveform_label.setAlignment(Qt.AlignCenter)
+        sent_waveform_label.setStyleSheet("font-size: 14px; color: #E0E0E0;")
+        phone_layout.addWidget(sent_waveform_label)
+        sent_waveform_widget = WaveformWidget(dynamic=False)
+        phone_layout.addWidget(sent_waveform_widget, alignment=Qt.AlignCenter)
+
+        return phone_container_widget, phone_display_frame, phone_button, waveform_widget, sent_waveform_widget, phone_status_label
 
     def update_phone_ui(self, phone_id):
-        """Update phone button and status label based on state."""
         phone = self.manager.phones[phone_id]
         other_phone = self.manager.phones[1 - phone_id]
         state = phone['state']
@@ -146,10 +182,9 @@ class PhoneUI(QMainWindow):
             button.setStyleSheet("background-color: #107C10;")
 
     def set_phone_state(self, client_id, state_str, number):
-        """Handle state changes from client."""
         state = self.manager.map_state(state_str)
         phone = self.manager.phones[client_id]
-        other_phone = self.manager.phones[1 - phone_id]
+        other_phone = self.manager.phones[1 - client_id]
         print(f"Setting state for Phone {client_id + 1}: {state}, number: {number}, is_initiator: {phone['is_initiator']}")
         phone['state'] = state
         if state == PhoneState.IN_CALL:
@@ -162,7 +197,7 @@ class PhoneUI(QMainWindow):
                 print(f"Phone {client_id + 1} (responder) starting handshake")
                 phone['client'].start_handshake(initiator=False, keypair=phone['keypair'], peer_pubkey=other_phone['public_key'])
             elif number == "HANDSHAKE_DONE":
-                self.manager.start_audio(client_id)
+                self.manager.start_audio(client_id, parent=self)  # Pass self as parent
         self.update_phone_ui(client_id)
 
     def settings_action(self):

protocol_prototype/DryBox/UI/phone_client.py

@@ -1,4 +1,3 @@
-# phone_client.py
 import socket
 import time
 import select
@@ -47,10 +46,15 @@ class PhoneClient(QThread):
                     self.state.process_command(self)
                     self.state.check_handshake_timeout(self)
                     if not self.state.handshake_in_progress:
+                        if self.sock is None:
+                            print(f"Client {self.client_id} socket is None, exiting inner loop")
+                            break
                         readable, _, _ = select.select([self.sock], [], [], 0.01)
                         if readable:
                             try:
-                                print(f"Client {self.client_id} attempting sock.recv")
+                                if self.sock is None:
+                                    print(f"Client {self.client_id} socket is None before recv, exiting")
+                                    break
                                 data = self.sock.recv(1024)
                                 if not data:
                                     print(f"Client {self.client_id} disconnected")
@@ -65,9 +69,16 @@ class PhoneClient(QThread):
                         self.msleep(20)
                         print(f"Client {self.client_id} yielding during handshake")
                     self.msleep(1)
+            except Exception as e:
+                print(f"Client {self.client_id} unexpected error in run loop: {e}")
+                self.state_changed.emit("CALL_END", "", self.client_id)
+                break
             finally:
                 if self.sock:
-                    self.sock.close()
+                    try:
+                        self.sock.close()
+                    except Exception as e:
+                        print(f"Client {self.client_id} error closing socket: {e}")
                     self.sock = None
 
     def send(self, message):
@@ -87,8 +98,13 @@ class PhoneClient(QThread):
     def stop(self):
         self.running = False
         if self.sock:
-            self.sock.close()
+            try:
+                self.sock.close()
+            except Exception as e:
+                print(f"Client {self.client_id} error closing socket in stop: {e}")
             self.sock = None
+        self.quit()
+        self.wait(1000)
 
     def start_handshake(self, initiator, keypair, peer_pubkey):
         self.state.start_handshake(initiator, keypair, peer_pubkey)

protocol_prototype/DryBox/UI/phone_manager.py

@@ -1,9 +1,8 @@
-# phone_manager.py
 import secrets
 from PyQt5.QtCore import QTimer
 from phone_client import PhoneClient
 from session import NoiseXKSession
-from phone_state import PhoneState
+from phone_state import PhoneState  # Added import
 
 class PhoneManager:
     def __init__(self):
@@ -11,7 +10,6 @@ class PhoneManager:
         self.handshake_done_count = 0
 
     def initialize_phones(self):
-        """Initialize phone clients and their keypairs."""
         for i in range(2):
             client = PhoneClient(i)
             keypair = NoiseXKSession.generate_keypair()
@@ -27,12 +25,10 @@ class PhoneManager:
             }
             self.phones.append(phone)
 
-        # Share public keys
         self.phones[0]['peer_public_key'] = self.phones[1]['public_key']
         self.phones[1]['peer_public_key'] = self.phones[0]['public_key']
 
     def phone_action(self, phone_id, ui_manager):
-        """Handle phone actions (call, answer, disconnect)."""
         phone = self.phones[phone_id]
         other_phone = self.phones[1 - phone_id]
         print(f"Phone {phone_id + 1} Action, current state: {phone['state']}, is_initiator: {phone['is_initiator']}")
@@ -47,7 +43,7 @@ class PhoneManager:
             phone['state'] = other_phone['state'] = PhoneState.IN_CALL
             phone['client'].send("IN_CALL")
         elif phone['state'] in [PhoneState.IN_CALL, PhoneState.CALLING]:
-            if not phone['client'].handshake_in_progress and phone['state'] != PhoneState.CALLING:
+            if not phone['client'].state.handshake_in_progress and phone['state'] != PhoneState.CALLING:
                 phone['state'] = other_phone['state'] = PhoneState.IDLE
                 phone['client'].send("CALL_END")
                 for p in [phone, other_phone]:
@@ -56,34 +52,47 @@ class PhoneManager:
             else:
                 print(f"Phone {phone_id + 1} cannot hang up during handshake or call setup")
         
-        # Update UI
         ui_manager.update_phone_ui(phone_id)
         ui_manager.update_phone_ui(1 - phone_id)
 
     def send_audio(self, phone_id):
-        """Send mock audio data."""
         phone = self.phones[phone_id]
-        if phone['state'] == PhoneState.IN_CALL and phone['client'].session and phone['client'].sock:
+        if phone['state'] == PhoneState.IN_CALL and phone['client'].state.session and phone['client'].sock:
             mock_audio = secrets.token_bytes(16)
             try:
-                phone['client'].session.send(phone['client'].sock, mock_audio)
+                self.update_sent_waveform(phone_id, mock_audio)
+                phone['client'].state.session.send(phone['client'].sock, mock_audio)
                 print(f"Client {phone_id} sent encrypted audio packet, length=32")
             except Exception as e:
                 print(f"Client {phone_id} failed to send audio: {e}")
 
-    def start_audio(self, client_id):
+    def start_audio(self, client_id, parent=None):
-        """Start audio timer after HANDSHAKE_DONE."""
         self.handshake_done_count += 1
         print(f"HANDSHAKE_DONE received for client {client_id}, count: {self.handshake_done_count}")
         if self.handshake_done_count == 2:
             for phone in self.phones:
                 if phone['state'] == PhoneState.IN_CALL:
                     if not phone['audio_timer'] or not phone['audio_timer'].isActive():
-                        phone['audio_timer'] = QTimer()
+                        phone['audio_timer'] = QTimer(parent)  # Parent to PhoneUI
                         phone['audio_timer'].timeout.connect(lambda pid=phone['id']: self.send_audio(pid))
                         phone['audio_timer'].start(100)
             self.handshake_done_count = 0
 
     def update_waveform(self, client_id, data):
-        """Update waveform with received audio data."""
+        self.phones[client_id]['waveform'].set_data(data)
-        print(f"Updating waveform for client_id {client_id}, data_length={len(data)}")
+
+    def update_sent_waveform(self, client_id, data):
+        self.phones[client_id]['sent_waveform'].set_data(data)
+
+    def map_state(self, state_str):
+        if state_str == "RINGING":
+            return PhoneState.RINGING
+        elif state_str in ["CALL_END", "CALL_DROPPED"]:
+            return PhoneState.IDLE
+        elif state_str == "IN_CALL":
+            return PhoneState.IN_CALL
+        elif state_str == "HANDSHAKE":
+            return PhoneState.IN_CALL
+        elif state_str == "HANDSHAKE_DONE":
+            return PhoneState.IN_CALL
+        return PhoneState.IDLE

protocol_prototype/DryBox/UI/session.py

@@ -126,7 +126,7 @@ class NoiseXKSession:
             raise RuntimeError("Handshake not complete")
         ct = self._send_cs.encrypt_with_ad(b'', plaintext)
         logging.debug(f"[ENCRYPT] {ct.hex()}")
-        self._dump_cipherstate("SEND→ after encrypt", self._send_cs)
+        # self._dump_cipherstate("SEND→ after encrypt", self._send_cs)
         self._send_all(sock, ct)
 
     def receive(self, sock: socket.socket) -> bytes:
@@ -137,7 +137,7 @@ class NoiseXKSession:
             raise RuntimeError("Handshake not complete")
         ct = self._recv_all(sock)
         logging.debug(f"[CIPHERTEXT] {ct.hex()}")
-        self._dump_cipherstate("RECV→ before decrypt", self._recv_cs)
+        # self._dump_cipherstate("RECV→ before decrypt", self._recv_cs)
         pt = self._recv_cs.decrypt_with_ad(b'', ct)
         logging.debug(f"[DECRYPT] {pt!r}")
         return pt
@@ -153,7 +153,7 @@ class NoiseXKSession:
             logging.debug(f"[DECRYPT] Stripping 2-byte length prefix from {len(ciphertext)}-byte input")
             ciphertext = ciphertext[2:]
         logging.debug(f"[CIPHERTEXT] {ciphertext.hex()}")
-        self._dump_cipherstate("DECRYPT→ before decrypt", self._recv_cs)
+        # self._dump_cipherstate("DECRYPT→ before decrypt", self._recv_cs)
         pt = self._recv_cs.decrypt_with_ad(b'', ciphertext)
         logging.debug(f"[DECRYPT] {pt!r}")
         return pt