InputData.py

import hashlib
import json
import re
import copy
from typing import Any, Optional, Union, Callable
import struct
import functools

from ..client import EthAppClient, EIP712FieldType
from ..signing_partners import CAL_COIN_META_PARTNER
from ..status_word import StatusWord


# global variables
app_client: EthAppClient = None  # type: ignore[assignment]
filtering_paths: dict = {}
filtering_tokens: list[dict] = []
filtering_calldatas: list[dict] = []
current_path: list[str] = []
sig_ctx: dict[str, Any] = {}


# From a string typename, extract the type and all the array depth
# Input  = "uint8[2][][4]"          |   "bool"
# Output = ('uint8', [2, None, 4])  |   ('bool', [])
def get_array_levels(typename):
    array_lvls = []
    regex = re.compile(r"(.*)\[([0-9]*)\]$")

    while True:
        result = regex.search(typename)
        if not result:
            break
        typename = result.group(1)

        level_size = result.group(2)
        if len(level_size) == 0:
            level_size = None
        else:
            level_size = int(level_size)
        array_lvls.insert(0, level_size)
    return (typename, array_lvls)


# From a string typename, extract the type and its size
# Input  = "uint64"         |   "string"
# Output = ('uint', 64)     |   ('string', None)
def get_typesize(typename):
    regex = re.compile(r"^(\w+?)(\d*)$")
    result = regex.search(typename)
    typename = result.group(1)
    typesize = result.group(2)
    if len(typesize) == 0:
        typesize = None
    else:
        typesize = int(typesize)
    return (typename, typesize)


def parse_int(typesize):
    return (EIP712FieldType.INT, int(typesize / 8))


def parse_uint(typesize):
    return (EIP712FieldType.UINT, int(typesize / 8))


def parse_address(typesize):
    return (EIP712FieldType.ADDRESS, None)


def parse_bool(typesize):
    return (EIP712FieldType.BOOL, None)


def parse_string(typesize):
    return (EIP712FieldType.STRING, None)


def parse_bytes(typesize):
    if typesize is not None:
        return (EIP712FieldType.FIX_BYTES, typesize)
    return (EIP712FieldType.DYN_BYTES, None)


# set functions for each type
parsing_type_functions = {}
parsing_type_functions["int"] = parse_int
parsing_type_functions["uint"] = parse_uint
parsing_type_functions["address"] = parse_address
parsing_type_functions["bool"] = parse_bool
parsing_type_functions["string"] = parse_string
parsing_type_functions["bytes"] = parse_bytes


def send_struct_def_field(typename, keyname):
    type_enum = None

    (typename, array_lvls) = get_array_levels(typename)
    (typename, typesize) = get_typesize(typename)

    if typename in parsing_type_functions:
        (type_enum, typesize) = parsing_type_functions[typename](typesize)
    else:
        type_enum = EIP712FieldType.CUSTOM
        typesize = None

    with app_client.eip712_send_struct_def_struct_field(type_enum,
                                                        typename,
                                                        typesize,
                                                        array_lvls,
                                                        keyname):
        pass

    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending field def {keyname} of type {typename}: {response.status}"

    return (typename, type_enum, typesize, array_lvls)


def encode_integer(value: Union[str, int], typesize: int) -> bytes:
    # Some are already represented as integers in the JSON, but most as strings
    if isinstance(value, str):
        value = int(value, 0)

    if value == 0:
        data = b'\x00'
    else:
        # biggest uint type accepted by struct.pack
        uint64_mask = 0xffffffffffffffff
        data = struct.pack(">QQQQ",
                           (value >> 192) & uint64_mask,
                           (value >> 128) & uint64_mask,
                           (value >> 64) & uint64_mask,
                           value & uint64_mask)
        data = data[len(data) - typesize:]
        data = data.lstrip(b'\x00')
    return data


def encode_int(value: str, typesize: int) -> bytes:
    return encode_integer(value, typesize)


def encode_uint(value: str, typesize: int) -> bytes:
    return encode_integer(value, typesize)


def encode_hex_string(value: str, size: int) -> bytes:
    assert value.startswith("0x")
    value = value[2:]
    if len(value) < (size * 2):
        value = value.rjust(size * 2, "0")
    assert len(value) == (size * 2)
    return bytes.fromhex(value)


def encode_address(value: str, typesize: int) -> bytes:
    return encode_hex_string(value, 20)


def encode_bool(value: str, typesize: int) -> bytes:
    return encode_integer(value, 1)


def encode_string(value: str, typesize: int) -> bytes:
    return value.encode()


def encode_bytes_fix(value: str, typesize: int) -> bytes:
    return encode_hex_string(value, typesize)


def encode_bytes_dyn(value: str, typesize: int) -> bytes:
    # length of the value string
    # - the length of 0x (2)
    # / by the length of one byte in a hex string (2)
    return encode_hex_string(value, int((len(value) - 2) / 2))


# set functions for each type
encoding_functions = {}
encoding_functions[EIP712FieldType.INT] = encode_int
encoding_functions[EIP712FieldType.UINT] = encode_uint
encoding_functions[EIP712FieldType.ADDRESS] = encode_address
encoding_functions[EIP712FieldType.BOOL] = encode_bool
encoding_functions[EIP712FieldType.STRING] = encode_string
encoding_functions[EIP712FieldType.FIX_BYTES] = encode_bytes_fix
encoding_functions[EIP712FieldType.DYN_BYTES] = encode_bytes_dyn


def send_all_filtering_tokens(tokens: list[dict]):
    for token in tokens:
        response = app_client.provide_token_metadata(token["ticker"],
                                                     bytes.fromhex(token["addr"][2:]),
                                                     token["decimals"],
                                                     token["chain_id"])
        assert response.status == StatusWord.SWO_SUCCESS, \
            f"Error sending token metadata for {token['ticker']}: {response.status}"


def send_filter(path: str, discarded: bool) -> Optional[Callable]:
    ret: Optional[Callable] = None
    assert path in filtering_paths.keys()

    if filtering_paths[path]["type"].startswith("amount_join_"):
        if "id" in filtering_paths[path].keys():
            join_id = filtering_paths[path]["id"]
        else:
            # Permit (ERC-2612)
            join_id = 0xff
        if filtering_paths[path]["type"].endswith("_token"):
            send_filtering_amount_join_token(path, join_id, discarded)
        elif filtering_paths[path]["type"].endswith("_value"):
            send_filtering_amount_join_value(path,
                                             join_id,
                                             filtering_paths[path]["name"],
                                             discarded)
    elif filtering_paths[path]["type"] == "datetime":
        send_filtering_datetime(path, filtering_paths[path]["name"], discarded)
    elif filtering_paths[path]["type"] == "trusted_name":
        send_filtering_trusted_name(path,
                                    filtering_paths[path]["name"],
                                    filtering_paths[path]["tn_type"],
                                    filtering_paths[path]["tn_source"],
                                    discarded)
    elif filtering_paths[path]["type"].startswith("calldata_"):
        calldata_index = filtering_paths[path]["index"]
        for calldata in filtering_calldatas:
            if calldata["index"] == calldata_index:
                if not calldata["sent"]:
                    send_filtering_calldata_info(calldata["index"],
                                                 calldata["value_flag"],
                                                 calldata["callee_flag"],
                                                 calldata["chain_id_flag"],
                                                 calldata["selector_flag"],
                                                 calldata["amount_flag"],
                                                 calldata["spender_flag"])
                    calldata["sent"] = True
                break
        if filtering_paths[path]["type"].endswith("_value"):
            send_filtering_calldata_value(path, calldata_index, discarded)
        elif filtering_paths[path]["type"].endswith("_callee"):
            send_filtering_calldata_callee(path, calldata_index, discarded)
        elif filtering_paths[path]["type"].endswith("_chain_id"):
            send_filtering_calldata_chain_id(path, calldata_index, discarded)
        elif filtering_paths[path]["type"].endswith("_selector"):
            send_filtering_calldata_selector(path, calldata_index, discarded)
        elif filtering_paths[path]["type"].endswith("_amount"):
            send_filtering_calldata_amount(path, calldata_index, discarded)
        elif filtering_paths[path]["type"].endswith("_spender"):
            send_filtering_calldata_spender(path, calldata_index, discarded)
        else:
            assert False
        calldata["path_count"] -= 1
        if calldata["path_count"] == 0:
            ret = calldata["handler"]
    elif filtering_paths[path]["type"] == "raw":
        send_filtering_raw(path, filtering_paths[path]["name"], discarded)
    else:
        assert False

    return ret


def send_struct_impl_field(value, field):
    assert not isinstance(value, list)
    assert field["enum"] != EIP712FieldType.CUSTOM

    callback: Optional[Callable] = None

    data = encoding_functions[field["enum"]](value, field["typesize"])

    if filtering_paths:
        path = ".".join(current_path)
        if path in filtering_paths.keys():
            callback = send_filter(path, False)

    with app_client.eip712_send_struct_impl_struct_field(data):
        pass

    if callback is not None:
        callback()

    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending field {field['name']} of type {field['type']}: {response.status}"


def evaluate_field(structs, data, field, lvls_left, new_level=True):
    array_lvls = field["array_lvls"]

    if new_level:
        current_path.append(field["name"])
    if len(array_lvls) > 0 and lvls_left > 0:
        with app_client.eip712_send_struct_impl_array(len(data)):
            pass
        response = app_client.response()
        assert response is not None
        assert response.status == StatusWord.SWO_SUCCESS, \
            f"Error sending array {field['name']} of type {field['type']}: {response.status}"
        if len(data) == 0:
            for path in filtering_paths.keys():
                dpath = ".".join(current_path) + ".[]"
                if path.startswith(dpath):
                    response = app_client.eip712_filtering_discarded_path(path)
                    assert response is not None
                    assert response.status == StatusWord.SWO_SUCCESS, \
                        f"Error sending discarded path {path}: {response.status}"
                    send_filter(path, True)
        idx = 0
        for subdata in data:
            current_path.append("[]")
            evaluate_field(structs, subdata, field, lvls_left - 1, False)
            current_path.pop()
            idx += 1
        if array_lvls[lvls_left - 1] is not None:
            assert array_lvls[lvls_left - 1] == idx, \
                f"Mismatch in array size! Got {idx}, expected {array_lvls[lvls_left - 1]}"
    else:
        if field["enum"] == EIP712FieldType.CUSTOM:
            send_struct_impl(structs, data, field["type"])
        else:
            send_struct_impl_field(data, field)
    if new_level:
        current_path.pop()


def send_struct_impl(structs, data, structname):
    # Check if it is a struct we don't known
    assert structname in structs.keys(), \
        f"Unknown struct {structname} in types definition"

    for f in structs[structname]:
        evaluate_field(structs, data[f["name"]], f, len(f["array_lvls"]))


def start_signature_payload(ctx: dict, magic: int) -> bytearray:
    to_sign = bytearray()
    # magic number so that signature for one type of filter can't possibly be
    # valid for another, defined in APDU specs
    to_sign.append(magic)
    to_sign += ctx["chainid"]
    to_sign += ctx["caddr"]
    to_sign += ctx["schema_hash"]
    return to_sign


# ledgerjs doesn't actually sign anything, and instead uses already pre-computed signatures
def send_filtering_message_info(display_name: str, filters_count: int):
    to_sign = start_signature_payload(sig_ctx, 183)
    to_sign.append(filters_count)
    to_sign += display_name.encode()

    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    with app_client.eip712_filtering_message_info(display_name, filters_count, sig):
        pass
    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering message info for {display_name}: {response.status}"


def send_filtering_amount_join_token(path: str, join_id: int, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 11)
    to_sign += path.encode()
    to_sign.append(join_id)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    with app_client.eip712_filtering_amount_join_token(join_id, sig, discarded):
        pass
    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering amount join token for {path} with token index {join_id}: {response.status}"


def send_filtering_amount_join_value(path: str, join_id: int, display_name: str, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 22)
    to_sign += path.encode()
    to_sign += display_name.encode()
    to_sign.append(join_id)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    with app_client.eip712_filtering_amount_join_value(join_id, display_name, sig, discarded):
        pass
    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering amount join value for {path} with token index {join_id}: {response.status}"


def send_filtering_datetime(path: str, display_name: str, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 33)
    to_sign += path.encode()
    to_sign += display_name.encode()
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    with app_client.eip712_filtering_datetime(display_name, sig, discarded):
        pass
    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering datetime for {path}: {response.status}"


def send_filtering_trusted_name(path: str,
                                display_name: str,
                                name_type: list[int],
                                name_source: list[int],
                                discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 44)
    to_sign += path.encode()
    to_sign += display_name.encode()
    for t in name_type:
        to_sign.append(t)
    for s in name_source:
        to_sign.append(s)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    with app_client.eip712_filtering_trusted_name(display_name, name_type, name_source, sig, discarded):
        pass
    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering trusted name for {path}: {response.status}"


def send_filtering_calldata_info(index: int,
                                 value_filter_flag: bool,
                                 callee_filter_flag: int,
                                 chain_id_filter_flag: bool,
                                 selector_filter_flag: bool,
                                 amount_filter_flag: bool,
                                 spender_filter_flag: int):
    to_sign = start_signature_payload(sig_ctx, 55)
    to_sign.append(index)
    to_sign.append(value_filter_flag)
    to_sign.append(int(callee_filter_flag))
    to_sign.append(chain_id_filter_flag)
    to_sign.append(selector_filter_flag)
    to_sign.append(amount_filter_flag)
    to_sign.append(int(spender_filter_flag))
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    response = app_client.eip712_filtering_calldata_info(index,
                                                         value_filter_flag,
                                                         callee_filter_flag,
                                                         chain_id_filter_flag,
                                                         selector_filter_flag,
                                                         amount_filter_flag,
                                                         spender_filter_flag,
                                                         sig)
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering calldata info : {response.status}"


def send_filtering_calldata_value(path: str, index: int, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 66)
    to_sign += path.encode()
    to_sign.append(index)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    response = app_client.eip712_filtering_calldata_value(index, sig, discarded)
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering calldata value for {path}: {response.status}"


def send_filtering_calldata_callee(path: str, index: int, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 77)
    to_sign += path.encode()
    to_sign.append(index)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    response = app_client.eip712_filtering_calldata_callee(index, sig, discarded)
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering calldata callee for {path}: {response.status}"


def send_filtering_calldata_chain_id(path: str, index: int, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 88)
    to_sign += path.encode()
    to_sign.append(index)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    response = app_client.eip712_filtering_calldata_chain_id(index, sig, discarded)
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering calldata callee for {path}: {response.status}"


def send_filtering_calldata_selector(path: str, index: int, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 99)
    to_sign += path.encode()
    to_sign.append(index)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    response = app_client.eip712_filtering_calldata_selector(index, sig, discarded)
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering calldata callee for {path}: {response.status}"


def send_filtering_calldata_amount(path: str, index: int, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 110)
    to_sign += path.encode()
    to_sign.append(index)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    response = app_client.eip712_filtering_calldata_amount(index, sig, discarded)
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering calldata callee for {path}: {response.status}"


def send_filtering_calldata_spender(path: str, index: int, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 121)
    to_sign += path.encode()
    to_sign.append(index)
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    response = app_client.eip712_filtering_calldata_spender(index, sig, discarded)
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering calldata callee for {path}: {response.status}"


# ledgerjs doesn't actually sign anything, and instead uses already pre-computed signatures
def send_filtering_raw(path: str, display_name: str, discarded: bool):
    to_sign = start_signature_payload(sig_ctx, 72)
    to_sign += path.encode()
    to_sign += display_name.encode()
    sig = CAL_COIN_META_PARTNER.sign(bytes(to_sign))
    with app_client.eip712_filtering_raw(display_name, sig, discarded):
        pass
    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending filtering raw for {path}: {response.status}"


def prepare_filtering(data_json, filtr_data):
    global filtering_paths
    global filtering_tokens
    global filtering_calldatas

    if "fields" in filtr_data:
        filtering_paths = filtr_data["fields"]
    else:
        filtering_paths = {}

    if "tokens" in filtr_data:
        filtering_tokens = filtr_data["tokens"]
    else:
        filtering_tokens = []

    if "calldatas" in filtr_data:
        filtering_calldatas = filtr_data["calldatas"]
        for calldata in filtering_calldatas:
            calldata["sent"] = False
            calldata["path_count"] = 0
            for path in filtering_paths.values():
                if path["type"].startswith("calldata_"):
                    if path["index"] == calldata["index"]:
                        calldata["path_count"] += 1
            if "handler" in calldata:
                if calldata["handler"] is not None:
                    calldata["handler"] = functools.partial(calldata["handler"], app_client, data_json)
            else:
                calldata["handler"] = None


def handle_optional_domain_values(domain):
    if "chainId" not in domain.keys():
        domain["chainId"] = 0
    if "verifyingContract" not in domain.keys():
        domain["verifyingContract"] = "0x0000000000000000000000000000000000000000"


def init_signature_context(sig_ctx, types, domain, filters):
    handle_optional_domain_values(domain)
    if "address" in filters:
        caddr = filters["address"]
    else:
        caddr = domain["verifyingContract"]
    if caddr.startswith("0x"):
        caddr = caddr[2:]
    sig_ctx["caddr"] = bytearray.fromhex(caddr)
    chainid = domain["chainId"]
    sig_ctx["chainid"] = bytearray()
    for i in range(8):
        sig_ctx["chainid"].append(chainid & (0xff << (i * 8)))
    sig_ctx["chainid"].reverse()

    # Order type fields
    for type_name in types.keys():
        for i in range(len(types[type_name])):
            types[type_name][i] = dict(sorted(types[type_name][i].items()))

    schema_str = json.dumps(types).replace(" ", "")
    schema_hash = hashlib.sha224(schema_str.encode())
    sig_ctx["schema_hash"] = bytearray.fromhex(schema_hash.hexdigest())


def process_data(aclient: EthAppClient,
                 data_json: dict,
                 filters: Optional[dict] = None) -> None:
    global app_client
    global current_path

    current_path = []
    # deepcopy because this function modifies the dict
    data_json = copy.deepcopy(data_json)
    app_client = aclient
    domain_typename = "EIP712Domain"
    message_typename = data_json["primaryType"]
    types = data_json["types"]
    domain = data_json["domain"]
    message = data_json["message"]

    if filters:
        init_signature_context(sig_ctx, types, domain, filters)

    # send types definition
    for key in types.keys():
        with app_client.eip712_send_struct_def_struct_name(key):
            pass
        response = app_client.response()
        assert response is not None
        assert response.status == StatusWord.SWO_SUCCESS, \
            f"Error sending struct def {key}: {response.status}"
        for f in types[key]:
            (f["type"], f["enum"], f["typesize"], f["array_lvls"]) = \
             send_struct_def_field(f["type"], f["name"])

    if filters:
        with app_client.eip712_filtering_activate():
            pass
        response = app_client.response()
        assert response is not None
        assert response.status == StatusWord.SWO_SUCCESS, \
            f"Error activating filtering: {response.status}"
        prepare_filtering(data_json, filters)
        send_all_filtering_tokens(filtering_tokens)

    # Send ledgerPKI certificate
    app_client.send_pki_certificate(CAL_COIN_META_PARTNER)

    # send domain implementation
    with app_client.eip712_send_struct_impl_root_struct(domain_typename):
        pass
    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending domain root struct {domain_typename}: {response.status}"
    send_struct_impl(types, domain, domain_typename)

    if filters:
        if filters and "name" in filters:
            send_filtering_message_info(filters["name"], len(filtering_paths))
        else:
            send_filtering_message_info(domain["name"], len(filtering_paths))

    # send message implementation
    with app_client.eip712_send_struct_impl_root_struct(message_typename):
        pass
    response = app_client.response()
    assert response is not None
    assert response.status == StatusWord.SWO_SUCCESS, \
        f"Error sending message root struct {message_typename}: {response.status}"
    send_struct_impl(types, message, message_typename)