From 8323b3e948f48249e0c24f25a414d183444f1404 Mon Sep 17 00:00:00 2001
From: jeffthibault <jdthibault2@gmail.com>
Date: Wed, 20 Jul 2022 14:51:44 -0400
Subject: [PATCH] implement basic protocol in nip-01

---
 nostr/__init__.py     |  0
 nostr/event.py        | 54 ++++++++++++++++++++++++++
 nostr/filter.py       | 82 ++++++++++++++++++++++++++++++++++++++++
 nostr/key.py          | 20 ++++++++++
 nostr/message_type.py |  9 +++++
 nostr/relay.py        | 88 +++++++++++++++++++++++++++++++++++++++++++
 nostr/subscription.py |  6 +++
 7 files changed, 259 insertions(+)
 create mode 100644 nostr/__init__.py
 create mode 100644 nostr/event.py
 create mode 100644 nostr/filter.py
 create mode 100644 nostr/key.py
 create mode 100644 nostr/message_type.py
 create mode 100644 nostr/relay.py
 create mode 100644 nostr/subscription.py

diff --git a/nostr/__init__.py b/nostr/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nostr/event.py b/nostr/event.py
new file mode 100644
index 0000000..a8f5ba9
--- /dev/null
+++ b/nostr/event.py
@@ -0,0 +1,54 @@
+import time
+import json
+from enum import IntEnum
+from secp256k1 import PrivateKey, PublicKey
+from hashlib import sha256
+
+class EventKind(IntEnum):
+    SET_METADATA = 0
+    TEXT_NOTE = 1
+    RECOMMEND_RELAY = 2
+
+class Event():
+    def __init__(
+            self, 
+            public_key: str, 
+            content: str, 
+            created_at: int=int(time.time()), 
+            kind: int=EventKind.TEXT_NOTE, 
+            tags: "list[list[str]]"=[], 
+            id: int=None, 
+            signature: str=None) -> None:
+        self.id = id if id != None else sha256(Event.serialize(public_key, created_at, kind, tags, content)).hexdigest()
+        self.public_key = public_key
+        self.content = content
+        self.created_at = created_at
+        self.kind = kind
+        self.tags = tags
+        self.signature = signature
+
+    @staticmethod
+    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=(',', ':'))
+        return data_str.encode()
+
+    def sign(self, sk: str) -> None:
+        private_key = PrivateKey(bytes.fromhex(sk))
+        sig = private_key.schnorr_sign(bytes.fromhex(self.id), None, raw=True)
+        self.signature = sig.hex()
+
+    def verify(self) -> bool:
+        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_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "pubkey": self.public_key,
+            "created_at": self.created_at,
+            "kind": self.kind,
+            "tags": self.tags,
+            "content": self.content,
+            "sig": self.signature
+        }
diff --git a/nostr/filter.py b/nostr/filter.py
new file mode 100644
index 0000000..1cd36dc
--- /dev/null
+++ b/nostr/filter.py
@@ -0,0 +1,82 @@
+from collections import UserList
+from .event import Event
+
+class Filter:
+    def __init__(
+            self, 
+            ids: "list[str]"=None, 
+            kinds: "list[int]"=None, 
+            authors: "list[str]"=None, 
+            since: int=None, 
+            until: int=None, 
+            tags: "dict[str, list[str]]"=None,
+            limit: int=None) -> None:
+        self.IDs = ids
+        self.kinds = kinds
+        self.authors = authors
+        self.since = since
+        self.until = until
+        self.tags = tags
+        self.limit = limit
+
+    def matches(self, event: Event) -> bool:
+        if self.IDs != None and event.id not in self.IDs:
+            return False
+        if self.kinds != None and event.kind not in self.kinds:
+            return False
+        if self.authors != None and event.public_key not in self.authors:
+            return False
+        if self.since != None and event.created_at < self.since:
+            return False
+        if self.until != None and event.created_at > self.until:
+            return False
+        if self.tags != None and len(event.tags) == 0:
+            return False
+        if self.tags != None:
+            e_tag_identifiers = [e_tag[0] for e_tag in event.tags] 
+            for f_tag, f_tag_values in self.tags.items():
+                if f_tag[1:] not in e_tag_identifiers:
+                    return False
+                for e_tag in event.tags:
+                    if e_tag[1] not in f_tag_values:
+                        return False
+        
+        return True
+
+    def to_json(self) -> dict:
+        res = {}
+        if self.IDs != None:
+            res["ids"] = self.IDs
+        if self.kinds != None:   
+            res["kinds"] = self.kinds
+        if self.authors != None:
+            res["authors"] = self.authors
+        if self.since != None:
+            res["since"] = self.since
+        if self.until != None:
+            res["until"] = self.until
+        if self.tags != None:
+            for tag, values in self.tags.items():
+                res[tag] = values
+        if self.limit != None:
+            res["limit"] = self.limit
+
+        return res
+                
+class Filters(UserList):
+    def __init__(self, initlist: "list[Filter]"=[]) -> None:
+        super().__init__(initlist)
+        self.data: "list[Filter]"
+
+    def match(self, event: Event):
+        for filter in self.data:
+            if filter.matches(event):
+                return True
+        return False
+
+    def to_json(self) -> list:
+        res = []
+        for filter in self.data:
+            res.append(filter.to_json())
+        return res
+        
\ No newline at end of file
diff --git a/nostr/key.py b/nostr/key.py
new file mode 100644
index 0000000..4274e36
--- /dev/null
+++ b/nostr/key.py
@@ -0,0 +1,20 @@
+from secp256k1 import PrivateKey
+
+def generate_private_key() -> str:
+    private_key = PrivateKey()
+    public_key = private_key.pubkey.serialize().hex()
+    while not public_key.startswith("02"):
+        private_key = PrivateKey()
+        public_key = private_key.pubkey.serialize().hex()
+    return private_key.serialize()
+
+def get_public_key(secret: str) -> str:
+    private_key = PrivateKey(bytes.fromhex(secret))
+    public_key = private_key.pubkey.serialize().hex()
+    return public_key[2:] # chop off sign byte
+
+def get_key_pair() -> tuple:
+    private_key = PrivateKey()
+    public_key = private_key.pubkey.serialize().hex()
+    return (private_key.serialize(), public_key[2:])
+    
\ No newline at end of file
diff --git a/nostr/message_type.py b/nostr/message_type.py
new file mode 100644
index 0000000..c35e099
--- /dev/null
+++ b/nostr/message_type.py
@@ -0,0 +1,9 @@
+class ClientMessageType:
+    EVENT = "EVENT"
+    REQUEST = "REQ"
+    CLOSE = "CLOSE"
+
+class RelayMessageType:
+    EVENT = "EVENT"
+    NOTICE = "NOTICE"
+    
\ No newline at end of file
diff --git a/nostr/relay.py b/nostr/relay.py
new file mode 100644
index 0000000..b563808
--- /dev/null
+++ b/nostr/relay.py
@@ -0,0 +1,88 @@
+import json
+import ssl
+from typing import Union
+from websocket import WebSocket, WebSocketConnectionClosedException, WebSocketTimeoutException
+from .event import Event
+from .filter import Filters
+from .message_type import RelayMessageType
+from .subscription import Subscription
+
+class RelayPolicy:
+    def __init__(self, should_read: bool=True, should_write: bool=True) -> None:
+        self.should_read = should_read
+        self.should_write = should_write
+
+class Relay:
+    def __init__(
+            self, 
+            url: str, 
+            policy: RelayPolicy,
+            ws: WebSocket=WebSocket(sslopt={"cert_reqs": ssl.CERT_NONE}),
+            subscriptions: dict[str, Subscription]={}) -> None:
+        self.url = url
+        self.policy = policy
+        self.ws = ws
+        self.subscriptions = subscriptions
+
+    def open_websocket_connection(self) -> None:
+        self.ws.connect(self.url, timeout=1)
+
+    def close_websocket_connection(self) -> None:
+        self.ws.close()
+
+    def add_subscription(self, id: str, filters: Filters) -> None:
+        self.subscriptions[id] = Subscription(id, filters)
+
+    def close_subscription(self, id: str) -> None:
+        self.subscriptions.pop(id)
+
+    def update_subscription(self, id: str, filters: Filters) -> None:
+        subscription = self.subscriptions[id]
+        subscription.filters = filters
+
+    def publish_message(self, message: str) -> None:
+        self.ws.send(message)
+
+    def get_message(self) -> Union[None, str]:
+        while True:
+            try:
+                message = self.ws.recv()
+                if not self._is_valid_message(message):
+                    continue
+                
+                return message
+                
+            except WebSocketConnectionClosedException:
+                print('received connection closed')
+                break
+            except WebSocketTimeoutException:
+                print('ws connection timed out')
+                break
+
+        return None
+
+    def _is_valid_message(self, message: str) -> bool:
+        if not message or message[0] != '[' or message[-1] != ']':
+            return False
+
+        message_json = json.loads(message)
+        message_type = message_json[0]
+        if message_type == RelayMessageType.NOTICE:
+            return True
+        if message_type == RelayMessageType.EVENT:
+            if not len(message_json) == 3:
+                return False
+            
+            subscription_id = message_json[1]
+            if subscription_id not in self.subscriptions:
+                return False
+
+            e = message_json[2]
+            event = Event(e['pubkey'], e['content'], e['created_at'], e['kind'], e['tags'], e['id'], e['sig'])
+            if not event.verify():
+                return False
+
+            if not self.subscriptions[subscription_id].filters.match(event):
+                return False
+
+            return True
diff --git a/nostr/subscription.py b/nostr/subscription.py
new file mode 100644
index 0000000..4d3814e
--- /dev/null
+++ b/nostr/subscription.py
@@ -0,0 +1,6 @@
+from .filter import Filters
+
+class Subscription:
+    def __init__(self, id: str, filters: Filters=None) -> None:
+        self.id = id
+        self.filters = filters