From 083d8877444965e02974aa693e402e1c65105c3c Mon Sep 17 00:00:00 2001 From: jeffthibault Date: Sun, 30 Oct 2022 15:57:20 -0400 Subject: [PATCH] refactor key.py: add PrivateKey and PublicKey classes --- nostr/key.py | 121 ++++++++++++++++++++++++--------------------------- 1 file changed, 56 insertions(+), 65 deletions(-) diff --git a/nostr/key.py b/nostr/key.py index f353233..8c21e5b 100644 --- a/nostr/key.py +++ b/nostr/key.py @@ -1,95 +1,86 @@ -import os +import secrets import base64 +import secp256k1 from cffi import FFI -from secp256k1 import PrivateKey, PublicKey from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes from cryptography.hazmat.primitives import padding from . import bech32 -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() +class PublicKey: + def __init__(self, raw_bytes: bytes) -> None: + self.raw_bytes= raw_bytes -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 bech32(self) -> str: + converted_bits = bech32.convertbits(self.raw_bytes, 8, 5) + return bech32.bech32_encode("npub", converted_bits, bech32.Encoding.BECH32) -def get_key_pair() -> tuple: - private_key = PrivateKey() - public_key = private_key.pubkey.serialize().hex() - return (private_key.serialize(), public_key[2:]) + def hex(self) -> str: + return self.raw_bytes.hex() -def bech32_encode_private_key(private_key: str) -> str: - converted_bits = bech32.convertbits(bytes.fromhex(private_key), 8, 5) - return bech32.bech32_encode("nsec", converted_bits, bech32.Encoding.BECH32) + def verify_signed_message_hash(self, hash: str, sig: str) -> bool: + pk = secp256k1.PublicKey(self.raw_bytes, True) + return pk.schnorr_verify(bytes.fromhex(hash), bytes.fromhex(sig), None, True) -def bech32_decode_private_key(private_key_bech32: str) -> str: - data = bech32.bech32_decode(private_key_bech32)[1] - return bytes(bech32.convertbits(data, 5, 8, False)).hex() +class PrivateKey: + def __init__(self, raw_secret: bytes=None) -> None: + if not raw_secret is None: + self.raw_secret = raw_secret + else: + self.raw_secret = secrets.token_bytes(32) -def bech32_encode_public_key(public_key: str) -> str: - converted_bits = bech32.convertbits(bytes.fromhex(public_key), 8, 5) - return bech32.bech32_encode("npub", converted_bits, bech32.Encoding.BECH32) + sk = secp256k1.PrivateKey(self.raw_secret) + self.public_key = PublicKey(sk.pubkey.serialize()[1:]) -def bech32_decode_public_key(public_key_bech32: str) -> str: - data = bech32.bech32_decode(public_key_bech32)[1] - return bytes(bech32.convertbits(data, 5, 8, False)).hex() + def bech32(self) -> str: + converted_bits = bech32.convertbits(self.raw_secret, 8, 5) + return bech32.bech32_encode("nsec", converted_bits, bech32.Encoding.BECH32) -def tweak_add_private_key(private_key: str, scalar: bytes) -> str: - sk = PrivateKey(bytes.fromhex(private_key)) - tweaked_secret = sk.tweak_add(scalar) - new_sk = PrivateKey(tweaked_secret) - return new_sk.serialize() + def hex(self) -> str: + return self.raw_secret.hex() -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), hashfn=copy_x).hex() + def tweak_add(self, scalar: bytes) -> bytes: + sk = secp256k1.PrivateKey(self.raw_secret) + return sk.tweak_add(scalar) -def encrypt_message(content: str, shared_secret: str) -> str: - iv = os.urandom(16) + def compute_shared_secret(self, public_key_hex: str) -> bytes: + pk = secp256k1.PublicKey(bytes.fromhex("02" + public_key_hex), True) + return pk.ecdh(self.raw_secret, hashfn=copy_x) - cipher = Cipher(algorithms.AES(bytes.fromhex(shared_secret)), modes.CBC(iv)) - padder = padding.PKCS7(128).padder() - padded_data = padder.update(content.encode()) + padder.finalize() + def encrypt_message(self, message: str, public_key_hex: str) -> str: + padder = padding.PKCS7(128).padder() + padded_data = padder.update(message.encode()) + padder.finalize() - encryptor = cipher.encryptor() - encrypted_message = encryptor.update(padded_data) + encryptor.finalize() + iv = secrets.token_bytes(16) + cipher = Cipher(algorithms.AES(self.compute_shared_secret(public_key_hex)), modes.CBC(iv)) - return f"{base64.b64encode(encrypted_message).decode()}?iv={base64.b64encode(iv).decode()}" + encryptor = cipher.encryptor() + encrypted_message = encryptor.update(padded_data) + encryptor.finalize() -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] + return f"{base64.b64encode(encrypted_message).decode()}?iv={base64.b64encode(iv).decode()}" - encrypted_content = base64.b64decode(encoded_content) - iv = base64.b64decode(encoded_iv) + def decrypt_message(self, encoded_message: str, public_key_hex: str) -> str: + encoded_data = encoded_message.split('?iv=') + encoded_content, encoded_iv = encoded_data[0], encoded_data[1] - cipher = Cipher(algorithms.AES(bytes.fromhex(shared_secret)), modes.CBC(iv)) - decryptor = cipher.decryptor() - decrypted_message = decryptor.update(encrypted_content) + decryptor.finalize() + iv = base64.b64decode(encoded_iv) + cipher = Cipher(algorithms.AES(self.compute_shared_secret(public_key_hex)), modes.CBC(iv)) + encrypted_content = base64.b64decode(encoded_content) - unpadder = padding.PKCS7(128).unpadder() - unpadded_data = unpadder.update(decrypted_message) + unpadder.finalize() + decryptor = cipher.decryptor() + decrypted_message = decryptor.update(encrypted_content) + decryptor.finalize() - return unpadded_data.decode() + unpadder = padding.PKCS7(128).unpadder() + unpadded_data = unpadder.update(decrypted_message) + unpadder.finalize() -def sign_message(hash: str, private_key: str) -> str: - sk = PrivateKey(bytes.fromhex(private_key)) - sig = sk.schnorr_sign(bytes.fromhex(hash), None, raw=True) - return sig.hex() + return unpadded_data.decode() -def verify_message(hash: str, sig: str, public_key: str) -> bool: - pk = PublicKey(bytes.fromhex("02" + public_key), True) - return pk.schnorr_verify(bytes.fromhex(hash), bytes.fromhex(sig), None, True) + def sign_message_hash(self, hash: bytes) -> str: + sk = secp256k1.PrivateKey(self.raw_secret) + sig = sk.schnorr_sign(hash, None, raw=True) + return sig.hex() ffi = FFI() @ffi.callback("int (unsigned char *, const unsigned char *, const unsigned char *, void *)") def copy_x(output, x32, y32, data): ffi.memmove(output, x32, 32) - return 1 - \ No newline at end of file + return 1 \ No newline at end of file