lan_chat.py

#!/usr/bin/env python3

import argparse
import socket
import threading
from datetime import datetime

DEFAULT_HOST = "0.0.0.0"
DEFAULT_PORT = 5000
BUFFER_SIZE = 4096
ENCODING = "utf-8"


def timestamp():
    return datetime.now().strftime("%H:%M:%S")


def send_line(sock, text):
    sock.sendall((text + "\n").encode(ENCODING))


def receive_lines(sock):
    buffer = b""

    while True:
        data = sock.recv(BUFFER_SIZE)

        if not data:
            return

        buffer += data

        while b"\n" in buffer:
            line, buffer = buffer.split(b"\n", 1)
            yield line.decode(ENCODING, errors="replace")


class ChatServer:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.clients = {}
        self.lock = threading.Lock()

    def broadcast(self, message, exclude=None):
        with self.lock:
            clients = list(self.clients.items())

        disconnected = []

        for sock, username in clients:
            if sock is exclude:
                continue

            try:
                send_line(sock, message)
            except OSError:
                disconnected.append(sock)

        for sock in disconnected:
            self.remove_client(sock)

    def remove_client(self, sock):
        with self.lock:
            username = self.clients.pop(sock, None)

        try:
            sock.close()
        except OSError:
            pass

        return username

    def handle_client(self, sock, address):
        username = None

        try:
            lines = receive_lines(sock)

            username = next(lines).strip()[:32]

            if not username:
                send_line(sock, "ERROR|Username cannot be empty.")
                return

            with self.lock:
                existing_names = {
                    name.lower() for name in self.clients.values()
                }

                if username.lower() in existing_names:
                    send_line(
                        sock,
                        "ERROR|That username is already in use."
                    )
                    return

                self.clients[sock] = username

            send_line(sock, f"SYSTEM|Connected as {username}.")
            self.broadcast(
                f"SYSTEM|{username} joined the chat.",
                exclude=sock
            )

            print(
                f"[+] {username} connected from "
                f"{address[0]}:{address[1]}"
            )

            for message in lines:
                message = message.rstrip("\r")

                if not message:
                    continue

                message = message[:1000]

                formatted = f"MSG|{username}|{message}"

                print(
                    f"[{timestamp()}] "
                    f"{username}: {message}"
                )

                self.broadcast(formatted)

        except (ConnectionError, OSError, StopIteration):
            pass

        finally:
            old_username = self.remove_client(sock)

            if old_username:
                self.broadcast(
                    f"SYSTEM|{old_username} left the chat."
                )

                print(f"[-] {old_username} disconnected")

    def run(self):
        with socket.socket(
            socket.AF_INET,
            socket.SOCK_STREAM
        ) as server:

            server.setsockopt(
                socket.SOL_SOCKET,
                socket.SO_REUSEADDR,
                1
            )

            server.bind((self.host, self.port))
            server.listen()

            try:
                local_ip = socket.gethostbyname(
                    socket.gethostname()
                )
            except socket.gaierror:
                local_ip = "<your LAN IP>"

            print("LocalChat server started")
            print(f"Listening on {local_ip}:{self.port}")
            print()
            print(
                "Other devices must be connected "
                "to the same LAN/Wi-Fi."
            )
            print("Press Ctrl+C to stop.")
            print()

            try:
                while True:
                    client, address = server.accept()

                    thread = threading.Thread(
                        target=self.handle_client,
                        args=(client, address),
                        daemon=True
                    )

                    thread.start()

            except KeyboardInterrupt:
                print("\nServer stopped.")


class ChatClient:
    def __init__(self, host, port, username):
        self.host = host
        self.port = port
        self.username = username
        self.running = True

    def receive_loop(self, sock):
        try:
            for line in receive_lines(sock):

                if line.startswith("MSG|"):
                    _, username, message = line.split(
                        "|",
                        2
                    )

                    print(
                        f"\r[{timestamp()}] "
                        f"{username}: {message}"
                    )

                    print("> ", end="", flush=True)

                elif line.startswith("SYSTEM|"):
                    print(
                        f"\r* {line[7:]}"
                    )

                    print("> ", end="", flush=True)

                elif line.startswith("ERROR|"):
                    print(
                        f"\rERROR: {line[6:]}"
                    )

                    self.running = False
                    return

        except (ConnectionError, OSError):
            pass

        finally:
            if self.running:
                print("\nConnection to server lost.")

            self.running = False

    def run(self):
        try:
            sock = socket.create_connection(
                (self.host, self.port),
                timeout=10
            )

            sock.settimeout(None)

        except OSError as error:
            print(
                f"Could not connect to "
                f"{self.host}:{self.port}"
            )

            print(f"Error: {error}")
            return

        with sock:
            send_line(sock, self.username)

            receiver = threading.Thread(
                target=self.receive_loop,
                args=(sock,),
                daemon=True
            )

            receiver.start()

            print(
                f"Connected to "
                f"{self.host}:{self.port} "
                f"as {self.username}"
            )

            print(
                "Type a message and press Enter."
            )

            print(
                "Type /quit to leave."
            )

            print()

            try:
                while self.running:
                    message = input("> ").strip()

                    if message == "/quit":
                        break

                    if message:
                        try:
                            send_line(
                                sock,
                                message[:1000]
                            )
                        except OSError:
                            break

            except (KeyboardInterrupt, EOFError):
                pass

            finally:
                self.running = False

        print("Disconnected.")


def main():
    parser = argparse.ArgumentParser(
        description="Chat with people on your local network."
    )

    subparsers = parser.add_subparsers(
        dest="mode",
        required=True
    )

    server_parser = subparsers.add_parser(
        "server",
        help="Start a LocalChat server"
    )

    server_parser.add_argument(
        "--host",
        default=DEFAULT_HOST
    )

    server_parser.add_argument(
        "--port",
        type=int,
        default=DEFAULT_PORT
    )

    client_parser = subparsers.add_parser(
        "join",
        help="Join a LocalChat server"
    )

    client_parser.add_argument(
        "host",
        help="LAN IP address of the server"
    )

    client_parser.add_argument(
        "--port",
        type=int,
        default=DEFAULT_PORT
    )

    client_parser.add_argument(
        "--name",
        help="Your username"
    )

    args = parser.parse_args()

    if args.mode == "server":
        ChatServer(
            args.host,
            args.port
        ).run()

    else:
        username = args.name

        if not username:
            username = input("Username: ").strip()

        if not username:
            parser.error(
                "Username cannot be empty"
            )

        ChatClient(
            args.host,
            args.port,
            username
        ).run()


if __name__ == "__main__":
    main()