From 4693c071b44ecb9820ebfa34477fcc095f3fcc05 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Tue, 13 Nov 2018 14:29:40 +0100
Subject: [PATCH] src/apps/wallet/sign_tx: implement OMNI parsing in OP_RETURN
 layout

---
 src/apps/wallet/sign_tx/layout.py     | 19 +++++++++++-----
 src/apps/wallet/sign_tx/omni.py       | 24 ++++++++++++++++++++
 tests/test_apps.wallet.signtx.omni.py | 32 +++++++++++++++++++++++++++
 3 files changed, 69 insertions(+), 6 deletions(-)
 create mode 100644 src/apps/wallet/sign_tx/omni.py
 create mode 100644 tests/test_apps.wallet.signtx.omni.py

diff --git a/src/apps/wallet/sign_tx/layout.py b/src/apps/wallet/sign_tx/layout.py
index 76685885c2..217b4057c6 100644
--- a/src/apps/wallet/sign_tx/layout.py
+++ b/src/apps/wallet/sign_tx/layout.py
@@ -6,7 +6,7 @@ from trezor.ui.text import Text
 from trezor.utils import chunks, format_amount
 
 from apps.common.confirm import confirm, hold_to_confirm
-from apps.wallet.sign_tx import addresses
+from apps.wallet.sign_tx import addresses, omni
 
 
 def format_coin_amount(amount, coin):
@@ -23,11 +23,18 @@ def split_op_return(data):
 
 async def confirm_output(ctx, output, coin):
     if output.script_type == OutputScriptType.PAYTOOPRETURN:
-        data = hexlify(output.op_return_data).decode()
-        if len(data) >= 18 * 5:
-            data = data[: (18 * 5 - 3)] + "..."
-        text = Text("OP_RETURN", ui.ICON_SEND, icon_color=ui.GREEN)
-        text.mono(*split_op_return(data))
+        data = output.op_return_data
+        if omni.is_valid(data):
+            # OMNI transaction
+            text = Text("OMNI transaction", ui.ICON_SEND, icon_color=ui.GREEN)
+            text.normal(omni.parse(data))
+        else:
+            # generic OP_RETURN
+            data = hexlify(data).decode()
+            if len(data) >= 18 * 5:
+                data = data[: (18 * 5 - 3)] + "..."
+            text = Text("OP_RETURN", ui.ICON_SEND, icon_color=ui.GREEN)
+            text.mono(*split_op_return(data))
     else:
         address = output.address
         address_short = addresses.address_short(coin, address)
diff --git a/src/apps/wallet/sign_tx/omni.py b/src/apps/wallet/sign_tx/omni.py
new file mode 100644
index 0000000000..0b6f63d7ae
--- /dev/null
+++ b/src/apps/wallet/sign_tx/omni.py
@@ -0,0 +1,24 @@
+from ustruct import unpack
+
+from trezor.utils import format_amount
+
+currencies = {1: "OMNI", 2: "tOMNI", 3: "MAID", 31: "USDT"}
+
+
+def is_valid(data: bytes) -> bool:
+    return data[:4] == b"omni"
+
+
+def parse(data: bytes) -> bool:
+    if not is_valid(data):
+        return None
+    tx_version, tx_type = unpack(">HH", data[4:8])
+    if tx_version == 0 and tx_type == 0 and len(data) == 20:  # OMNI simple send
+        currency, amount = unpack(">IQ", data[8:20])
+        return "Simple send of %s %s" % (
+            format_amount(amount, 8),
+            currencies.get(currency, "UNKN"),
+        )
+    else:
+        # unknown OMNI transaction
+        return "Unknown transaction"
diff --git a/tests/test_apps.wallet.signtx.omni.py b/tests/test_apps.wallet.signtx.omni.py
new file mode 100644
index 0000000000..4d1c26ece9
--- /dev/null
+++ b/tests/test_apps.wallet.signtx.omni.py
@@ -0,0 +1,32 @@
+from common import *
+
+from apps.wallet.sign_tx.omni import is_valid, parse
+
+class TestSignTxOmni(unittest.TestCase):
+
+    def test_is_valid(self):
+        VECTORS = {
+             "6f6d6e69000000000000001f000000002b752ee0": True,
+             "6f6d6e69000000000000001f0000000020c85580": True,
+             "0f6d6e69000000000000001f0000000020c85580": False,
+             "6f6d6e69000000000000001f0000000020c8558000": True,
+             "6f6d6e69000000000000001f0000000020c855": True,
+        }
+        for k, v in VECTORS.items():
+            k = unhexlify(k)
+            self.assertEqual(is_valid(k), v)
+
+    def test_parse(self):
+        VECTORS = {
+             "6f6d6e69000000000000001f000000002b752ee0": "Simple send of 7.291 USDT",
+             "6f6d6e69000000000000001f0000000020c85580": "Simple send of 5.5 USDT",
+             "6f6d6e690000000000000003000000002b752ee0": "Simple send of 7.291 MAID",
+             "6f6d6e690000000000000000000000002b752ee0": "Simple send of 7.291 UNKN",
+             "6f6d6e6901000000": "Unknown transaction",
+        }
+        for k, v in VECTORS.items():
+            k = unhexlify(k)
+            self.assertEqual(parse(k), v)
+
+if __name__ == '__main__':
+    unittest.main()