relative import

This commit is contained in:
callebtc
2023-02-08 10:10:12 +01:00
parent c4fbeb83d7
commit 557a01697e
2 changed files with 30 additions and 15 deletions

View File

@@ -3,7 +3,7 @@ import json
from dataclasses import dataclass, field from dataclasses import dataclass, field
from enum import IntEnum from enum import IntEnum
from typing import List from typing import List
from secp256k1 import PrivateKey, PublicKey from secp256k1 import PublicKey
from hashlib import sha256 from hashlib import sha256
from .message_type import ClientMessageType from .message_type import ClientMessageType

View File

@@ -6,8 +6,8 @@ from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding from cryptography.hazmat.primitives import padding
from hashlib import sha256 from hashlib import sha256
from nostr.delegation import Delegation from .delegation import Delegation
from nostr.event import EncryptedDirectMessage, Event, EventKind from .event import EncryptedDirectMessage, Event, EventKind
from . import bech32 from . import bech32
@@ -28,14 +28,14 @@ class PublicKey:
@classmethod @classmethod
def from_npub(cls, npub: str): def from_npub(cls, npub: str):
""" Load a PublicKey from its bech32/npub form """ """Load a PublicKey from its bech32/npub form"""
hrp, data, spec = bech32.bech32_decode(npub) hrp, data, spec = bech32.bech32_decode(npub)
raw_public_key = bech32.convertbits(data, 5, 8)[:-1] raw_public_key = bech32.convertbits(data, 5, 8)[:-1]
return cls(bytes(raw_public_key)) return cls(bytes(raw_public_key))
class PrivateKey: class PrivateKey:
def __init__(self, raw_secret: bytes=None) -> None: def __init__(self, raw_secret: bytes = None) -> None:
if not raw_secret is None: if not raw_secret is None:
self.raw_secret = raw_secret self.raw_secret = raw_secret
else: else:
@@ -46,7 +46,7 @@ class PrivateKey:
@classmethod @classmethod
def from_nsec(cls, nsec: str): def from_nsec(cls, nsec: str):
""" Load a PrivateKey from its bech32/nsec form """ """Load a PrivateKey from its bech32/nsec form"""
hrp, data, spec = bech32.bech32_decode(nsec) hrp, data, spec = bech32.bech32_decode(nsec)
raw_secret = bech32.convertbits(data, 5, 8)[:-1] raw_secret = bech32.convertbits(data, 5, 8)[:-1]
return cls(bytes(raw_secret)) return cls(bytes(raw_secret))
@@ -71,22 +71,28 @@ class PrivateKey:
padded_data = padder.update(message.encode()) + padder.finalize() padded_data = padder.update(message.encode()) + padder.finalize()
iv = secrets.token_bytes(16) iv = secrets.token_bytes(16)
cipher = Cipher(algorithms.AES(self.compute_shared_secret(public_key_hex)), modes.CBC(iv)) cipher = Cipher(
algorithms.AES(self.compute_shared_secret(public_key_hex)), modes.CBC(iv)
)
encryptor = cipher.encryptor() encryptor = cipher.encryptor()
encrypted_message = encryptor.update(padded_data) + encryptor.finalize() encrypted_message = encryptor.update(padded_data) + encryptor.finalize()
return f"{base64.b64encode(encrypted_message).decode()}?iv={base64.b64encode(iv).decode()}" return f"{base64.b64encode(encrypted_message).decode()}?iv={base64.b64encode(iv).decode()}"
def encrypt_dm(self, dm: EncryptedDirectMessage) -> None: def encrypt_dm(self, dm: EncryptedDirectMessage) -> None:
dm.content = self.encrypt_message(message=dm.cleartext_content, public_key_hex=dm.recipient_pubkey) 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: def decrypt_message(self, encoded_message: str, public_key_hex: str) -> str:
encoded_data = encoded_message.split('?iv=') encoded_data = encoded_message.split("?iv=")
encoded_content, encoded_iv = encoded_data[0], encoded_data[1] encoded_content, encoded_iv = encoded_data[0], encoded_data[1]
iv = base64.b64decode(encoded_iv) iv = base64.b64decode(encoded_iv)
cipher = Cipher(algorithms.AES(self.compute_shared_secret(public_key_hex)), modes.CBC(iv)) cipher = Cipher(
algorithms.AES(self.compute_shared_secret(public_key_hex)), modes.CBC(iv)
)
encrypted_content = base64.b64decode(encoded_content) encrypted_content = base64.b64decode(encoded_content)
decryptor = cipher.decryptor() decryptor = cipher.decryptor()
@@ -110,7 +116,9 @@ class PrivateKey:
event.signature = self.sign_message_hash(bytes.fromhex(event.id)) event.signature = self.sign_message_hash(bytes.fromhex(event.id))
def sign_delegation(self, delegation: Delegation) -> None: def sign_delegation(self, delegation: Delegation) -> None:
delegation.signature = self.sign_message_hash(sha256(delegation.delegation_token.encode()).digest()) delegation.signature = self.sign_message_hash(
sha256(delegation.delegation_token.encode()).digest()
)
def __eq__(self, other): def __eq__(self, other):
return self.raw_secret == other.raw_secret return self.raw_secret == other.raw_secret
@@ -122,9 +130,12 @@ def mine_vanity_key(prefix: str = None, suffix: str = None) -> PrivateKey:
while True: while True:
sk = PrivateKey() sk = PrivateKey()
if prefix is not None and not sk.public_key.bech32()[5:5+len(prefix)] == prefix: if (
prefix is not None
and not sk.public_key.bech32()[5 : 5 + len(prefix)] == prefix
):
continue continue
if suffix is not None and not sk.public_key.bech32()[-len(suffix):] == suffix: if suffix is not None and not sk.public_key.bech32()[-len(suffix) :] == suffix:
continue continue
break break
@@ -132,7 +143,11 @@ def mine_vanity_key(prefix: str = None, suffix: str = None) -> PrivateKey:
ffi = FFI() ffi = FFI()
@ffi.callback("int (unsigned char *, const unsigned char *, const unsigned char *, void *)")
@ffi.callback(
"int (unsigned char *, const unsigned char *, const unsigned char *, void *)"
)
def copy_x(output, x32, y32, data): def copy_x(output, x32, y32, data):
ffi.memmove(output, x32, 32) ffi.memmove(output, x32, 32)
return 1 return 1