client.py

"""
SMS-STYLE CHAT CLIENT - Real-time bidirectional encrypted chat.
"""

import socket
import struct
import os
import threading
from crypto.engine import CryptoEngine

# Global locks for thread-safe socket operations
send_lock = threading.Lock()
recv_lock = threading.Lock()

def send_message(sock, data):
    """Send length-prefixed message (thread-safe)."""
    with send_lock:
        length = struct.pack('!I', len(data))
        sock.sendall(length + data)

def receive_message(sock):
    """Receive length-prefixed message (thread-safe)."""
    with recv_lock:
        length_data = sock.recv(4)
        if not length_data:
            return None
        length = struct.unpack('!I', length_data)[0]
        data = b''
        while len(data) < length:
            chunk = sock.recv(length - len(data))
            if not chunk:
                return None
            data += chunk
        return data

def main():
    print("MINIMAL DEMO CLIENT")
    print("="*50)
    
    crypto = CryptoEngine()
    
    # Connect to server
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    print("[*] Connecting to localhost:5000...")
    client.connect(('localhost', 5000))
    print("[+] Connected!")
    
    # ECDH Key Exchange
    print("\n[*] Performing ECDH key exchange...")
    my_private, my_public = crypto.generate_ecdh_keypair()
    my_public_bytes = crypto.serialize_public_key(my_public)
    
    # Receive server's public key
    server_public_bytes = receive_message(client)
    server_public = crypto.deserialize_public_key(server_public_bytes)
    print("[*] Received server public key")
    
    # Send my public key
    send_message(client, my_public_bytes)
    print("[*] Sent public key")
    
    # Compute shared secret and derive session key
    shared_secret = crypto.compute_shared_secret(my_private, server_public)
    session_key = crypto.derive_session_key(shared_secret, b"demo_salt", b"demo_session")
    print("[+] Session key established!")
    print(f"[*] Session key: {session_key.hex()[:32]}...")
    
    print("\n" + "="*50)
    print("SMS-STYLE CHAT ACTIVE")
    print("Type anytime - messages appear instantly!")
    print("="*50)
    
    from common.types import EncryptedMessage
    active = True
    
    # Receive thread - listens for incoming messages
    def receive_loop():
        nonlocal active
        while active:
            try:
                encrypted_data = receive_message(client)
                if not encrypted_data:
                    print("\n[*] Connection closed")
                    active = False
                    break
                
                encrypted_msg = EncryptedMessage.deserialize(encrypted_data)
                plaintext_bytes = crypto.decrypt_message(encrypted_msg, session_key)
                plaintext = plaintext_bytes.decode('utf-8')
                
                if plaintext.lower() == 'quit':
                    print("\n[*] Server disconnected")
                    active = False
                    break
                
                print(f"\n[Server]: {plaintext}")
                print("You: ", end="", flush=True)
            except Exception as e:
                if active:
                    print(f"\n[!] Error: {e}")
                active = False
                break
    
    # Start receive thread
    recv_thread = threading.Thread(target=receive_loop, daemon=True)
    recv_thread.start()
    
    # Send loop - main thread handles user input
    print("\nYou: ", end="", flush=True)
    while active:
        try:
            message = input().strip()
            
            if not message:
                print("You: ", end="", flush=True)
                continue
            
            message_bytes = message.encode('utf-8')
            encrypted_msg = crypto.encrypt_message(message_bytes, session_key)
            send_message(client, encrypted_msg.serialize())
            
            if message.lower() == 'quit':
                active = False
                break
            
            print("You: ", end="", flush=True)
        except (KeyboardInterrupt, EOFError):
            print("\n[*] Disconnecting...")
            active = False
            break
        except Exception as e:
            print(f"\n[!] Error: {e}")
            active = False
            break
    
    recv_thread.join(timeout=1)
    
    client.close()
    print("\n[+] Client disconnected")

if __name__ == "__main__":
    main()