A chat application using authenticated encryption. (for reading group!) https://ekzhang.mmm.page/hsrg

aeadchat.py

"""A chat application using authenticated encryption.

Client usage:

    pip3 install pynacl redis rich textual typer
    python3 aeadchat.py <HOST> <PORT>

This connects to a Redis server as the underlying message broker / transport.

---

## (Advanced) How to start the server

For demonstration purposes doing this is enough.

    $ redis-server
    $ bore local 6379 --to bore.pub

Then you can run the chat application connected to bore.pub:$PORT for the
automatically provided port number.

Be warned that this is insecure and exposes your local Redis server to the
world. ONLY RUN THE SERVER IF YOU UNDERSTAND THIS RISK.
"""

from datetime import datetime
from typing import Optional

import nacl.bindings
import nacl.exceptions
import nacl.utils
import redis.asyncio as aioredis
import typer
from nacl.pwhash import argon2i
from rich.text import Text
from rich.tree import Tree
from textual.app import App, ComposeResult
from textual.reactive import reactive
from textual.widget import Widget
from textual.widgets import Input


NONCE_BYTES = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_NPUBBYTES


def encrypt(key: bytes, plaintext: bytes) -> bytes:
    """Encrypt a message using a secret key."""
    nonce = nacl.utils.random(NONCE_BYTES)
    return nonce + nacl.bindings.crypto_aead_xchacha20poly1305_ietf_encrypt(
        plaintext, None, nonce, key
    )


def decrypt(key: bytes, ciphertext: bytes) -> bytes:
    """Decrypt a message using a secret key."""
    nonce, ciphertext = ciphertext[:NONCE_BYTES], ciphertext[NONCE_BYTES:]
    return nacl.bindings.crypto_aead_xchacha20poly1305_ietf_decrypt(
        ciphertext, None, nonce, key
    )


async def send_message(r: aioredis.Redis, key: Optional[bytes], message: str) -> None:
    data = message.encode("utf-8")
    if key:
        data = encrypt(key, data)
    await r.xadd("chat", {"message": data}, maxlen=1000)


async def recv_messages(
    r: aioredis.Redis, key: Optional[bytes], last_id: str
) -> tuple[str, list[tuple[int, str]]]:
    events = await r.xread(streams={"chat": last_id}, count=100, block=0)
    assert len(events) == 1
    assert events[0][0] == b"chat"
    messages = []
    for e_id, e_item in events[0][1]:
        last_id = e_id
        ts = int(e_id.split(b"-")[0])
        data = e_item[b"message"]
        try:
            if key:
                data = decrypt(key, data)
            messages.append((ts, data.decode("utf-8", errors="replace")))
        except nacl.exceptions.CryptoError as exc:
            messages.append((0, "Error: " + str(exc)))
    return last_id, messages


class TextView(Widget):
    messages: reactive[list[tuple[int, str]]] = reactive([])

    def render(self) -> Tree:
        tree = Tree("", hide_root=True)
        for ts, msg in reversed(self.messages):
            entry = Text()
            if ts:
                time_str = datetime.fromtimestamp(ts / 1e3).strftime("[%-I:%M:%S %p]")
                entry.append(time_str, style="cyan")
                entry.append(" " + msg)
            else:
                entry.append(msg, style="red")
            tree.add(entry)
        return tree


class AeadChat(App):
    CSS = """
        Screen {
            layout: vertical;
        }
        .chatarea {
            height: 1fr;
            padding: 0 2;
            border: solid black;
        }
    """

    def __init__(self, r: aioredis.Redis, key: Optional[bytes]) -> None:
        super().__init__()
        self.r = r
        self.key = key
        self.last_id = "$"

    def on_mount(self) -> None:
        self.set_interval(0.01, self.listen)

    def compose(self) -> ComposeResult:
        """Create child widgets for the app."""
        yield TextView(classes="chatarea")
        yield Input(placeholder="Send a message…")

    async def listen(self) -> None:
        self.last_id, incoming = await recv_messages(self.r, self.key, self.last_id)
        view = self.query_one(TextView)
        view.messages = [*view.messages, *incoming]
        view.refresh(layout=True)

    async def on_input_submitted(self, event: Input.Submitted) -> None:
        await send_message(self.r, self.key, event.value)
        event.input.value = ""


def main(
    host: str,
    port: int,
    password: str = typer.Option("", prompt=True, hide_input=True),
) -> None:
    if password:
        key = argon2i.kdf(
            nacl.bindings.crypto_aead_xchacha20poly1305_ietf_KEYBYTES,
            password.encode("utf-8"),
            b"a" * argon2i.SALTBYTES,
            opslimit=argon2i.OPSLIMIT_MIN,
            memlimit=argon2i.MEMLIMIT_MIN,
        )
    else:
        key = None
    r = aioredis.Redis(host=host, port=port)
    AeadChat(r, key).run()


if __name__ == "__main__":
    typer.run(main)