From a1cd0b094988ebee01b6d3bd3c6a8754cb6c1667 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Wed, 3 Feb 2021 13:39:26 +0100
Subject: [PATCH] fix(tests): update shamir_mnemonic usage

---
 tests/click_tests/reset.py                      |  9 ++++-----
 .../test_msg_resetdevice_slip39_advanced.py     |  9 ++++-----
 .../test_msg_resetdevice_slip39_basic.py        | 11 ++++-------
 tests/device_tests/test_reset_backup.py         | 17 +++++++----------
 4 files changed, 19 insertions(+), 27 deletions(-)

diff --git a/tests/click_tests/reset.py b/tests/click_tests/reset.py
index 2ff682b03..550e793cc 100644
--- a/tests/click_tests/reset.py
+++ b/tests/click_tests/reset.py
@@ -1,4 +1,4 @@
-import shamir_mnemonic as shamir
+from shamir_mnemonic import shamir
 
 from trezorlib import messages
 
@@ -72,7 +72,6 @@ def confirm_words(debug, words):
 def validate_mnemonics(mnemonics, expected_ems):
     # We expect these combinations to recreate the secret properly
     # In case of click tests the mnemonics are always XofX so no need for combinations
-    ms = shamir.combine_mnemonics(mnemonics)
-    identifier, iteration_exponent, _, _, _ = shamir._decode_mnemonics(mnemonics)
-    ems = shamir._encrypt(ms, b"", iteration_exponent, identifier)
-    assert ems == expected_ems
+    groups = shamir.decode_mnemonics(mnemonics)
+    ems = shamir.recover_ems(groups)
+    assert expected_ems == ems.ciphertext
diff --git a/tests/device_tests/test_msg_resetdevice_slip39_advanced.py b/tests/device_tests/test_msg_resetdevice_slip39_advanced.py
index a36ab725e..cb2fbb21b 100644
--- a/tests/device_tests/test_msg_resetdevice_slip39_advanced.py
+++ b/tests/device_tests/test_msg_resetdevice_slip39_advanced.py
@@ -17,7 +17,7 @@
 from unittest import mock
 
 import pytest
-import shamir_mnemonic as shamir
+from shamir_mnemonic import shamir
 
 from trezorlib import device, messages as proto
 from trezorlib.exceptions import TrezorFailure
@@ -186,7 +186,6 @@ def validate_mnemonics(mnemonics, threshold, expected_ems):
     # 3of5 shares 3of5 groups
     # TODO: test all possible group+share combinations?
     test_combination = mnemonics[0:3] + mnemonics[5:8] + mnemonics[10:13]
-    ms = shamir.combine_mnemonics(test_combination)
-    identifier, iteration_exponent, _, _, _ = shamir._decode_mnemonics(test_combination)
-    ems = shamir._encrypt(ms, b"", iteration_exponent, identifier)
-    assert ems == expected_ems
+    groups = shamir.decode_mnemonics(test_combination)
+    ems = shamir.recover_ems(groups)
+    assert expected_ems == ems.ciphertext
diff --git a/tests/device_tests/test_msg_resetdevice_slip39_basic.py b/tests/device_tests/test_msg_resetdevice_slip39_basic.py
index a600f11d3..4482379d3 100644
--- a/tests/device_tests/test_msg_resetdevice_slip39_basic.py
+++ b/tests/device_tests/test_msg_resetdevice_slip39_basic.py
@@ -18,8 +18,7 @@ from itertools import combinations
 from unittest import mock
 
 import pytest
-import shamir_mnemonic as shamir
-from shamir_mnemonic import MnemonicError
+from shamir_mnemonic import MnemonicError, shamir
 
 from trezorlib import device, messages as proto
 from trezorlib.exceptions import TrezorFailure
@@ -144,11 +143,9 @@ class TestMsgResetDeviceT2:
 def validate_mnemonics(mnemonics, threshold, expected_ems):
     # We expect these combinations to recreate the secret properly
     for test_group in combinations(mnemonics, threshold):
-        # TODO: HOTFIX, we should fix this properly by modifying and unifying the python-shamir-mnemonic API
-        ms = shamir.combine_mnemonics(test_group)
-        identifier, iteration_exponent, _, _, _ = shamir._decode_mnemonics(test_group)
-        ems = shamir._encrypt(ms, b"", iteration_exponent, identifier)
-        assert ems == expected_ems
+        groups = shamir.decode_mnemonics(test_group)
+        ems = shamir.recover_ems(groups)
+        assert expected_ems == ems.ciphertext
     # We expect these combinations to raise MnemonicError
     for test_group in combinations(mnemonics, threshold - 1):
         with pytest.raises(
diff --git a/tests/device_tests/test_reset_backup.py b/tests/device_tests/test_reset_backup.py
index 75b087c20..64069d6dc 100644
--- a/tests/device_tests/test_reset_backup.py
+++ b/tests/device_tests/test_reset_backup.py
@@ -18,7 +18,7 @@
 from unittest import mock
 
 import pytest
-import shamir_mnemonic as shamir
+from shamir_mnemonic import shamir
 
 from trezorlib import device, messages
 from trezorlib.messages import BackupType, ButtonRequestType as B
@@ -108,11 +108,9 @@ def backup_flow_slip39_basic(client):
         )
         device.backup(client)
 
-    mnemonics = mnemonics[:3]
-    ms = shamir.combine_mnemonics(mnemonics)
-    identifier, iteration_exponent, _, _, _ = shamir._decode_mnemonics(mnemonics)
-    secret = shamir._encrypt(ms, b"", iteration_exponent, identifier)
-    return secret
+    groups = shamir.decode_mnemonics(mnemonics[:3])
+    ems = shamir.recover_ems(groups)
+    return ems.ciphertext
 
 
 def backup_flow_slip39_advanced(client):
@@ -172,10 +170,9 @@ def backup_flow_slip39_advanced(client):
         device.backup(client)
 
     mnemonics = mnemonics[0:3] + mnemonics[5:8] + mnemonics[10:13]
-    ms = shamir.combine_mnemonics(mnemonics)
-    identifier, iteration_exponent, _, _, _ = shamir._decode_mnemonics(mnemonics)
-    secret = shamir._encrypt(ms, b"", iteration_exponent, identifier)
-    return secret
+    groups = shamir.decode_mnemonics(mnemonics)
+    ems = shamir.recover_ems(groups)
+    return ems.ciphertext
 
 
 VECTORS = [