From fb1deb615613d6dcc4785e97372df7a2db254ee4 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Tue, 25 Aug 2020 13:15:52 +0200
Subject: [PATCH] python/trezorctl: improve 'set homescreen' command

---
 python/src/trezorlib/cli/settings.py | 92 ++++++++++++++++++++--------
 1 file changed, 67 insertions(+), 25 deletions(-)

diff --git a/python/src/trezorlib/cli/settings.py b/python/src/trezorlib/cli/settings.py
index 212cae355..63e4603a5 100644
--- a/python/src/trezorlib/cli/settings.py
+++ b/python/src/trezorlib/cli/settings.py
@@ -16,9 +16,15 @@
 
 import click
 
-from .. import device, messages
+from .. import device, firmware, messages, toif
 from . import ChoiceType, with_client
 
+try:
+    from PIL import Image
+except ImportError:
+    Image = None
+
+
 ROTATION = {"north": 0, "east": 90, "south": 180, "west": 270}
 SAFETY_LEVELS = {
     "strict": messages.SafetyCheckLevel.Strict,
@@ -26,6 +32,51 @@ SAFETY_LEVELS = {
 }
 
 
+def image_to_t1(filename: str) -> bytes:
+    if Image is None:
+        raise click.ClickException(
+            "Image library is missing. Please install via 'pip install Pillow'."
+        )
+
+    image = Image.open(filename)
+    if image.size != (128, 64):
+        raise click.ClickException("Wrong size of the image - should be 128x64")
+
+    image = image.convert("1")
+    return image.tobytes("raw", "1")
+
+
+def image_to_tt(filename: str) -> bytes:
+    if filename.endswith(".toif"):
+        try:
+            toif_image = toif.load(filename)
+        except Exception as e:
+            raise click.ClickException("TOIF file is corrupted") from e
+
+    elif Image is None:
+        raise click.ClickException(
+            "Image library is missing. Please install via 'pip install Pillow'."
+        )
+
+    else:
+        try:
+            image = Image.open(filename)
+            toif_image = toif.from_image(image)
+        except Exception as e:
+            raise click.ClickException(
+                "Failed to convert image to Trezor format"
+            ) from e
+
+    if toif_image.size != (144, 144):
+        raise click.ClickException("Wrong size of image - should be 144x144")
+
+    if toif_image.mode != firmware.ToifMode.full_color:
+        raise click.ClickException("Wrong image mode - should be full_color")
+
+    toif_image = toif.from_image(image)
+    return toif_image.to_bytes()
+
+
 @click.group(name="set")
 def cli():
     """Device settings."""
@@ -108,36 +159,27 @@ def flags(client, flags):
 
 
 @cli.command()
-@click.argument(
-    "filename", type=click.Path(dir_okay=False, readable=True), required=False
-)
+@click.argument("filename")
 @click.option(
     "-f", "--filename", "_ignore", is_flag=True, hidden=True, expose_value=False
 )
 @with_client
 def homescreen(client, filename):
-    """Set new homescreen."""
-    if filename is None:
-        img = b"\x00"
-    elif filename.endswith(".toif"):
-        img = open(filename, "rb").read()
-        if img[:8] != b"TOIf\x90\x00\x90\x00":
-            raise click.ClickException("File is not a TOIF file with size of 144x144")
+    """Set new homescreen.
+
+    To revert to default homescreen, use 'trezorctl set homescreen default'
+    """
+    if filename == "default":
+        img = b""
     else:
-        from PIL import Image
-
-        im = Image.open(filename)
-        if im.size != (128, 64):
-            raise click.ClickException("Wrong size of the image")
-        im = im.convert("1")
-        pix = im.load()
-        img = bytearray(1024)
-        for j in range(64):
-            for i in range(128):
-                if pix[i, j]:
-                    o = i + j * 128
-                    img[o // 8] |= 1 << (7 - o % 8)
-        img = bytes(img)
+        # use Click's facility to validate the path for us
+        param = click.Path(dir_okay=False, readable=True, exists=True)
+        param.convert(filename, None, None)
+        if client.features.model == "1":
+            img = image_to_t1(filename)
+        else:
+            img = image_to_tt(filename)
+
     return device.apply_settings(client, homescreen=img)