Polkadot (Blockchain Platform)

# From: https://github.com/polkascan/py-scale-codec/blob/master/scalecodec/utils/ss58.py

""" SS58 is a simple address format designed for Substrate based chains.
    Encoding/decoding according to specification on
    https://docs.substrate.io/reference/address-formats/

"""
from typing import Optional, Union

import base58
from hashlib import blake2b


def ss58_decode(address: str, valid_address_type: Optional[int] = None) -> str:
    """
    Decodes given SS58 encoded address to an account ID
    Parameters
    ----------
    address: e.g. EaG2CRhJWPb7qmdcJvy3LiWdh26Jreu9Dx6R1rXxPmYXoDk
    valid_address_type

    Returns
    -------
    Decoded string AccountId
    """

    # Check if address is already decoded
    if address.startswith('0x'):
        return address

    if address == '':
        raise ValueError("Empty address provided")

    checksum_prefix = b'SS58PRE'

    address_decoded = base58.b58decode(address)

    if address_decoded[0] & 0b0100_0000:
        address_type_length = 2
        address_type = ((address_decoded[0] & 0b0011_1111) << 2) | (address_decoded[1] >> 6) | \
                       ((address_decoded[1] & 0b0011_1111) << 8)
    else:
        address_type_length = 1
        address_type = address_decoded[0]

    if address_type in [46, 47]:
        raise ValueError(f"{address_type} is a reserved SS58 format")

    if valid_address_type is not None and address_type != valid_address_type:
        raise ValueError("Invalid SS58 format")

    # Determine checksum length according to length of address string
    if len(address_decoded) in [3, 4, 6, 10]:
        checksum_length = 1
    elif len(address_decoded) in [5, 7, 11, 34 + address_type_length, 35 + address_type_length]:
        checksum_length = 2
    elif len(address_decoded) in [8, 12]:
        checksum_length = 3
    elif len(address_decoded) in [9, 13]:
        checksum_length = 4
    elif len(address_decoded) in [14]:
        checksum_length = 5
    elif len(address_decoded) in [15]:
        checksum_length = 6
    elif len(address_decoded) in [16]:
        checksum_length = 7
    elif len(address_decoded) in [17]:
        checksum_length = 8
    else:
        raise ValueError("Invalid address length")

    checksum = blake2b(checksum_prefix + address_decoded[0:-checksum_length]).digest()

    if checksum[0:checksum_length] != address_decoded[-checksum_length:]:
        raise ValueError("Invalid checksum")

    return address_decoded[address_type_length:len(address_decoded)-checksum_length].hex()


def ss58_encode(account_id: Union[str, bytes], address_type: int = 42) -> str:
    """
    Encodes an account ID to an Substrate address according to provided address_type

    Parameters
    ----------
    account_id: the public key, e.g. 0x192c3c7e5789b461fbf1c7f614ba5eed0b22efc507cda60a5e7fda8e046bcdce
    address_type: the address type, e.g. 0 for Polkadot, 2 for Kusama, 42 for generic Substrate. The prefix in https://github.com/paritytech/ss58-registry/blob/main/ss58-registry.json

    Returns
    -------
    str
    """
    checksum_prefix = b'SS58PRE'

    if address_type < 0 or address_type > 16383 or address_type in [46, 47]:
        raise ValueError("Invalid value for ss58_format")

    if type(account_id) is bytes or type(account_id) is bytearray:
        account_id_bytes = account_id
    else:
        account_id_bytes = bytes.fromhex(account_id.replace('0x', ''))

    if len(account_id_bytes) in [32, 33]:
        # Checksum size is 2 bytes for public key
        checksum_length = 2
    elif len(account_id_bytes) in [1, 2, 4, 8]:
        # Checksum size is 1 byte for account index
        checksum_length = 1
    else:
        raise ValueError("Invalid length for address")

    if address_type < 64:
        address_type_bytes = bytes([address_type])
    else:
        address_type_bytes = bytes([
            ((address_type & 0b0000_0000_1111_1100) >> 2) | 0b0100_0000,
            (address_type >> 8) | ((address_type & 0b0000_0000_0000_0011) << 6)
        ])

    input_bytes = address_type_bytes + account_id_bytes
    checksum = blake2b(checksum_prefix + input_bytes).digest()

    return base58.b58encode(input_bytes + checksum[:checksum_length]).decode()


if __name__ == "__main__":
    pub_key = "0x192c3c7e5789b461fbf1c7f614ba5eed0b22efc507cda60a5e7fda8e046bcdce"
    polkadot_address = ss58_encode(pub_key, 0)
    print("Polkadot address:", polkadot_address)  # 1a1LcBX6hGPKg5aQ6DXZpAHCCzWjckhea4sz3P1PvL3oc4F

    substrate_address = ss58_encode(pub_key, 42)
    print("Substrate address:", substrate_address)  # 5CdiCGvTEuzut954STAXRfL8Lazs3KCZa5LPpkPeqqJXdTHp