client.py

import socket
import threading
import ssl
import os
import time
from protocol import send_json, recv_json
from encryption_utils import encrypt_message, decrypt_message, set_key

HOST = "10.207.26.156"
PORT = 5000

def receive_messages(sock):
    buffer = ""

    while True:
        try:
            messages, buffer = recv_json(sock, buffer)

            for data in messages:
                msg_type = data.get("type")

                if msg_type == "message":
                    print(f"[{data['sender']}]: {data['content']}")

                elif msg_type == "private":
                    print(f"[PRIVATE from {data['from']}]: {data['content']}")

                elif msg_type == "secret":
                    decrypted = decrypt_message(data["content"])
                    print(f"[SECRET from {data['from']}]: {decrypted}")

                elif msg_type == "info":
                    print(f"[INFO]: {data['message']}")

                elif msg_type == "file":
                    os.makedirs("received_files", exist_ok=True)

                    filename = "received_files/" + data["filename"]
                    filesize = data["filesize"]
                    remaining = filesize

                    with open(filename, "wb") as f:
                        while remaining > 0:
                            chunk = sock.recv(min(4096, remaining))
                            f.write(chunk)
                            remaining -= len(chunk)

                    print(f"[FILE RECEIVED]: {filename}")

        except Exception as e:
            print("Receive error:", e)
            break

def start_client():
    key = b'oGk8-xGWT0PhL0ek9rhQ1xgGjkVaawX8EhAHfy5bHmQ='
    set_key(key)

    context = ssl.create_default_context()
    context.check_hostname = False
    context.verify_mode = ssl.CERT_NONE

    raw_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    secure_sock = context.wrap_socket(raw_sock)

    secure_sock.connect((HOST, PORT))
    print("Connected with cipher:", secure_sock.cipher())

    threading.Thread(target=receive_messages, args=(secure_sock,), daemon=True).start()

    username = input("Enter username: ")
    room = input("Enter room: ")

    send_json(secure_sock, {
        "type": "join",
        "username": username,
        "room": room
    })

    while True:
        msg = input()

        if msg == "/exit":
            secure_sock.close()
            break

        elif msg.startswith("/private"):
            parts = msg.split(" ", 2)
            if len(parts) < 3:
                print("Usage: /private user message")
                continue

            send_json(secure_sock, {
                "type": "private",
                "to": parts[1],
                "content": parts[2]
            })

        elif msg.startswith("/secret"):
            parts = msg.split(" ", 2)
            if len(parts) < 3:
                print("Usage: /secret user message")
                continue

            encrypted = encrypt_message(parts[2])

            send_json(secure_sock, {
                "type": "secret",
                "to": parts[1],
                "content": encrypted
            })

        elif msg.startswith("/file"):
            parts = msg.split(" ", 1)
            if len(parts) < 2 or not os.path.exists(parts[1]):
                print("Invalid file")
                continue

            filepath = parts[1]
            filesize = os.path.getsize(filepath)
            filename = os.path.basename(filepath)

            send_json(secure_sock, {
                "type": "file",
                "filename": filename,
                "filesize": filesize,
                "room": room
            })

            time.sleep(0.1)

            with open(filepath, "rb") as f:
                while chunk := f.read(4096):
                    secure_sock.sendall(chunk)

        else:
            send_json(secure_sock, {
                "type": "message",
                "room": room,
                "content": msg
            })

if __name__ == "__main__":
    start_client()