add EncryptedDirectMessage class; simplify Event class (#39)

Author: kdmukai

README.md

@@ -11,6 +11,7 @@ public_key = private_key.public_key
 print(f"Private key: {private_key.bech32()}")
 print(f"Public key: {public_key.bech32()}")
 ```
+
 **Connect to relays**
 ```python
 import json
@@ -30,6 +31,7 @@ while relay_manager.message_pool.has_notices():
 
 relay_manager.close_connections()
 ```
+
 **Publish to relays**
 ```python
 import json 
@@ -48,14 +50,46 @@ time.sleep(1.25) # allow the connections to open
 
 private_key = PrivateKey()
 
-event = Event(private_key.public_key.hex(), "Hello Nostr")
+event = Event("Hello Nostr")
 private_key.sign_event(event)
 
 relay_manager.publish_event(event)
 time.sleep(1) # allow the messages to send
 
 relay_manager.close_connections()
 ```
+
+**Reply to a note**
+```python
+from nostr.event import Event
+
+reply = Event(
+  content="Hey, that's a great point!",
+)
+
+# create 'e' tag reference to the note you're replying to
+reply.add_event_ref(original_note_id)
+
+# create 'p' tag reference to the pubkey you're replying to
+reply.add_pubkey_ref(original_note_author_pubkey)
+
+private_key.sign_event(reply)
+relay_manager.publish_event(reply)
+```
+
+**Send a DM**
+```python
+from nostr.event import EncryptedDirectMessage
+
+dm = EncryptedDirectMessage(
+  recipient_pubkey=recipient_pubkey,
+  cleartext_content="Secret message!"
+)
+private_key.sign_event(dm)
+relay_manager.publish_event(dm)
+```
+
+
 **Receive events from relays**
 ```python
 import json
@@ -112,7 +146,6 @@ delegation = Delegation(
 identity_pk.sign_delegation(delegation)
 
 event = Event(
-    delegatee_pk.public_key.hex(),
     "Hello, NIP-26!",
     tags=[delegation.get_tag()],
 )

nostr/event.py

@@ -1,12 +1,15 @@
 import time
 import json
+from dataclasses import dataclass, field
 from enum import IntEnum
+from typing import List
 from secp256k1 import PrivateKey, PublicKey
 from hashlib import sha256
 
 from nostr.message_type import ClientMessageType
 
 
+
 class EventKind(IntEnum):
     SET_METADATA = 0
     TEXT_NOTE = 1
@@ -16,41 +19,58 @@ class EventKind(IntEnum):
     DELETE = 5
 
 
-class Event():
-    def __init__(
-            self, 
-            public_key: str, 
-            content: str, 
-            created_at: int = None, 
-            kind: int=EventKind.TEXT_NOTE, 
-            tags: "list[list[str]]"=[], 
-            id: str=None, 
-            signature: str=None) -> None:
-        if not isinstance(content, str):
+
+@dataclass
+class Event:
+    content: str = None
+    public_key: str = None
+    created_at: int = None
+    kind: int = EventKind.TEXT_NOTE
+    tags: List[List[str]] = field(default_factory=list)  # Dataclasses require special handling when the default value is a mutable type
+    signature: str = None
+
+
+    def __post_init__(self):
+        if self.content is not None and not isinstance(self.content, str):
+            # DMs initialize content to None but all other kinds should pass in a str
             raise TypeError("Argument 'content' must be of type str")
 
-        self.public_key = public_key
-        self.content = content
-        self.created_at = created_at or int(time.time())
-        self.kind = kind
-        self.tags = tags
-        self.signature = signature
-        self.id = id or Event.compute_id(self.public_key, self.created_at, self.kind, self.tags, self.content)
+        if self.created_at is None:
+            self.created_at = int(time.time())
+
 
     @staticmethod
-    def serialize(public_key: str, created_at: int, kind: int, tags: "list[list[str]]", content: str) -> bytes:
+    def serialize(public_key: str, created_at: int, kind: int, tags: List[List[str]], content: str) -> bytes:
         data = [0, public_key, created_at, kind, tags, content]
         data_str = json.dumps(data, separators=(',', ':'), ensure_ascii=False)
         return data_str.encode()
 
+
     @staticmethod
-    def compute_id(public_key: str, created_at: int, kind: int, tags: "list[list[str]]", content: str) -> str:
+    def compute_id(public_key: str, created_at: int, kind: int, tags: List[List[str]], content: str):
         return sha256(Event.serialize(public_key, created_at, kind, tags, content)).hexdigest()
 
+
+    @property
+    def id(self) -> str:
+        # Always recompute the id to reflect the up-to-date state of the Event
+        return Event.compute_id(self.public_key, self.created_at, self.kind, self.tags, self.content)
+
+
+    def add_pubkey_ref(self, pubkey:str):
+        """ Adds a reference to a pubkey as a 'p' tag """
+        self.tags.append(['p', pubkey])
+
+
+    def add_event_ref(self, event_id:str):
+        """ Adds a reference to an event_id as an 'e' tag """
+        self.tags.append(['e', event_id])
+
+
     def verify(self) -> bool:
-        pub_key = PublicKey(bytes.fromhex("02" + self.public_key), True) # add 02 for schnorr (bip340)
-        event_id = Event.compute_id(self.public_key, self.created_at, self.kind, self.tags, self.content)
-        return pub_key.schnorr_verify(bytes.fromhex(event_id), bytes.fromhex(self.signature), None, raw=True)
+        pub_key = PublicKey(bytes.fromhex("02" + self.public_key), True)  # add 02 for schnorr (bip340)
+        return pub_key.schnorr_verify(bytes.fromhex(self.id), bytes.fromhex(self.signature), None, raw=True)
+
 
     def to_message(self) -> str:
         return json.dumps(
@@ -67,3 +87,37 @@ def to_message(self) -> str:
                 }
             ]
         )
+
+
+
+@dataclass
+class EncryptedDirectMessage(Event):
+    recipient_pubkey: str = None
+    cleartext_content: str = None
+    reference_event_id: str = None
+
+
+    def __post_init__(self):
+        if self.content is not None:
+            self.cleartext_content = self.content
+            self.content = None
+
+        if self.recipient_pubkey is None:
+            raise Exception("Must specify a recipient_pubkey.")
+
+        self.kind = EventKind.ENCRYPTED_DIRECT_MESSAGE
+        super().__post_init__()
+
+        # Must specify the DM recipient's pubkey in a 'p' tag
+        self.add_pubkey_ref(self.recipient_pubkey)
+
+        # Optionally specify a reference event (DM) this is a reply to
+        if self.reference_event_id is not None:
+            self.add_event_ref(self.reference_event_id)
+
+
+    @property
+    def id(self) -> str:
+        if self.content is None:
+            raise Exception("EncryptedDirectMessage `id` is undefined until its message is encrypted and stored in the `content` field")
+        return super().id

nostr/key.py

@@ -7,7 +7,7 @@
 from hashlib import sha256
 
 from nostr.delegation import Delegation
-from nostr.event import Event
+from nostr.event import EncryptedDirectMessage, Event, EventKind
 from . import bech32
 
 
@@ -77,6 +77,9 @@ def encrypt_message(self, message: str, public_key_hex: str) -> str:
         encrypted_message = encryptor.update(padded_data) + encryptor.finalize()
 
         return f"{base64.b64encode(encrypted_message).decode()}?iv={base64.b64encode(iv).decode()}"
+    
+    def encrypt_dm(self, dm: EncryptedDirectMessage) -> None:
+        dm.content = self.encrypt_message(message=dm.cleartext_content, public_key_hex=dm.recipient_pubkey)
 
     def decrypt_message(self, encoded_message: str, public_key_hex: str) -> str:
         encoded_data = encoded_message.split('?iv=')
@@ -100,6 +103,10 @@ def sign_message_hash(self, hash: bytes) -> str:
         return sig.hex()
 
     def sign_event(self, event: Event) -> None:
+        if event.kind == EventKind.ENCRYPTED_DIRECT_MESSAGE and event.content is None:
+            self.encrypt_dm(event)
+        if event.public_key is None:
+            event.public_key = self.public_key.hex()
         event.signature = self.sign_message_hash(bytes.fromhex(event.id))
 
     def sign_delegation(self, delegation: Delegation) -> None:
@@ -108,6 +115,7 @@ def sign_delegation(self, delegation: Delegation) -> None:
     def __eq__(self, other):
         return self.raw_secret == other.raw_secret
 
+
 def mine_vanity_key(prefix: str = None, suffix: str = None) -> PrivateKey:
     if prefix is None and suffix is None:
         raise ValueError("Expected at least one of 'prefix' or 'suffix' arguments")
@@ -122,6 +130,7 @@ def mine_vanity_key(prefix: str = None, suffix: str = None) -> PrivateKey:
 
     return sk
 
+
 ffi = FFI()
 @ffi.callback("int (unsigned char *, const unsigned char *, const unsigned char *, void *)")
 def copy_x(output, x32, y32, data):

test/test_event.py

@@ -1,14 +1,93 @@
-from nostr.event import Event
-from nostr.key import PrivateKey
+import pytest
 import time
+from nostr.event import Event, EncryptedDirectMessage
+from nostr.key import PrivateKey
+
+
+
+class TestEvent:
+    def test_event_default_time(self):
+        """
+            ensure created_at default value reflects the time at Event object instantiation
+            see: https://github.com/jeffthibault/python-nostr/issues/23
+        """
+        event1 = Event(content='test event')
+        time.sleep(1.5)
+        event2 = Event(content='test event')
+        assert event1.created_at < event2.created_at
+    
+
+    def test_content_only_instantiation(self):
+        """ should be able to create an Event by only specifying content without kwarg """
+        event = Event("Hello, world!")
+        assert event.content is not None
+
+
+    def test_event_id_recomputes(self):
+        """ should recompute the Event.id to reflect the current Event attrs """
+        event = Event(content="some event")
+
+        # id should be computed on the fly
+        event_id = event.id
+
+        event.created_at += 10
+
+        # Recomputed id should now be different
+        assert event.id != event_id
+    
+
+    def test_add_event_ref(self):
+        """ should add an 'e' tag for each event_ref added """
+        some_event_id = "some_event_id"
+        event = Event(content="Adding an 'e' tag")
+        event.add_event_ref(some_event_id)
+        assert ['e', some_event_id] in event.tags
+
+
+    def test_add_pubkey_ref(self):
+        """ should add a 'p' tag for each pubkey_ref added """
+        some_pubkey = "some_pubkey"
+        event = Event(content="Adding a 'p' tag")
+        event.add_pubkey_ref(some_pubkey)
+        assert ['p', some_pubkey] in event.tags
+
+
+
+class TestEncryptedDirectMessage:
+    def setup_class(self):
+        self.sender_pk = PrivateKey()
+        self.sender_pubkey = self.sender_pk.public_key.hex()
+        self.recipient_pk = PrivateKey()
+        self.recipient_pubkey = self.recipient_pk.public_key.hex()
+
+
+    def test_content_field_moved_to_cleartext_content(self):
+        """ Should transfer `content` field data to `cleartext_content` """
+        dm = EncryptedDirectMessage(content="My message!", recipient_pubkey=self.recipient_pubkey)
+        assert dm.content is None
+        assert dm.cleartext_content is not None
+    
+
+    def test_nokwarg_content_allowed(self):
+        """ Should allow creating a new DM w/no `content` nor `cleartext_content` kwarg """
+        dm = EncryptedDirectMessage("My message!", recipient_pubkey=self.recipient_pubkey)
+        assert dm.cleartext_content is not None
+    
+
+    def test_recipient_p_tag(self):
+        """ Should generate recipient 'p' tag """
+        dm = EncryptedDirectMessage(cleartext_content="Secret message!", recipient_pubkey=self.recipient_pubkey)
+        assert ['p', self.recipient_pubkey] in dm.tags
+    
+
+    def test_unencrypted_dm_has_undefined_id(self):
+        """ Should raise Exception if `id` is requested before DM is encrypted """
+        dm = EncryptedDirectMessage(cleartext_content="My message!", recipient_pubkey=self.recipient_pubkey)
+
+        with pytest.raises(Exception) as e:
+            dm.id
+        assert "undefined" in str(e)
 
-def test_event_default_time():
-    """
-        ensure created_at default value reflects the time at Event object instantiation
-        see: https://github.com/jeffthibault/python-nostr/issues/23
-    """
-    public_key = PrivateKey().public_key.hex()
-    event1 = Event(public_key=public_key, content='test event')
-    time.sleep(1.5)
-    event2 = Event(public_key=public_key, content='test event')
-    assert event1.created_at < event2.created_at
+        # But once we encrypt it, we can request its id
+        self.sender_pk.encrypt_dm(dm)
+        assert dm.id is not None