enc-pk.py

import hashlib
import binascii
import base58
import unicodedata
import scrypt

from Crypto import Random
from Crypto.Cipher import AES

from neo.Wallets.KeyPair import KeyPair
from neo.SmartContract.Contract import Contract

SCRYPT_ITERATIONS = 16384
SCRYPT_BLOCKSIZE = 8
SCRYPT_PARALLEL_FACTOR = 8
SCRYPT_KEY_LEN_BYTES = 64

NEP_HEADER = bytearray([0x01, 0x42])
NEP_FLAG = bytearray([0xe0])


def gen_encrypted_key(pwd, private_key=None):
    """
    * https://github.com/neo-project/proposals/blob/master/nep-2.mediawiki#encryption-steps
    * https://github.com/CityOfZion/neon-wallet-react-native/blob/master/app/api/crypto/index.js#L133
    """
    if not private_key:
        private_key = bytes(Random.get_random_bytes(32))
        # print("privkey:", binascii.hexlify(private_key))

    account = KeyPair(priv_key=private_key)
    contract = Contract.CreateSignatureContract(account.PublicKey)
    address = contract.Address

    # hash address twice
    address_hash_tmp = hashlib.sha256(address.encode('utf-8')).digest()
    address_hash2 = hashlib.sha256(address_hash_tmp).digest()
    # print(address_hash2)

    # first 4 bytes
    addressHash = address_hash2[:4]
    # print("addressHash:", addressHash, len(addressHash), type(addressHash))

    pwd_normalized = bytes(unicodedata.normalize('NFC', pwd), 'utf-8')
    # print('pwd-norm:', pwd_normalized, type(pwd_normalized))
    derived = scrypt.hash(pwd_normalized, addressHash,
                          N=SCRYPT_ITERATIONS,
                          r=SCRYPT_BLOCKSIZE,
                          p=SCRYPT_PARALLEL_FACTOR,
                          buflen=SCRYPT_KEY_LEN_BYTES)

    derived1 = derived[:32]
    derived2 = derived[32:]
    # print("derived :", derived)
    # print("derived1:", derived1)
    # print("derived2:", derived2)

    def xor(a, b):
        assert isinstance(a, bytes)
        assert isinstance(b, bytes)
        res = bytearray()
        for i in range(len(a)):
            res.append(a[i] ^ b[i])
        return bytes(res)

    xor_ed = xor(private_key, derived1)
    # print("xor: ", xor_ed)

    cypher = AES.new(derived2, AES.MODE_ECB)
    encrypted = cypher.encrypt(xor_ed)
    # print("encrypted:", encrypted)

    assembled = bytearray()
    assembled.extend(NEP_HEADER)
    assembled.extend(NEP_FLAG)
    assembled.extend(addressHash)
    assembled.extend(encrypted)
    # print("assembled:", assembled)

    encrypted_key = base58.b58encode_check(bytes(assembled))
    return encrypted_key, address


def gen_encrypted_key_from_wif(wif, pwd):
    private_key = KeyPair.PrivateKeyFromWIF(wif)
    return gen_encrypted_key(pwd, private_key)


def test():
    #
    # Test 1: Normal inputs
    #
    # inputs
    pwd = "TestingOneTwoThree"
    my_privatekey = "cbf4b9f70470856bb4f40f80b87edb90865997ffee6df315ab166d713af433a5"

    # expected outputs
    target_address = "AStZHy8E6StCqYQbzMqi4poH7YNDHQKxvt"
    target_encrypted_key = "6PYVPVe1fQznphjbUxXP9KZJqPMVnVwCx5s5pr5axRJ8uHkMtZg97eT5kL"

    # run
    encrypted_key, address = gen_encrypted_key(pwd, binascii.unhexlify(my_privatekey))
    assert(address == target_address)
    assert(encrypted_key == target_encrypted_key)

    #
    # Test 2: Unicode+Emoji password
    #
    # inputs
    pwd = "hellö♥️"
    my_privatekey = "03eb20a711f93c04459000c62cc235f9e9da82382206b812b07fd2f81779aa42"

    # expected outputs
    target_address = "AXQUduANGZF4e7wDazVAtyRLHwMounaUMA"
    target_encrypted_key = "6PYWdv8bP9vbfGsNnjzDawCoXCYpk4rnWG8xTZrvdzx6FjB6jv4H9MM586"

    # run
    encrypted_key, address = gen_encrypted_key(pwd, binascii.unhexlify(my_privatekey))
    assert(address == target_address)
    assert(encrypted_key == target_encrypted_key)

    #
    # Test 3: Create encrypted key from WIF
    #
    # inputs
    pwd = "asdasdasdasd"
    wif = "L4bsJCYWjwKugCMjqQSokgfFtCT7qDWybAHaUFoMhUTJzgsXDGmA"

    # expected outputs
    target_address = "ANeSmVGUM7rH1iMMDFBjNXdCjTdqjFxT5U"
    target_encrypted_key = "6PYR1BKNwF1Xpe6LGbkDZcJYRVijZ1QPmMRj7wF6N75FLvLdwGKBwXWErg"

    # run
    encrypted_key, address = gen_encrypted_key_from_wif(wif, pwd)
    assert(address == target_address)
    assert(encrypted_key == target_encrypted_key)


test()
print("Tests ok!")