From 61e4f79c90519ddff03c4e2fc9a07bfa83c0b4ba Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Thu, 28 Nov 2019 16:50:08 +0100
Subject: [PATCH] core: use common implementation for toif_convert

---
 core/tools/toif_convert.py | 160 +++++--------------------------------
 1 file changed, 21 insertions(+), 139 deletions(-)

diff --git a/core/tools/toif_convert.py b/core/tools/toif_convert.py
index 6a528e0b1f..377ca09612 100755
--- a/core/tools/toif_convert.py
+++ b/core/tools/toif_convert.py
@@ -1,153 +1,35 @@
 #!/usr/bin/env python3
-import io
-import struct
-import sys
-import zlib
-from os.path import basename
 
+import click
 from PIL import Image
 
+from trezorlib._internal import toif
 
-def png2toif(input_data):
-    def process_rgb(w, h, pix):
-        data = bytes()
-        for j in range(h):
-            for i in range(w):
-                r, g, b = pix[i, j]
-                c = ((r & 0xF8) << 8) | ((g & 0xFC) << 3) | ((b & 0xF8) >> 3)
-                data += struct.pack(">H", c)
-        return data
 
-    def process_grayscale(w, h, pix):
-        data = bytes()
-        for j in range(h):
-            for i in range(w // 2):
-                l1, l2 = pix[i * 2, j], pix[i * 2 + 1, j]
-                c = (l1 & 0xF0) | (l2 >> 4)
-                data += struct.pack(">B", c)
-        return data
-
-    im = Image.open(io.BytesIO(input_data))
-    w, h = im.size
-    if im.mode not in ["RGB", "L"]:
-        raise ValueError("Unknown mode:", im.mode)
-    if im.mode == "L":
-        if w % 2 > 0:
-            raise ValueError("PNG file must have width divisible by 2")
-
-    pix = im.load()
-    if im.mode == "RGB":
-        pixeldata = process_rgb(w, h, pix)
-    else:
-        pixeldata = process_grayscale(w, h, pix)
-    z = zlib.compressobj(level=9, wbits=10)
-    zdata = z.compress(pixeldata) + z.flush()
-    zdata = zdata[2:-4]  # strip header and checksum
-
-    data = b""
-    if im.mode == "RGB":
-        data += b"TOIf"
-    else:
-        data += b"TOIg"
-    data += struct.pack("<HH", w, h)
-    data += struct.pack("<I", len(zdata))
-    data += zdata
-
-    return data
+@click.command()
+@click.argument("infile", type=click.File("rb"))
+@click.argument("outfile", type=click.File("wb"))
+def toif_convert(infile, outfile):
+    """Convert any image format to/from TOIF or vice-versa.
 
+    \b
+    Examples:
+      toif_convert.py somefile.jpg outfile.toif
+      toif_convert.py infile.toif outfile.png
     """
-    with open(ofn_py, "wb") as f:
-        f.write(("%s = %s\n" % (bname, data)).encode())
-    with open(ofn_h, "wt") as f:
-        f.write("// clang-format off\n")
-        f.write("static const uint8_t toi_%s[] = {\n" % bname)
-        f.write("    // magic\n")
-        if im.mode == "RGB":
-            f.write("    'T', 'O', 'I', 'f',\n")
-        else:
-            f.write("    'T', 'O', 'I', 'g',\n")
-        f.write("    // width (16-bit), height (16-bit)\n")
-        f.write(
-            "    0x%02x, 0x%02x, 0x%02x, 0x%02x,\n"
-            % (w & 0xFF, w >> 8, h & 0xFF, h >> 8)
-        )
-        l = len(zdata)
-        f.write("    // compressed data length (32-bit)\n")
-        f.write(
-            "    0x%02x, 0x%02x, 0x%02x, 0x%02x,\n"
-            % (l & 0xFF, (l >> 8) & 0xFF, (l >> 16) & 0xFF, l >> 24)
-        )
-        f.write("    // compressed data\n")
-        f.write("   ")
-        for b in zdata:
-            f.write(" 0x%02x," % b)
-        f.write("\n};\n")
-    """
+    if infile.name.endswith(".toif") or infile.name == "-":
+        toi = toif.from_bytes(infile.read())
+        im = toi.to_image()
+        im.save(outfile)
 
+    elif outfile.name.endswith(".toif") or outfile.name == "-":
+        im = Image.open(infile)
+        toi = toif.from_image(im)
+        outfile.write(toi.to_bytes())
 
-def toif2png(data):
-    def process_rgb(w, h, data):
-        pix = bytearray(w * h * 3)
-        for i in range(w * h):
-            c = (data[i * 2] << 8) + data[i * 2 + 1]
-            pix[i * 3 + 0] = (c & 0xF800) >> 8
-            pix[i * 3 + 1] = (c & 0x07C0) >> 3
-            pix[i * 3 + 2] = (c & 0x001F) << 3
-        return bytes(pix)
-
-    def process_grayscale(w, h, data):
-        pix = bytearray(w * h)
-        for i in range(w * h // 2):
-            pix[i * 2 + 0] = data[i] & 0xF0
-            pix[i * 2 + 1] = (data[i] & 0x0F) << 4
-        return bytes(pix)
-
-    if data[:4] != b"TOIf" and data[:4] != b"TOIg":
-        raise ValueError("Unknown TOIF header")
-    format = data[:4].decode().lower()
-    w, h = struct.unpack("<HH", data[4:8])
-    l = struct.unpack("<I", data[8:12])[0]
-    data = data[12:]
-    if len(data) != l:
-        raise ValueError("Compressed data length mismatch (%d vs %d)" % (len(data), l))
-    data = zlib.decompress(data, -10)
-
-    if format == "toif":
-        if len(data) != w * h * 2:
-            raise ValueError(
-                "Uncompressed data length mismatch (%d vs %d)" % (len(data), w * h * 2)
-            )
-        pix = process_rgb(w, h, data)
-        img = Image.frombuffer("RGB", (w, h), pix, "raw", "RGB", 0, 1)
-        ret = io.BytesIO()
-        img.save(ret, format="PNG")
-        return ret.getvalue()
-
-    if format == "toig":
-        if len(data) != w * h // 2:
-            raise ValueError(
-                "Uncompressed data length mismatch (%d vs %d)" % (len(data), w * h // 2)
-            )
-        pix = process_grayscale(w, h, data)
-        img = Image.frombuffer("L", (w, h), pix, "raw", "L", 0, 1)
-        ret = io.BytesIO()
-        img.save(ret, format="PNG")
-        return ret.getvalue()
-
-
-def main():
-    if len(sys.argv) < 2:
-        print("Usage: toi_convert [input.png|input.toif] [output.toif|output.png]")
-        return 1
-    ifn, ofn = sys.argv[1], sys.argv[2]
-    if ifn.endswith(".png") and ofn.endswith(".toif"):
-        open(ofn, "wb").write(png2toif(open(ifn, "rb").read()))
-    elif ifn.endswith(".toif") and ofn.endswith(".png"):
-        open(ofn, "wb").write(toif2png(open(ifn, "rb").read()))
     else:
-        print("Must provide a valid combination of PNG/TOIF or TOIF/PNG files")
-        return 2
+        raise click.ClickException("At least one of the arguments must end with .toif")
 
 
 if __name__ == "__main__":
-    main()
+    toif_convert()