From 53c524c90a681b7ddf4e846541efe361a302ac1e Mon Sep 17 00:00:00 2001
From: jeffthibault <jdthibault2@gmail.com>
Date: Mon, 25 Jul 2022 09:44:28 -0400
Subject: [PATCH] add nip-04 encrypted dm support

---
 nostr/key.py     | 35 ++++++++++++++++++++++++++++++++++-
 requirements.txt |  1 +
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/nostr/key.py b/nostr/key.py
index 4274e36..c5afaae 100644
--- a/nostr/key.py
+++ b/nostr/key.py
@@ -1,4 +1,8 @@
-from secp256k1 import PrivateKey
+import os
+import base64
+from secp256k1 import PrivateKey, PublicKey
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives import padding
 
 def generate_private_key() -> str:
     private_key = PrivateKey()
@@ -17,4 +21,33 @@ def get_key_pair() -> tuple:
     private_key = PrivateKey()
     public_key = private_key.pubkey.serialize().hex()
     return (private_key.serialize(), public_key[2:])
+
+def compute_shared_secret(sender_private_key: str, receiver_public_key: str) -> str:
+    public_key = PublicKey(bytes.fromhex("02" + receiver_public_key), True)
+    return public_key.ecdh(bytes.fromhex(sender_private_key)).hex() 
+
+def encrypt_message(content: str, shared_secret: str) -> str:
+    iv = os.urandom(16)
+
+    cipher = Cipher(algorithms.AES(bytes.fromhex(shared_secret)), modes.CBC(iv))
+    encryptor = cipher.encryptor()
+    
+    padder = padding.PKCS7(128).padder()
+    padded_data = padder.update(content.encode()) + padder.finalize()
+    encrypted_message = encryptor.update(padded_data) + encryptor.finalize()
+    
+    return f"{base64.b64encode(encrypted_message).decode()}?iv={base64.b64encode(iv).decode()}"
+
+def decrypt_message(encoded_message: str, shared_secret: str) -> str:
+    encoded_data = encoded_message.split('?iv=')
+    encoded_content, encoded_iv = encoded_data[0], encoded_data[1]
+
+    encrypted_content = base64.b64decode(encoded_content)
+    iv = base64.b64decode(encoded_iv)
+
+    cipher = Cipher(algorithms.AES(bytes.fromhex(shared_secret)), modes.CBC(iv))
+    decryptor = cipher.decryptor()
+    decrypted_message = decryptor.update(encrypted_content) + decryptor.finalize()
+
+    return decrypted_message.decode()
     
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 68d3690..bfc81af 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 cffi==1.15.0
+cryptography==37.0.4
 pycparser==2.21
 secp256k1==0.14.0
 websocket-client==1.2.3