arno/trezor-firmware
1
0
Fork 0
mirror of https://github.com/trezor/trezor-firmware.git synced 2025-05-29 12:18:51 +00:00
Code Issues Releases Wiki Activity

python: make PIL optional in the toif module, add size check

Browse Source
This commit is contained in:
matejcik 2020-08-25 13:14:47 +02:00
parent eddaeb1280
commit a14634c389
2 changed files with 17 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
python/src/trezorlib/debuglink.py
View File

@ -446,20 +446,6 @@ class TrezorClientDebugLink(TrezorClient):
def _raw_read(self): def _raw_read(self):
__tracebackhide__ = True # for pytest # pylint: disable=W0612 __tracebackhide__ = True # for pytest # pylint: disable=W0612
# if SCREENSHOT and self.debug:
# from PIL import Image
# layout = self.debug.state().layout
# im = Image.new("RGB", (128, 64))
# pix = im.load()
# for x in range(128):
# for y in range(64):
# rx, ry = 127 - x, 63 - y
# if (ord(layout[rx + (ry / 8) * 128]) & (1 << (ry % 8))) > 0:
# pix[x, y] = (255, 255, 255)
# im.save("scr%05d.png" % self.screenshot_id)
# self.screenshot_id += 1
resp = super()._raw_read() resp = super()._raw_read()
resp = self._filter_message(resp) resp = self._filter_message(resp)
self._check_request(resp) self._check_request(resp)

28
python/src/trezorlib/toif.py
View File

@ -4,12 +4,13 @@ from typing import Sequence, Tuple
import attr import attr
from . import firmware
try: try:
from PIL import Image from PIL import Image
except ImportError: except ImportError:
Image = None Image = None
from . import firmware
RGBPixel = Tuple[int, int, int] RGBPixel = Tuple[int, int, int]
@ -66,19 +67,14 @@ class Toif:
size = attr.ib() # type: Tuple[int, int] size = attr.ib() # type: Tuple[int, int]
data = attr.ib() # type: bytes data = attr.ib() # type: bytes
def _expected_data_length(self) -> int: @data.validator
def check_data_size(self, _, value):
width, height = self.size width, height = self.size
if self.mode is firmware.ToifMode.grayscale: if self.mode is firmware.ToifMode.grayscale:
return width * height // 2 expected_size = width * height // 2
else: else:
return width * height * 2 expected_size = width * height * 2
def to_image(self) -> "Image":
if Image is None:
raise RuntimeError("PIL is not available. Please install via 'pip install Pillow'")
uncompressed = _decompress(self.data) uncompressed = _decompress(self.data)
expected_size = self._expected_data_length()
if len(uncompressed) != expected_size: if len(uncompressed) != expected_size:
raise ValueError( raise ValueError(
"Uncompressed data is {} bytes, expected {}".format( "Uncompressed data is {} bytes, expected {}".format(
@ -86,6 +82,14 @@ class Toif:
) )
) )
def to_image(self) -> "Image":
if Image is None:
raise RuntimeError(
"PIL is not available. Please install via 'pip install Pillow'"
)
uncompressed = _decompress(self.data)
if self.mode is firmware.ToifMode.grayscale: if self.mode is firmware.ToifMode.grayscale:
pil_mode = "L" pil_mode = "L"
raw_data = _to_grayscale(uncompressed) raw_data = _to_grayscale(uncompressed)
@ -118,7 +122,9 @@ def load(filename: str) -> Toif:
def from_image(image: "Image", background=(0, 0, 0, 255)) -> Toif: def from_image(image: "Image", background=(0, 0, 0, 255)) -> Toif:
if Image is None: if Image is None:
raise RuntimeError("PIL is not available. Please install via 'pip install Pillow'") raise RuntimeError(
"PIL is not available. Please install via 'pip install Pillow'"
)
if image.mode == "RGBA": if image.mode == "RGBA":
background = Image.new("RGBA", image.size, background) background = Image.new("RGBA", image.size, background)