add EncryptedDirectMessage class; simplify Event class (#39)

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 @@ class Event():
                 }
             ]
         )
+
+
+
+@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