server.py

#!/usr/bin/env python3
"""
NETwasel - Deliverable #4
Server
"""

import socket
import threading
import argparse
import sys
import ssl
import os
import json
import hashlib

USERS_DB = os.path.join(os.getcwd(), "users.json")


class ChatServer:
    """
    A simple TCP chat server for NETwasel.
    - Accepts client connections
    - Receives messages
    - Broadcasts them to all connected clients
    """

    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self.server_socket = None
        self.clients = {}  # {socket: username}
        self.usernames = {}  # {username: socket}
        self.display_names = {}  # {display_name: username}
        self.name_index = {}  # {display_name: list of usernames}
        self.lock = threading.Lock()  # protect shared dict

    def start(self):
        """Start the TLS-secured chat server and begin listening for connections."""
        # Spawn a thread for the client
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        try:
            self.server_socket.bind((self.host, self.port))
        except OSError as e:
            print(f"[ERROR] Could not bind to {self.host}:{self.port} -> {e}")
            sys.exit(1)

        self.server_socket.listen(5)

        context = self.create_tls_context()
        print(f"[INFO] NETwasel TLS server running on {self.host}:{self.port}")

        # Accept clients in loop
        try:
            while True:
                client_conn, client_addr = self.server_socket.accept()
                # Wrap client socket in TLS
                try:
                    tls_client = context.wrap_socket(client_conn, server_side=True)
                    print(f"[INFO] TLS connection from {client_addr}")

                    # Spawn a thread for the client
                    thread = threading.Thread(
                        target=self.handle_client, args=(tls_client,)
                    )
                    thread.daemon = True
                    thread.start()
                except ssl.SSLError as e:
                    print(f"[ERROR] TLS handshake failed: {e}")
                    client_conn.close()
        except KeyboardInterrupt:
            print("\n[INFO] Server shutting down...")
        finally:
            self.shutdown()

    def create_tls_context(self):
        """Sets up a modern TLS context using the server's cert and key."""
        cert_file = os.path.join(os.getcwd(), "NETwasel.crt")
        key_file = os.path.join(os.getcwd(), "NETwasel.key")

        context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        context.minimum_version = ssl.TLSVersion.TLSv1_2  # Enforce modern TLS
        context.check_hostname = False
        context.verify_mode = ssl.CERT_NONE
        context.load_cert_chain(certfile=cert_file, keyfile=key_file)
        return context

    def send_key_to_client(self, client_sock, pubkey):
        try:
            if isinstance(pubkey, str):
                pubkey_bytes = pubkey.encode("utf-8")
            else:
                pubkey_bytes = pubkey
            client_sock.sendall(pubkey_bytes)
            client_sock.sendall(b"ENDKEY\n")
        except Exception as e:
            print(f"[ERROR] Failed to send key: {e}")
            try:
                client_sock.sendall(b"ERROR Failed to send key\nENDKEY\n")
            except:
                pass

    
    def _drain_file_data(self, sock, size, max_drain=10_000_000):
        """Safely drain a known or capped amount of incoming file data."""
        to_drain = min(size, max_drain)
        drained = 0
        try:
            sock.settimeout(0.2)
            while drained < to_drain:
                chunk = sock.recv(min(1024, to_drain - drained))
                if not chunk:
                    break
                drained += len(chunk)
        except (socket.timeout, ssl.SSLError):
            pass
        finally:
            sock.settimeout(None)

    def handle_client(self, client_sock: socket.socket):
        """Handle communication with a connected client."""
        try:
            # First message must be LOGIN or REGISTER
            login_msg = client_sock.recv(1024).decode("utf-8").strip()
            parts = login_msg.split(" ", 2)

            if len(parts) != 3 or parts[0] not in ("LOGIN", "REGISTER"):
                client_sock.sendall(
                    b"ERROR Expected LOGIN or REGISTER <username> <password>\n"
                )
                client_sock.close()
                return

            cmd, username, password = parts
            users = load_users()

            if cmd == "LOGIN":
                if username not in users:
                    client_sock.sendall(b"ERROR Invalid username or password\n")
                    client_sock.close()
                    return
                stored = users[username]
                if hash_password(password, stored["salt"]) != stored["hash"]:
                    client_sock.sendall(b"ERROR Invalid username or password\n")
                    client_sock.close()
                    return
                client_sock.sendall(b"OK\n")

            elif cmd == "REGISTER":
                if username in users:
                    client_sock.sendall(b"ERROR Username already exists\n")
                    client_sock.close()
                    return
                users[username] = create_user(password)
                save_users(users)
                client_sock.sendall(b"OK\n")

            # Proceed to receive HELLO <name>
            name_msg = client_sock.recv(1024).decode("utf-8").strip()
            if not name_msg.startswith("HELLO"):
                client_sock.sendall(b"Expected HELLO <name>\n")
                client_sock.close()
                return

            name = name_msg[6:].strip() or username

            with self.lock:
                self.clients[client_sock] = username
                self.usernames[username] = client_sock
                self.display_names[username] = name
                self.name_index.setdefault(name, []).append(username)

            welcome = f"[INFO] {name} has joined NETwasel\n"
            self.broadcast(welcome, exclude=client_sock)
            client_sock.sendall(b"[INFO] Welcome to NETwasel chat!\n")


            buffer = b""
            # Main receive loop
            while True:
                while b"\n" not in buffer:
                    try:
                        chunk = client_sock.recv(1024)
                        if not chunk:
                            print("[INFO] Client disconnected cleanly.")
                            self.remove_client(client_sock)
                            return
                        buffer += chunk
                    except ssl.SSLError as e:
                        print(f"[SSL ERROR] {e}")
                        self.remove_client(client_sock)
                        return
                    except ConnectionResetError:
                        print("[INFO] Client forcibly closed the connection.")
                        self.remove_client(client_sock)
                        return
                    except Exception as e:
                        print(f"[ERROR] Unexpected server error: {e}")
                        self.remove_client(client_sock)
                        return
                
                line, buffer = buffer.split(b"\n", 1)
                msg = line.decode("utf-8", errors="ignore").strip()
                
                # Handles private messaging
                if msg.startswith("PRIVMSG "):
                    parts = msg.split(" ", 2)
                    if len(parts) < 3:
                        client_sock.sendall(b"ERROR Invalid private message format\n")
                        continue

                    target_user, message_body = parts[1], parts[2]

                    with self.lock:
                        target_sock = self.usernames.get(target_user)
                        sender_user = self.clients.get(client_sock, "unknown")
                        sender_name = self.display_names.get(sender_user, sender_user)

                    if target_sock:
                        try:
                            # Send to recipient
                            target_sock.sendall(
                                f"[PM from {sender_name}]: {message_body}\n".encode(
                                    "utf-8"
                                )
                            )
                            # Confirm to sender
                            client_sock.sendall(
                                f"[PM to {target_user}]: {message_body}\n".encode(
                                    "utf-8"
                                )
                            )
                        except:
                            client_sock.sendall(b"ERROR Failed to deliver message\n")
                    else:
                        client_sock.sendall(b"ERROR User not online\n")
                    continue

                # Handles encrypted private messaging
                if msg.startswith("ENCMSG "):
                    parts = msg.split(" ", 2)
                    if len(parts) < 3:
                        client_sock.sendall(b"ERROR Invalid ENCMSG format\n")
                        continue

                    target_user, b64cipher = parts[1], parts[2]
                    with self.lock:
                        target_sock = self.usernames.get(target_user)
                        sender_user = self.clients.get(client_sock, "unknown")
                        sender_name = self.display_names.get(sender_user, sender_user)

                    if not target_sock:
                        client_sock.sendall(b"ERROR Recipient not online\n")
                        continue

                    try:
                        forward_msg = f"ENCMSG {sender_user} {b64cipher}\n"
                        target_sock.sendall(f"[ENCRYPTED PM from {sender_name}]\n".encode("utf-8"))
                        target_sock.sendall(forward_msg.encode("utf-8"))
                    except:
                        client_sock.sendall(b"ERROR Failed to send encrypted message\n")
                    continue

                # Handles /whois command
                if msg.startswith("/whois "):
                    display = msg.split(" ", 1)[1].strip()
                    with self.lock:
                        matches = self.name_index.get(display, [])
                    if matches:
                        client_sock.sendall(
                            f"Users with name '{display}': {', '.join(matches)}\n".encode(
                                "utf-8"
                            )
                        )
                    else:
                        client_sock.sendall(
                            f"No users found with name '{display}'\n".encode("utf-8")
                        )
                    continue

                # Handles private/group file transfer
                if msg.startswith("SEND_FILE ") or msg.startswith("SEND_ALL "):
                    parts = msg.split(" ")
                    broadcast = msg.startswith("SEND_ALL")

                    try:
                        if broadcast:
                            if len(parts) != 3:
                                raise ValueError("Invalid SEND_ALL format")
                            filename = parts[1]
                            filesize = int(parts[2])
                            recipients = [
                                sock for sock in self.clients if sock != client_sock
                            ]
                        else:
                            if len(parts) != 4:
                                raise ValueError("Invalid SEND_FILE format")
                            target_user = parts[1]
                            filename = parts[2]
                            filesize = int(parts[3])
                            with self.lock:
                                target_sock = self.usernames.get(target_user)
                            if not target_sock:
                                client_sock.sendall(b"ERROR Target user not online\n")
                                self._drain_file_data(client_sock, filesize)
                                client_sock.sendall(b"SYNC\n")
                                continue
                            recipients = [target_sock]
                    except ValueError:
                        client_sock.sendall(b"ERROR Invalid file header\n")
                        self._drain_file_data(client_sock, 4096)
                        client_sock.sendall(b"SYNC\n")
                        continue

                    # Notify recipients
                    sender_user = self.clients.get(client_sock, "unknown")
                    sender_name = self.display_names.get(sender_user, sender_user)
                    header = f"[FILE from {sender_name}]: {filename} ({filesize} bytes)\n".encode(
                        "utf-8"
                    )

                    for r in recipients:
                        try:
                            r.sendall(header)
                            r.sendall(
                                f"SEND_FILE {filename} {filesize}\n".encode("utf-8")
                            )
                        except:
                            pass

                    # Relay binary file
                    remaining = filesize
                    try:
                        while remaining > 0:
                            chunk = client_sock.recv(min(1024, remaining))
                            if not chunk:
                                break
                            for r in recipients:
                                try:
                                    r.sendall(chunk)
                                except:
                                    pass
                            remaining -= len(chunk)
                        client_sock.sendall(b"[INFO] File sent successfully\n")
                    except Exception as e:
                        client_sock.sendall(
                            f"ERROR File transfer failed: {e}\n".encode("utf-8")
                        )
                    continue

                # Handles encrypted private file transfer
                if msg.startswith("ENCFILE "):
                    try:
                        parts = msg.split(" ", 5)
                        if len(parts) != 5:
                            client_sock.sendall(b"ERROR Invalid ENCFILE header\n")
                            continue

                        target_user = parts[1]
                        filename = parts[2]
                        key_len = int(parts[3])
                        file_len = int(parts[4])
                        total_len = key_len + file_len

                        with self.lock:
                            target_sock = self.usernames.get(target_user)
                            sender_user = self.clients.get(client_sock, "unknown")
                            sender_name = self.display_names.get(
                                sender_user, sender_user
                            )

                        if not target_sock:
                            client_sock.sendall(b"ERROR Recipient not online\n")
                            self._drain_file_data(client_sock, total_len)
                            continue

                        # Notify recipient
                        header = f"[ENCRYPTED FILE from {sender_name}]: {filename} (encrypted)\n".encode(
                            "utf-8"
                        )
                        relay_cmd = f"ENCFILE {sender_user} {filename} {key_len} {file_len}\n".encode(
                            "utf-8"
                        )

                        target_sock.sendall(header)
                        target_sock.sendall(relay_cmd)

                        # Relay key + file
                        remaining = total_len
                        while remaining > 0:
                            chunk = client_sock.recv(min(1024, remaining))
                            if not chunk:
                                break
                            target_sock.sendall(chunk)
                            remaining -= len(chunk)

                        client_sock.sendall(
                            b"[INFO] Encrypted file relayed successfully\n"
                        )

                    except Exception as e:
                        client_sock.sendall(
                            f"ERROR Failed to relay encrypted file: {e}\n".encode(
                                "utf-8"
                            )
                        )
                    continue

                # Handles key uploads
                if msg.startswith("UPLOADKEY "):
                    parts = msg.split()
                    if len(parts) != 2:
                        client_sock.sendall(b"ERROR Invalid UPLOADKEY format\n")
                        continue

                    target_user = parts[1]
                    buffer = []

                    while True:
                        line = client_sock.recv(1024).decode("utf-8")
                        if line.strip() == "ENDKEY":
                            break
                        buffer.append(line)

                    public_key_pem = "".join(buffer)

                    users = load_users()
                    if target_user in users:
                        if "public_key" in users[target_user]:
                            client_sock.sendall(b"ERROR Public key already exists and cannot be overwritten\n")
                        else:
                            users[target_user]["public_key"] = public_key_pem
                            save_users(users)
                            client_sock.sendall(b"[INFO] Public key saved\n")
                    else:
                        client_sock.sendall(b"ERROR User not found\n")
                    continue

                # Handles user key's retrieval
                if msg.startswith("GETKEY "):
                    parts = msg.split()
                    if len(parts) != 2:
                        try:
                            client_sock.sendall(b"ERROR Invalid GETKEY format\n")
                        except Exception as e:
                            print(f"[ERROR] Failed to send error response: {e}")
                        continue

                    lookup_user = parts[1]
                    users = load_users()
                    user_entry = users.get(lookup_user)

                    try:
                        if user_entry and "public_key" in user_entry:
                            pubkey = user_entry["public_key"]
                            if isinstance(pubkey, str):
                                pubkey_bytes = pubkey.encode("utf-8")
                            else:
                                pubkey_bytes = pubkey

                            # Construct structured response
                            header = f"KEYDATA {lookup_user} {len(pubkey_bytes)}\n".encode("utf-8")
                            client_sock.sendall(header)
                            client_sock.sendall(pubkey_bytes)
                        else:
                            client_sock.sendall(b"ERROR Public key not found\n")
                    except Exception as e:
                        print(f"[ERROR] Failed during GETKEY send: {e}")
                    continue


                # Handles quitting
                if msg.lower() == "/quit":
                    break

                full_msg = f"{name}: {msg}\n"
                self.broadcast(full_msg, exclude=None)

        except ConnectionResetError:
            pass
        finally:
            self.remove_client(client_sock)

    def broadcast(self, message: str, exclude=None):
        """Send a message to all connected clients (optionally excluding one)."""
        with self.lock:
            for sock, _ in self.clients.items():
                if sock == exclude:
                    continue
                try:
                    sock.sendall(message.encode("utf-8"))
                except (BrokenPipeError, ConnectionResetError):
                    self.remove_client(sock)

    def remove_client(self, client_sock: socket.socket):
        """Remove a disconnected client from the dictionary."""
        with self.lock:
            name = self.clients.pop(client_sock, "Unknown")
        print(f"[INFO] {name} disconnected")
        client_sock.close()
        leave_msg = f"[INFO] {name} has left NETwasel\n"
        self.broadcast(leave_msg, exclude=None)

    def shutdown(self):
        """Shut down the server and close all sockets."""
        with self.lock:
            for sock in self.clients.keys():
                sock.close()
            self.clients.clear()
        if self.server_socket:
            self.server_socket.close()


def load_users():
    if os.path.exists(USERS_DB):
        with open(USERS_DB, "r") as f:
            return json.load(f)
    return {}


def save_users(users):
    with open(USERS_DB, "w") as f:
        os.umask(0o077)
        json.dump(users, f, indent=2)
    os.chmod(USERS_DB, 0o600)


def hash_password(password, salt_hex):
    salt = bytes.fromhex(salt_hex)
    return hashlib.pbkdf2_hmac("sha256", password.encode(), salt, 100_000).hex()


def create_user(password):
    salt = os.urandom(16).hex()
    return {"salt": salt, "hash": hash_password(password, salt)}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="NETwasel Chat Server")
    parser.add_argument(
        "--host", default="0.0.0.0", help="Host to bind to (default: 0.0.0.0)"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=5050,
        help="Port to bind to (5000-5100 required on stu)",
    )
    args = parser.parse_args()

    # Enforce allowed port range
    if not (5000 <= args.port <= 5100):
        print("[ERROR] Port must be in range 5000–5100 on stu.cs.jmu.edu")
        sys.exit(1)

    server = ChatServer(args.host, args.port)
    server.start()