From 2ee57da5d873cc47145571eb75a7c799d7dab471 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Wed, 4 Jul 2018 18:08:05 +0200
Subject: [PATCH] src/apps/wallet/sign_tx: fix last commit, add bip115 test

---
 src/apps/wallet/sign_tx/scripts.py       |  2 +-
 src/apps/wallet/sign_tx/writers.py       |  2 +-
 tests/test_apps.wallet.signtx.scripts.py | 20 ++++++++++++++++++++
 3 files changed, 22 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_apps.wallet.signtx.scripts.py

diff --git a/src/apps/wallet/sign_tx/scripts.py b/src/apps/wallet/sign_tx/scripts.py
index d2bf850272..8fd5ed2bda 100644
--- a/src/apps/wallet/sign_tx/scripts.py
+++ b/src/apps/wallet/sign_tx/scripts.py
@@ -1,7 +1,7 @@
 from trezor.crypto.hashlib import ripemd160, sha256
 from trezor.messages.MultisigRedeemScriptType import MultisigRedeemScriptType
 from apps.wallet.sign_tx.multisig import multisig_get_pubkeys
-from apps.wallet.sign_tx.writers import bytearray_with_cap, write_bytes, write_varint, write_op_push
+from apps.wallet.sign_tx.writers import bytearray_with_cap, write_bytes, write_varint, write_op_push, write_scriptnum
 
 
 class ScriptsError(ValueError):
diff --git a/src/apps/wallet/sign_tx/writers.py b/src/apps/wallet/sign_tx/writers.py
index 634cbb1873..a59e72a7eb 100644
--- a/src/apps/wallet/sign_tx/writers.py
+++ b/src/apps/wallet/sign_tx/writers.py
@@ -85,7 +85,7 @@ def write_scriptnum(w, n: int):
         w.append(n & 0xFF)
         w.append((n >> 8) & 0xFF)
         w.append((n >> 16) & 0xFF)
-    elif:
+    else:
         w.append(4)
         w.append(n & 0xFF)
         w.append((n >> 8) & 0xFF)
diff --git a/tests/test_apps.wallet.signtx.scripts.py b/tests/test_apps.wallet.signtx.scripts.py
new file mode 100644
index 0000000000..9ec833cfc3
--- /dev/null
+++ b/tests/test_apps.wallet.signtx.scripts.py
@@ -0,0 +1,20 @@
+from common import *
+
+from apps.wallet.sign_tx.scripts import script_replay_protection_bip115
+
+class TestSigntxScripts(unittest.TestCase):
+    # pylint: disable=C0301
+
+    def test_script_replay_protection_bip115(self):
+        vectors=[
+            ('206ec9b310745775c20cbe5bae8751daeb7f086cf913399d4f7634ef2a0000000003122005b4', '6ec9b310745775c20cbe5bae8751daeb7f086cf913399d4f7634ef2a00000000', 335890),
+            ('20caaa71b60cf893c1604b38e5af1bdc322dbb31818239088647272d1400000000030e2005b4', 'caaa71b60cf893c1604b38e5af1bdc322dbb31818239088647272d1400000000', 335886),
+        ]
+        for out, hsh, height in vectors:
+            hsh = unhexlify(hsh)
+            res = hexlify(script_replay_protection_bip115(hsh, height)).decode()
+            self.assertEqual(out, res)
+
+
+if __name__ == '__main__':
+    unittest.main()