diff --git a/src/trezor/crypto/der.py b/src/trezor/crypto/der.py index 8ad0564d63..6019116e13 100644 --- a/src/trezor/crypto/der.py +++ b/src/trezor/crypto/der.py @@ -1,4 +1,4 @@ -def encode_length(l: int): +def encode_length(l: int) -> bytes: if l < 0x80: return bytes([l]) elif l <= 0xFF: @@ -8,13 +8,13 @@ def encode_length(l: int): else: raise ValueError -def encode_int(i: bytes): +def encode_int(i: bytes) -> bytes: i = i.lstrip(b'\x00') if i[0] >= 0x80: i = b'\x00' + i return b'\x02' + encode_length(len(i)) + i -def encode_seq(seq: tuple): +def encode_seq(seq: tuple) -> bytes: res = b'' for i in seq: res += encode_int(i) diff --git a/src/trezor/crypto/rlp.py b/src/trezor/crypto/rlp.py index 13fb58d09c..4105ddfc64 100644 --- a/src/trezor/crypto/rlp.py +++ b/src/trezor/crypto/rlp.py @@ -1,4 +1,4 @@ -def to_binary(x: int): +def int_to_bytes(x: int) -> bytes: if x == 0: return b'' r = bytearray() @@ -7,27 +7,28 @@ def to_binary(x: int): x //= 256 return bytes(reversed(r)) -def encode_length(l: int, offset: int): +def encode_length(l: int, is_list: bool) -> bytes: + offset = 0xC0 if is_list else 0x80 if l < 56: return bytes([l + offset]) elif l < 256 ** 8: - bl = to_binary(l) + bl = int_to_bytes(l) return bytes([len(bl) + offset + 55]) + bl else: raise ValueError('Input too long') -def encode(data): +def encode(data) -> bytes: if isinstance(data, int): - return encode(to_binary(data)) + return encode(int_to_bytes(data)) if isinstance(data, bytes): if len(data) == 1 and ord(data) < 128: return data else: - return encode_length(len(data), 128) + data + return encode_length(len(data), is_list=False) + data elif isinstance(data, list): output = b'' for item in data: output += encode(item) - return encode_length(len(output), 192) + output + return encode_length(len(output), is_list=True) + output else: raise TypeError('Invalid input')