safe_dm.py

#!/usr/bin/env python
"""
Program that demostrates the protocol negotiation explained in https://github.com/nostr-protocol/nips/pull/351

requires aionostr: pip install aionostr
"""
import asyncio
from aionostr.key import PrivateKey, PublicKey
from aionostr.event import Event
from aionostr import Manager
from hashlib import sha256
import secrets
import time
import json



async def main(relays, to_pubkey=None):
    me = PrivateKey()
    print(f"My pubkey {me.public_key.hex()}\n")

    try:
        if to_pubkey:
            print(f"Intiating private DM with {to_pubkey}")
            await request_private_chat(
                me,
                to_pubkey,
                relays=relays,
            )
        else:
            print(f"In another terminal, run: dm.py {me.public_key.hex()}")
            await wait_for_dm_request(me, relays=relays)
    except KeyboardInterrupt:
        return


def make_dm(
    sender_privatekey, recipient_pubkey, content, kind=4, tags=None, expiration=86400
):
    tags = tags or []
    if expiration:
        tags.append(["expiration", str(int(time.time() + expiration))])

    dm = Event(
        pubkey=sender_privatekey.public_key.hex(),
        kind=kind,
        tags=tags,
        content=sender_privatekey.encrypt_message(content, recipient_pubkey),
    )
    dm.sign(sender_privatekey.hex())
    return dm


def make_safe_dm_request(
    sender_privatekey,
    recipient_pubkey,
    relays=["wss://nos.lol"],
    kind=20004,
    token=None,
    request_private_key=None,
    request_dm_kind=4,
    print_request=False,
):
    tags = [["r", relay] for relay in relays]
    tags.append(["p", recipient_pubkey])

    request = Event(
        pubkey=sender_privatekey.public_key.hex(),
        tags=tags,
        content=token or secrets.token_hex(),
        kind=kind,
    )
    request.sign(sender_privatekey.hex())

    if print_request:
        print("\n-------\nContents of request (will be encrypted):")
        print(json.dumps(request.to_json_object(), indent=4))
        print("\n------")

    if not request_private_key:
        tags = [["p", recipient_pubkey]]
    else:
        tags = []

    request_private_key = request_private_key or PrivateKey()
    envelope = make_dm(
        request_private_key,
        recipient_pubkey,
        str(request),
        tags=tags,
        kind=request_dm_kind,
    )

    return envelope, get_conversation_key(request.content, recipient_pubkey)


def get_conversation_key(token, pubkey):
    # the request content is a token used to create a new private key
    return PrivateKey(sha256((token + pubkey).encode()).digest())


def receive_dm_request(recipient_privatekey, envelope):
    try:
        request_data = recipient_privatekey.decrypt_message(
            envelope.content, envelope.pubkey
        )
    except ValueError:
        # print(f"Failed to decrypt from {envelope.pubkey}")
        return
    try:
        request_event = Event(**json.loads(request_data))
    except json.JSONEncoder:
        # not a valid request
        return
    if request_event.verify():
        # this ensures that the request is addressed to me
        if all(request_event.has_tag("p", recipient_privatekey.public_key.hex())):
            conversation_key = get_conversation_key(
                request_event.content, recipient_privatekey.public_key.hex()
            )
            return (
                conversation_key,
                request_event.pubkey,
                request_event.kind,
                [tag[1] for tag in request_event.tags if tag[0] == "r"],
            )


async def send_private_dm(relays, queue, from_private_key, to_pubkey, to_kind):
    print(f"Sender connecting to {relays}")
    async with Manager(relays, private_key=from_private_key.hex()) as manager:
        while True:
            dm_content = await queue.get()
            dm = make_dm(from_private_key, to_pubkey, content=dm_content, kind=to_kind)
            await manager.add_event(dm)


async def receive_private_dm(relays, private_key, receive_pubkey, kind, since=1):
    query = {"kinds": [kind], "limit": 10, "since": since}
    if kind == 4:
        query["#p"] = [private_key.public_key.hex()]
    else:
        query["authors"] = [private_key.public_key.hex()]
    print(f"Waiting for events on {relays} {query}")
    async with Manager(relays, private_key=private_key.hex()) as manager:
        async for event in manager.get_events(query, only_stored=False):
            dm_request = receive_dm_request(private_key, event)
            if dm_request:
                to_key, pub_key, to_kind, to_relays = dm_request
                print(
                    f"Received chat request from {pub_key} -> {to_kind} @ {to_relays}  {to_key.public_key.hex()}"
                )

                if input("Continue conversation? [Y/n] ").strip() != "n":
                    return to_key, pub_key, to_kind, to_relays
            else:
                try:
                    message = (
                        private_key.decrypt_message(event.content, receive_pubkey)
                        .strip()
                        .rjust(80)
                    )
                    # message = f"{event.pubkey[:4]}: {private_key.decrypt_message(event.content, receive_pubkey)}"
                    print(message)
                except ValueError:
                    print(f"Failed to decrypt {event}")
                    continue


async def request_private_chat(
    from_private_key,
    to_public_key,
    relays=["ws://localhost:6969"],
    conversation_kind=20004,
):
    envelope, receive_key = make_safe_dm_request(
        from_private_key,
        to_public_key,
        relays=relays,
        request_dm_kind=4,
        kind=conversation_kind,
    )

    async with Manager(relays, private_key=from_private_key.hex()) as manager:
        print(f"Sending request to {envelope.pubkey} {envelope.kind}")
        await manager.add_event(envelope)

    send_key, _, send_kind, send_relays = await receive_private_dm(
        relays, receive_key, None, conversation_kind, since=envelope.created_at
    )

    await start_private_chat(
        receive_key, send_key, relays, send_relays, conversation_kind, send_kind
    )


async def start_private_chat(
    receive_key, send_key, receive_relays, send_relays, receive_kind, send_kind
):
    send_queue = asyncio.Queue()

    receive_task = asyncio.create_task(
        receive_private_dm(
            receive_relays,
            receive_key,
            send_key.public_key.hex(),
            receive_kind,
            since=int(time.time()) - 1,
        )
    )

    send_task = asyncio.create_task(
        send_private_dm(
            send_relays, send_queue, send_key, receive_key.public_key.hex(), send_kind
        )
    )

    await asyncio.sleep(0.5)
    prompt = Prompt()
    print(
        f"Connected {receive_key.public_key.hex()} to {send_key.public_key.hex()}\n========================\n"
    )
    while True:
        response = await prompt("")
        if response is None:
            break
        elif response:
            await send_queue.put(response)


async def wait_for_dm_request(from_private_key, relays=["wss://nostr.mom"]):
    conversation_kind = 20044

    send_key, to_pubkey, send_kind, send_relays = await receive_private_dm(
        relays, from_private_key, None, 4
    )

    envelope, receive_key = make_safe_dm_request(
        from_private_key,
        send_key.public_key.hex(),
        relays=relays,
        request_dm_kind=send_kind,
        request_private_key=send_key,
        kind=conversation_kind,
    )

    async with Manager(send_relays, private_key=send_key.hex()) as manager:
        print(f"Sending response to {envelope.pubkey} {envelope.kind}")
        await manager.add_event(envelope)

    await start_private_chat(
        receive_key, send_key, relays, send_relays, conversation_kind, send_kind
    )


class Prompt:
    def __init__(self, loop=None):
        import sys

        self.loop = loop or asyncio.get_event_loop()
        self.q = asyncio.Queue()
        self.loop.add_reader(sys.stdin, self.got_input)

    def got_input(self):
        asyncio.ensure_future(self.q.put(sys.stdin.readline()), loop=self.loop)

    async def __call__(self, msg, end="", flush=True):
        print(msg, end=end, flush=flush)
        try:
            return (await self.q.get()).rstrip("\n")
        except:
            return


def demo():
    alice = PrivateKey(
        bytes.fromhex(
            "58483cd4ca00512e24ae55d1938b5e488a6a12cda576dc315793baa5f65642f1"
        )
    )
    alice_pubkey = alice.public_key.hex()
    bob = PrivateKey(
        bytes.fromhex(
            "3b450ab159dfbd9cce9e384a03a54587d92d494c5bf64d1f29c5e02a5ae81ff1"
        )
    )
    bob_pubkey = bob.public_key.hex()

    # alice wants to send to bob
    # alice will listen on wss://relayalice.com
    relays = ["wss://relayalice.com"]
    envelope, bob_alice_key = make_safe_dm_request(alice, bob_pubkey, relays)

    print(f"Alice pubkey: {alice_pubkey}")
    print(f"Bob pubkey: {bob_pubkey}")

    print("\nFirst, Alice sends Bob a DM request:")
    print(json.dumps(envelope.to_json_object(), indent=4))

    print("\nBob receives the DM...")

    bob_alice_key, requestor_pubkey, convo_kind, convo_relays = receive_dm_request(
        bob, envelope
    )
    print(
        f"Alice ({requestor_pubkey}) wants to talk to Bob using {bob_alice_key.public_key.hex()} at {convo_kind} -> {convo_relays}"
    )

    print(f"Bob sends a request to Alice at her chosen location...")
    print(f"Bob would like to use kind=20044 at wss://relaybob.com")
    envelope, alice_bob_key = make_safe_dm_request(
        bob,
        alice.public_key.hex(),
        ["wss://relaybob.com"],
        kind=20044,
        ephemeral_privatekey=bob_alice_key,
        conversation_kind=20004,
        print_request=True,
    )
    print(json.dumps(envelope.to_json_object(), indent=4))

    print(f"\nAlice receives and validates the request at wss://relayalice.com...")
    alice_bob_key, requestor_pubkey, convo_kind2, convo_relays2 = receive_dm_request(
        alice, envelope
    )

    print(
        f"\nBob ({requestor_pubkey}) wants to talk to Alice using {alice_bob_key.public_key.hex()} at {convo_kind2} -> {convo_relays2}"
    )

    print(f"\nFinally, Alice starts the conversation, posting to wss://relaybob.com:")
    print(
        json.dumps(
            make_dm(
                alice_bob_key,
                bob_alice_key.public_key.hex(),
                content="Hello Bob!",
                kind=convo_kind2,
            ).to_json_object(),
            indent=4,
        )
    )


if __name__ == "__main__":
    import sys

    # relays = ["wss://relay-pub.deschooling.us/"]
    relays = ["wss://nostr.mom/"]
    # relays = ["wss://nos.lol/"]
    # relays = ["ws://localhost:6969"]
    if len(sys.argv) >= 2:
        to_pubkey = sys.argv[1]
    else:
        to_pubkey = None
    asyncio.run(main(relays, to_pubkey))