hive-nectar 0.2.9__py3-none-any.whl

nectarbase/memo.py

@@ -0,0 +1,242 @@
+import hashlib
+import struct
+from binascii import hexlify, unhexlify
+from typing import Any, Tuple
+
+from Cryptodome.Cipher import AES
+
+from nectargraphenebase.account import PublicKey
+from nectargraphenebase.base58 import base58decode, base58encode
+from nectargraphenebase.types import varintdecode
+
+from .objects import Memo
+
+default_prefix = "STM"
+
+
+def get_shared_secret(priv: Any, pub: PublicKey) -> str:
+    """Derive the share secret between ``priv`` and ``pub``
+    :param `Base58` priv: Private Key
+    :param `Base58` pub: Public Key
+    :return: Shared secret
+    :rtype: hex
+    The shared secret is generated such that::
+        Pub(Alice) * Priv(Bob) = Pub(Bob) * Priv(Alice)
+    """
+    pub_point = pub.point()
+    priv_point = int(repr(priv), 16)
+    res = pub_point * priv_point
+    res_hex = "%032x" % res.x()
+    # Zero padding
+    res_hex = "0" * (64 - len(res_hex)) + res_hex
+    return res_hex
+
+
+def init_aes(shared_secret: str, nonce: int) -> Any:
+    """Initialize AES instance
+    :param hex shared_secret: Shared Secret to use as encryption key
+    :param int nonce: Random nonce
+    :return: AES instance
+    :rtype: AES
+    """
+    " Shared Secret "
+    ss = hashlib.sha512(unhexlify(shared_secret)).digest()
+    " Seed "
+    seed = bytes(str(nonce), "ascii") + hexlify(ss)
+    seed_digest = hexlify(hashlib.sha512(seed).digest()).decode("ascii")
+    " AES "
+    key = unhexlify(seed_digest[0:64])
+    iv = unhexlify(seed_digest[64:96])
+    return AES.new(key, AES.MODE_CBC, iv)
+
+
+def init_aes_bts(shared_secret: str, nonce: int) -> Any:
+    """Initialize AES instance
+    :param hex shared_secret: Shared Secret to use as encryption key
+    :param int nonce: Random nonce
+    :return: AES instance
+    :rtype: AES
+    """
+    " Shared Secret "
+    ss = hashlib.sha512(unhexlify(shared_secret)).digest()
+    " Seed "
+    seed = bytes(str(nonce), "ascii") + hexlify(ss)
+    seed_digest = hexlify(hashlib.sha512(seed).digest()).decode("ascii")
+    " AES "
+    key = unhexlify(seed_digest[0:64])
+    iv = unhexlify(seed_digest[64:96])
+    return AES.new(key, AES.MODE_CBC, iv)
+
+
+def init_aes2(shared_secret: str, nonce: int) -> Tuple[Any, int]:
+    """Initialize AES instance
+    :param hex shared_secret: Shared Secret to use as encryption key
+    :param int nonce: Random nonce
+    """
+    shared_secret = hashlib.sha512(unhexlify(shared_secret)).hexdigest()
+    # Seed
+    ss = unhexlify(shared_secret)
+    n = struct.pack("<Q", int(nonce))
+    encryption_key = hashlib.sha512(n + ss).hexdigest()
+    # Check'sum'
+    check = hashlib.sha256(unhexlify(encryption_key)).digest()
+    check = struct.unpack_from("<I", check[:4])[0]
+    # AES
+    key = unhexlify(encryption_key[0:64])
+    iv = unhexlify(encryption_key[64:96])
+    return AES.new(key, AES.MODE_CBC, iv), check
+
+
+def _pad(s: bytes, BS: int) -> bytes:
+    numBytes = BS - len(s) % BS
+    return s + numBytes * struct.pack("B", numBytes)
+
+
+def _unpad(s: bytes, BS: int) -> bytes:
+    count = s[-1]
+    if s[-count::] == count * struct.pack("B", count):
+        return s[:-count]
+    return s
+
+
+def encode_memo_bts(priv: Any, pub: PublicKey, nonce: int, message: str) -> str:
+    """Encode a message with a shared secret between Alice and Bob
+
+    :param PrivateKey priv: Private Key (of Alice)
+    :param PublicKey pub: Public Key (of Bob)
+    :param int nonce: Random nonce
+    :param str message: Memo message
+    :return: Encrypted message
+    :rtype: hex
+
+    """
+    shared_secret = get_shared_secret(priv, pub)
+    aes = init_aes_bts(shared_secret, nonce)
+    " Checksum "
+    raw = bytes(message, "utf8")
+    checksum = hashlib.sha256(raw).digest()
+    raw = checksum[0:4] + raw
+    " Padding "
+    raw = _pad(raw, 16)
+    " Encryption "
+    return hexlify(aes.encrypt(raw)).decode("ascii")
+
+
+def decode_memo_bts(priv: Any, pub: PublicKey, nonce: int, message: str) -> str:
+    """Decode a message with a shared secret between Alice and Bob
+
+    :param PrivateKey priv: Private Key (of Bob)
+    :param PublicKey pub: Public Key (of Alice)
+    :param int nonce: Nonce used for Encryption
+    :param bytes message: Encrypted Memo message
+    :return: Decrypted message
+    :rtype: str
+    :raise ValueError: if message cannot be decoded as valid UTF-8
+           string
+
+    """
+    shared_secret = get_shared_secret(priv, pub)
+    aes = init_aes_bts(shared_secret, nonce)
+    " Encryption "
+    raw = bytes(message, "ascii")
+    cleartext = aes.decrypt(unhexlify(raw))
+    " Checksum "
+    checksum = cleartext[0:4]
+    message_bytes = cleartext[4:]
+    message_bytes = _unpad(message_bytes, 16)
+    " Verify checksum "
+    check = hashlib.sha256(
+        message_bytes if isinstance(message_bytes, bytes) else message_bytes.encode("utf-8")
+    ).digest()[0:4]
+    if check != checksum:  # pragma: no cover
+        raise ValueError("checksum verification failure")
+    return message_bytes.decode("utf8") if isinstance(message_bytes, bytes) else message_bytes
+
+
+def encode_memo(priv: Any, pub: PublicKey, nonce: int, message: str, **kwargs: Any) -> str:
+    """Encode a message with a shared secret between Alice and Bob
+
+    :param PrivateKey priv: Private Key (of Alice)
+    :param PublicKey pub: Public Key (of Bob)
+    :param int nonce: Random nonce
+    :param str message: Memo message
+    :return: Encrypted message
+    :rtype: hex
+    """
+    shared_secret = get_shared_secret(priv, pub)
+    aes, check = init_aes2(shared_secret, nonce)
+    " Padding "
+    raw = bytes(message, "utf8")
+    raw = _pad(raw, 16)
+    " Encryption "
+    cipher = hexlify(aes.encrypt(raw)).decode("ascii")
+    prefix = kwargs.pop("prefix", default_prefix)
+    s = {
+        "from": format(priv.pubkey, prefix),
+        "to": format(pub, prefix),
+        "nonce": nonce,
+        "check": check,
+        "encrypted": cipher,
+        "prefix": prefix,
+    }
+    tx = Memo(**s)
+    return "#" + base58encode(hexlify(bytes(tx)).decode("ascii"))
+
+
+def extract_memo_data(message: str) -> Tuple[PublicKey, PublicKey, str, int, bytes]:
+    """Returns the stored pubkey keys, nonce, checksum and encrypted message of a memo"""
+    raw = base58decode(message[1:])
+    from_key = PublicKey(raw[:66])
+    raw = raw[66:]
+    to_key = PublicKey(raw[:66])
+    raw = raw[66:]
+    nonce = str(struct.unpack_from("<Q", unhexlify(raw[:16]))[0])
+    raw = raw[16:]
+    check = struct.unpack_from("<I", unhexlify(raw[:8]))[0]
+    raw = raw[8:]
+    cipher = unhexlify(raw)
+    return from_key, to_key, nonce, check, cipher
+
+
+def decode_memo(priv: Any, message: str) -> str:
+    """Decode a message with a shared secret between Alice and Bob
+
+    :param PrivateKey priv: Private Key (of Bob)
+    :param base58encoded message: Encrypted Memo message
+    :return: Decrypted message
+    :rtype: str
+    :raise ValueError: if message cannot be decoded as valid UTF-8
+           string
+    """
+    # decode structure
+    from_key, to_key, nonce, check, cipher = extract_memo_data(message)
+
+    if repr(to_key) == repr(priv.pubkey):
+        shared_secret = get_shared_secret(priv, from_key)
+    elif repr(from_key) == repr(priv.pubkey):
+        shared_secret = get_shared_secret(priv, to_key)
+    else:
+        raise ValueError("Incorrect PrivateKey")
+
+    # Init encryption
+    aes, checksum = init_aes2(shared_secret, int(nonce))
+    # Check
+    if not check == checksum:
+        raise AssertionError("Checksum failure")
+    # Encryption
+    # remove the varint prefix (FIXME, long messages!)
+    numBytes = 16 - len(cipher) % 16
+    n = 16 - numBytes
+    message_bytes = cipher[n:]
+    message_bytes = aes.decrypt(message_bytes)
+    message_bytes = _unpad(message_bytes, 16)
+    n = varintdecode(message_bytes)
+    if (len(message_bytes) - n) > 0 and (len(message_bytes) - n) < 8:
+        message_part = message_bytes[len(message_bytes) - n :]
+        if isinstance(message_part, bytes):
+            message_part = message_part.decode("utf8")
+        return "#" + message_part
+    else:
+        return "#" + (
+            message_bytes.decode("utf8") if isinstance(message_bytes, bytes) else message_bytes
+        )