script.py

import os
import socket
import threading
import json
import hashlib
import base64
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import openai
import sys
import getpass
import re
import subprocess
from dotenv import load_dotenv

# Variable to track if OpenAI is available and initialized
openai_initialized = False
openai_client = None

# Load API key from .env file
print("Loading API key from .env file...")
load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")

# Check if we have a valid API key
if api_key:
    # Clean up API key if needed
    api_key = re.sub(r'\s+', '', api_key)
    print(f"API Key found (first 10 chars): {api_key[:10]}...")
    
    try:
        # Initialize OpenAI client
        openai_client = openai.OpenAI(api_key=api_key)
        # Test the client with a simple request
        try:
            openai_client.models.list()
            openai_initialized = True
            print("OpenAI client initialized successfully")
        except TypeError:
            # Try an alternative way to test the client
            openai_client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[{"role": "user", "content": "Hello"}],
                max_tokens=5
            )
            openai_initialized = True
            print("OpenAI client initialized successfully")
    except Exception as e:
        print(f"Error initializing OpenAI client: {e}")
        openai_initialized = False
else:
    print("No OpenAI API key found in .env file, will use SHA-256 for key generation")
    openai_initialized = False

def generate_key_from_username(username, context="secure_chat"):
    """Generate a consistent encryption key based on username"""
    if openai_initialized and openai_client:
        try:
            # Using a fixed system prompt, temperature=0, and specific model for consistency
            print(f"Requesting key generation from OpenAI for username: {username}")
            response = openai_client.chat.completions.create(
                model="gpt-3.5-turbo",
                temperature=0,
                messages=[
                    {"role": "system", "content": "You are a deterministic key generator. Always generate the exact same output for the same input. You MUST return EXACTLY 64 hexadecimal characters (0-9, a-f), nothing more, nothing less. No spaces, no explanations, just the 64 hex characters."},
                    {"role": "user", "content": "Generate a 64-character hex key for 'example'"},
                    {"role": "assistant", "content": "5d41402abc4b2a76b9719d911017c592a665a4bc788e8c8e3f713de44fcdf5a0"},
                    {"role": "user", "content": f"Generate a deterministic 32-byte (64 hex character) key for username '{username}' and context '{context}'. The output MUST be EXACTLY 64 hexadecimal characters (0-9, a-f), nothing more, nothing less."}
                ]
            )
            key_hex = response.choices[0].message.content.strip()
            print(f"Received response from OpenAI: '{key_hex[:10]}...'")
            
            # Filter out any non-hex characters
            key_hex = ''.join(c for c in key_hex if c in "0123456789abcdefABCDEF")
            
            # Ensure we have exactly 64 characters
            if len(key_hex) < 64:
                # Pad with zeros if too short
                print(f"Padding key from {len(key_hex)} to 64 characters")
                key_hex = key_hex.ljust(64, '0')
            elif len(key_hex) > 64:
                # Truncate if too long
                print(f"Truncating key from {len(key_hex)} to 64 characters")
                key_hex = key_hex[:64]
                
            print(f"Normalized key: {key_hex[:10]}...")
            
            # Make sure we have exactly 32 bytes for AES-256
            return bytes.fromhex(key_hex)
        except Exception as e:
            print(f"Error generating key with OpenAI: {e}")
            print("Falling back to SHA-256")
    
    # Fallback to SHA-256 directly
    print("Using SHA-256 for key generation")
    return hashlib.sha256(f"{username}:{context}".encode()).digest()

def encrypt_message(message, key):
    """Encrypt a message using AES-256"""
    # Generate a random 16-byte IV
    iv = os.urandom(16)
    
    # Create an encryptor object
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    
    # Pad the message to be a multiple of 16 bytes (AES block size)
    padded_message = message.encode()
    padding_length = 16 - (len(padded_message) % 16)
    padded_message += bytes([padding_length]) * padding_length
    
    # Encrypt the padded message
    ciphertext = encryptor.update(padded_message) + encryptor.finalize()
    
    # Return the IV and ciphertext as a single base64-encoded string
    return base64.b64encode(iv + ciphertext).decode()

def decrypt_message(encrypted_message, key):
    """Decrypt a message using AES-256"""
    try:
        # Decode the base64-encoded message
        encrypted_data = base64.b64decode(encrypted_message.encode())
        
        # Extract the IV (first 16 bytes) and ciphertext
        iv = encrypted_data[:16]
        ciphertext = encrypted_data[16:]
        
        # Create a decryptor object
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
        decryptor = cipher.decryptor()
        
        # Decrypt the ciphertext
        padded_message = decryptor.update(ciphertext) + decryptor.finalize()
        
        # Remove padding
        padding_length = padded_message[-1]
        message = padded_message[:-padding_length]
        
        return message.decode()
    except Exception as e:
        return f"[Error decrypting message: {e}]"

class SecureChatClient:
    def __init__(self, username, host='127.0.0.1', port=12345, force_sha256=False):
        self.username = username
        self.host = host
        self.port = port
        self.socket = None
        self.connected = False
        self.force_sha256 = force_sha256
        self.using_ai_keys = False
        self.decryption_failures = 0  # Track decryption failures
        
        # Use a consistent shared context for all communication
        self.shared_context = "secure_chat_application"
        
        # Generate encryption key
        if force_sha256:
            print(f"[*] Using SHA-256 for key generation (forced)")
            self.encryption_key = hashlib.sha256(f"{username}:{self.shared_context}".encode()).digest()
        else:
            # Generate key for this user
            print(f"[*] Generating encryption key for '{username}'...")
            self.encryption_key = generate_key_from_username(username, self.shared_context)
            # Check if we're likely using an AI-generated key
            if openai_initialized and not self.force_sha256:
                self.using_ai_keys = True
                print("[*] Using AI-generated keys for encryption/decryption")
            else:
                print("[*] Using SHA-256 keys (AI generation not available)")
        
        # Store keys for other users
        self.peer_keys = {}
    
    def start_server(self):
        """Start a server socket to wait for connections"""
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server_socket.bind((self.host, self.port))
        server_socket.listen(1)
        
        print(f"[*] Waiting for incoming connection on {self.host}:{self.port}")
        self.socket, addr = server_socket.accept()
        print(f"[*] Accepted connection from {addr[0]}:{addr[1]}")
        self.connected = True
        
        # Start a thread to receive messages
        threading.Thread(target=self.receive_messages, daemon=True).start()
        
    def connect_to_server(self, server_host, server_port):
        """Connect to an existing server"""
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        print(f"[*] Connecting to {server_host}:{server_port}")
        self.socket.connect((server_host, server_port))
        print("[*] Connected successfully")
        self.connected = True
        
        # Start a thread to receive messages
        threading.Thread(target=self.receive_messages, daemon=True).start()
    
    def switch_to_sha256(self):
        """Switch to SHA-256 for encryption/decryption"""
        if self.force_sha256:
            print("[*] Already using SHA-256 keys")
            return
            
        self.force_sha256 = True
        self.using_ai_keys = False
        
        # Note: We're now primarily using the conversation key, but we'll update
        # our personal key as well for backwards compatibility
        self.encryption_key = hashlib.sha256(f"{self.username}:{self.shared_context}".encode()).digest()
        
        print(f"[*] ⚠️  SWITCHED TO SHA-256 KEY: {self.encryption_key.hex()[:8]}...")
        print(f"[*] AI-generated keys were not working properly.")
        print(f"[*] Using conversation key: {hashlib.sha256(f'chat:{self.shared_context}'.encode()).digest().hex()[:8]}...")
        print(f"[*] Tell the other person to also type '/sha256' if you're still having issues.")
        
        # Clear peer keys to regenerate them as SHA-256
        self.peer_keys = {}
    
    def generate_peer_key(self, peer_username):
        """Generate a key for the peer based on their username"""
        if peer_username not in self.peer_keys:
            if self.force_sha256:
                # Use SHA-256 if forced
                self.peer_keys[peer_username] = hashlib.sha256(f"{peer_username}:{self.shared_context}".encode()).digest()
                print(f"[*] Generated SHA-256 key for peer '{peer_username}'")
            else:
                # Otherwise use OpenAI
                print(f"[*] Generating encryption key for peer '{peer_username}'...")
                self.peer_keys[peer_username] = generate_key_from_username(peer_username, self.shared_context)
        return self.peer_keys[peer_username]
    
    def send_message(self, message):
        """Encrypt and send a message"""
        if not self.connected:
            print("[!] Not connected")
            return
        
        # For message sending, ALWAYS use a key derived from the CONVERSATION
        # This ensures both sides are using the same key
        conversation_key = hashlib.sha256(f"chat:{self.shared_context}".encode()).digest()
        
        # Encrypt with the conversation key, not our personal key
        encrypted_message = encrypt_message(message, conversation_key)
        packet = json.dumps({
            "sender": self.username,
            "content": encrypted_message,
            "using_conversation_key": True
        })
        self.socket.sendall(f"{packet}\n".encode())
        print(f"[*] Message sent with conversation key: {conversation_key.hex()[:8]}...")
    
    def receive_messages(self):
        """Continuously receive and decrypt messages"""
        buffer = ""
        input_prompt_displayed = False
        
        # The conversation key is constant for the entire chat
        conversation_key = hashlib.sha256(f"chat:{self.shared_context}".encode()).digest()
        
        while self.connected:
            try:
                data = self.socket.recv(4096).decode()
                if not data:
                    print("[!] Connection closed by peer")
                    self.connected = False
                    break
                
                buffer += data
                
                # Process complete messages in the buffer
                while "\n" in buffer:
                    message, buffer = buffer.split("\n", 1)
                    try:
                        packet = json.loads(message)
                        sender = packet["sender"]
                        encrypted_content = packet["content"]
                        
                        print(f"[*] Received message from {sender}")
                        
                        # Always decrypt with the conversation key
                        print(f"[*] Decrypting with conversation key: {conversation_key.hex()[:8]}...")
                        decrypted_content = decrypt_message(encrypted_content, conversation_key)
                        
                        # If that fails, try the old methods as fallback
                        if decrypted_content.startswith("[Error decrypting message:"):
                            print(f"[*] Conversation key failed, trying fallback methods...")
                            
                            # Try with sender's key
                            sender_key = self.generate_peer_key(sender)
                            decrypted_content = decrypt_message(encrypted_content, sender_key)
                            
                            # If that fails, try with our key
                            if decrypted_content.startswith("[Error decrypting message:"):
                                decrypted_content = decrypt_message(encrypted_content, self.encryption_key)
                                
                                # If both fail, try SHA-256 fallback
                                if decrypted_content.startswith("[Error decrypting message:"):
                                    fallback_key = hashlib.sha256(f"{sender}:{self.shared_context}".encode()).digest()
                                    decrypted_content = decrypt_message(encrypted_content, fallback_key)
                        
                        # Clear the current line if user was typing something
                        if not input_prompt_displayed:
                            sys.stdout.write("\r\033[K")  # Clear line
                        else:
                            sys.stdout.write("\n")
                            input_prompt_displayed = False
                        
                        # Print the received message
                        if decrypted_content.startswith("[Error decrypting message:"):
                            print(f"{sender}: [Failed to decrypt message]")
                            self.decryption_failures += 1
                            
                            # Auto-switch to SHA-256 after multiple failures
                            if self.decryption_failures >= 3 and not self.force_sha256:
                                print(f"[!] Multiple decryption failures detected")
                                self.switch_to_sha256()
                        else:
                            print(f"{sender}: {decrypted_content}")
                            # Reset failure counter on success
                            self.decryption_failures = 0
                        
                        # Restore the input prompt
                        sys.stdout.write("You: ")
                        sys.stdout.flush()
                        input_prompt_displayed = True
                    except json.JSONDecodeError:
                        print(f"[!] Received invalid JSON: {message}")
                    except Exception as e:
                        print(f"[!] Error processing message: {e}")
            
            except Exception as e:
                print(f"[!] Error receiving messages: {e}")
                self.connected = False
                break
    
    def close(self):
        """Close the socket connection"""
        if self.socket:
            self.socket.close()
        self.connected = False

def main():
    print("=== Secure Chat with AI-Generated Keys ===")
    print("NOTE: Both users use a shared conversation key for encryption/decryption.")
    print("      Messages are encrypted end-to-end and never stored.")
    
    # Display available commands
    print("\nAvailable commands:")
    print("  /sha256  - Switch to SHA-256 keys (use if messages aren't decrypting)")
    print("  /help    - Show available commands")
    print("  exit     - Exit the application")
    
    username = input("\nEnter your username: ")
    
    # Debug mode for key consistency check
    if username.lower() == "debug":
        debug_key_generation()
        return
    
    # Force SHA-256 mode if requested
    force_sha256 = False
    force_mode = input("Force SHA-256 for all keys? (y/n) [default: n]: ").lower()
    if force_mode == 'y':
        force_sha256 = True
        print("[*] Using SHA-256 keys for all encryption/decryption")
        print("[*] This ensures both users can decrypt each other's messages")
    
    mode = input("(1) Start a new chat or (2) Join existing chat? [1/2]: ")
    
    client = SecureChatClient(username, force_sha256=force_sha256)
    
    if mode == "1":
        client.start_server()
    elif mode == "2":
        # Use a clear variable for server host to avoid concatenation issues
        server_host = input("Enter server IP address [default: 127.0.0.1]: ")
        if not server_host:
            server_host = "127.0.0.1"
        
        # Same for port - parse it carefully
        port_input = input("Enter server port [default: 12345]: ")
        if not port_input:
            server_port = 12345
        else:
            try:
                server_port = int(port_input)
            except ValueError:
                print("Invalid port number, using default 12345")
                server_port = 12345
        
        client.connect_to_server(server_host, server_port)
    else:
        print("Invalid choice")
        return
    
    print("\n[*] Type your messages below ('exit' to quit)")
    print("[*] Using a shared conversation key for encryption/decryption")
    print("[*] Both users should be able to decrypt each other's messages")
    
    if not force_sha256:
        print("[*] If messages aren't being decrypted correctly, type '/sha256' to switch to SHA-256 keys")
        print("[*] The system will automatically switch to SHA-256 after 3 failed decryptions")
    
    try:
        while client.connected:
            message = input("You: ")
            if message.lower() == 'exit':
                break
            elif message.lower() == '/sha256':
                client.switch_to_sha256()
                print("[*] Switched to SHA-256 encryption")
                continue
            elif message.lower() == '/help':
                print("\nAvailable commands:")
                print("  /sha256  - Switch to SHA-256 keys (use if messages aren't decrypting)")
                print("  /help    - Show this help message")
                print("  exit     - Exit the application")
                continue
            client.send_message(message)
    except KeyboardInterrupt:
        print("\n[*] Exiting...")
    finally:
        client.close()

def debug_key_generation():
    """Debug function to check key generation consistency"""
    print("\n=== DEBUG MODE: KEY GENERATION TEST ===")
    test_username = input("Enter a test username: ")
    shared_context = "secure_chat_application"
    
    print("\nGenerating key 10 times to check consistency:")
    keys = []
    for i in range(10):
        key = generate_key_from_username(test_username, shared_context)
        key_hex = key.hex()
        keys.append(key_hex)
        print(f"Key {i+1}: {key_hex[:10]}...")
    
    if all(k == keys[0] for k in keys):
        print("\n✅ SUCCESS: All generated keys are identical!")
    else:
        print("\n❌ ERROR: Keys are not consistent!")
        
    # Also test SHA-256 for comparison
    print("\nSHA-256 key (for reference):")
    sha_key = hashlib.sha256(f"{test_username}:{shared_context}".encode()).digest().hex()
    print(f"SHA key: {sha_key[:10]}...")
    
    return

if __name__ == "__main__":
    main()