mirror of
https://github.com/trezor/trezor-firmware.git
synced 2024-11-22 15:38:11 +00:00
146 lines
4.1 KiB
Python
146 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import typing as t
|
|
from hashlib import sha256
|
|
|
|
from trezorlib import cosi, definitions, messages, protobuf
|
|
|
|
PRIVATE_KEYS_DEV = [byte * 32 for byte in (b"\xdd", b"\xde", b"\xdf")]
|
|
|
|
|
|
def sign_with_privkeys(digest: bytes, privkeys: t.Sequence[bytes]) -> bytes:
|
|
"""Locally produce a CoSi signature."""
|
|
pubkeys = [cosi.pubkey_from_privkey(sk) for sk in privkeys]
|
|
nonces = [cosi.get_nonce(sk, digest, i) for i, sk in enumerate(privkeys)]
|
|
|
|
global_pk = cosi.combine_keys(pubkeys)
|
|
global_R = cosi.combine_keys(R for _, R in nonces)
|
|
|
|
sigs = [
|
|
cosi.sign_with_privkey(digest, sk, global_pk, r, global_R)
|
|
for sk, (r, _) in zip(privkeys, nonces)
|
|
]
|
|
|
|
return cosi.combine_sig(global_R, sigs)
|
|
|
|
|
|
def make_network(
|
|
chain_id: int = 0,
|
|
slip44: int = 0,
|
|
symbol: str = "FAKE",
|
|
name: str = "Fake network",
|
|
) -> messages.EthereumNetworkInfo:
|
|
return messages.EthereumNetworkInfo(
|
|
chain_id=chain_id,
|
|
slip44=slip44,
|
|
symbol=symbol,
|
|
name=name,
|
|
)
|
|
|
|
|
|
def make_token(
|
|
symbol: str = "FAKE",
|
|
decimals: int = 18,
|
|
address: bytes = b"",
|
|
chain_id: int = 0,
|
|
name: str = "Fake token",
|
|
) -> messages.EthereumTokenInfo:
|
|
return messages.EthereumTokenInfo(
|
|
symbol=symbol,
|
|
decimals=decimals,
|
|
address=address,
|
|
chain_id=chain_id,
|
|
name=name,
|
|
)
|
|
|
|
|
|
def make_payload(
|
|
data_type: messages.EthereumDefinitionType = messages.EthereumDefinitionType.NETWORK,
|
|
timestamp: int = 0xFFFF_FFFF,
|
|
message: messages.EthereumNetworkInfo
|
|
| messages.EthereumTokenInfo
|
|
| bytes = make_network(),
|
|
) -> bytes:
|
|
if isinstance(message, bytes):
|
|
message_bytes = message
|
|
else:
|
|
writer = io.BytesIO()
|
|
protobuf.dump_message(writer, message)
|
|
message_bytes = writer.getvalue()
|
|
|
|
payload = definitions.DefinitionPayload(
|
|
magic=b"trzd1",
|
|
data_type=data_type,
|
|
timestamp=timestamp,
|
|
data=message_bytes,
|
|
)
|
|
return payload.build()
|
|
|
|
|
|
def sign_payload(
|
|
payload: bytes,
|
|
merkle_neighbors: list[bytes],
|
|
threshold: int = 3,
|
|
) -> tuple[bytes, bytes]:
|
|
digest = sha256(b"\x00" + payload).digest()
|
|
merkle_proof = []
|
|
for item in merkle_neighbors:
|
|
left, right = min(digest, item), max(digest, item)
|
|
digest = sha256(b"\x01" + left + right).digest()
|
|
merkle_proof.append(digest)
|
|
|
|
merkle_proof = len(merkle_proof).to_bytes(1, "little") + b"".join(merkle_proof)
|
|
signature = sign_with_privkeys(digest, PRIVATE_KEYS_DEV[:threshold])
|
|
sigmask = 0
|
|
for i in range(threshold):
|
|
sigmask |= 1 << i
|
|
sigmask_byte = sigmask.to_bytes(1, "little")
|
|
return merkle_proof, sigmask_byte + signature
|
|
|
|
|
|
def encode_network(
|
|
network: messages.EthereumNetworkInfo | None = None,
|
|
chain_id: int = 0,
|
|
slip44: int = 0,
|
|
symbol: str = "FAKE",
|
|
name: str = "Fake network",
|
|
) -> bytes:
|
|
if network is None:
|
|
network = make_network(chain_id, slip44, symbol, name)
|
|
payload = make_payload(
|
|
data_type=messages.EthereumDefinitionType.NETWORK, message=network
|
|
)
|
|
proof, signature = sign_payload(payload, [])
|
|
return payload + proof + signature
|
|
|
|
|
|
def encode_token(
|
|
token: messages.EthereumTokenInfo | None = None,
|
|
symbol: str = "FakeTok",
|
|
decimals: int = 18,
|
|
address: t.AnyStr = b"",
|
|
chain_id: int = 0,
|
|
name: str = "Fake token",
|
|
) -> bytes:
|
|
if token is None:
|
|
if isinstance(address, str):
|
|
if address.startswith("0x"):
|
|
address = address[2:]
|
|
address = bytes.fromhex(address) # type: ignore (typechecker is lying)
|
|
token = make_token(symbol, decimals, address, chain_id, name) # type: ignore (typechecker is lying)
|
|
payload = make_payload(
|
|
data_type=messages.EthereumDefinitionType.TOKEN, message=token
|
|
)
|
|
proof, signature = sign_payload(payload, [])
|
|
return payload + proof + signature
|
|
|
|
|
|
def make_defs(
|
|
network: bytes | None, token: bytes | None
|
|
) -> messages.EthereumDefinitions:
|
|
return messages.EthereumDefinitions(
|
|
encoded_network=network,
|
|
encoded_token=token,
|
|
)
|