Merge branch 'pro-bootloader'

Add tests for the Nitrokey Pro bootloader support.
Allow to reconnect to the device during the test execution
(e.g. for a reinsertion).

Connected: https://github.com/Nitrokey/nitrokey-pro-firmware/issues/69

8 files changed, 211 insertions, 119 deletions

diff --git a/unittest/conftest.py b/unittest/conftest.py
index 49b4f02..17d9ef5 100644
--- a/unittest/conftest.py
+++ b/unittest/conftest.py
@@ -20,6 +20,7 @@ SPDX-License-Identifier: LGPL-3.0
 """
 
 import pytest
+import os, sys
 
 from misc import ffi, gs
 
@@ -82,47 +83,8 @@ def C(request=None):
 
 
 def get_library(request, allow_offline=False):
-    fp = '../NK_C_API.h'
-
-    declarations = []
-    with open(fp, 'r') as f:
-        declarations = f.readlines()
-
-    cnt = 0
-    a = iter(declarations)
-    for declaration in a:
-        if declaration.strip().startswith('NK_C_API') \
-                or declaration.strip().startswith('struct'):
-            declaration = declaration.replace('NK_C_API', '').strip()
-            while ');' not in declaration and '};' not in declaration:
-                declaration += (next(a)).strip()+'\n'
-            ffi.cdef(declaration, override=True)
-            cnt += 1
-    print('Imported {} declarations'.format(cnt))
-
-    C = None
-    import os, sys
-    path_build = os.path.join("..", "build")
-    paths = [
-            os.environ.get('LIBNK_PATH', None),
-            os.path.join(path_build,"libnitrokey.so"),
-            os.path.join(path_build,"libnitrokey.dylib"),
-            os.path.join(path_build,"libnitrokey.dll"),
-            os.path.join(path_build,"nitrokey.dll"),
-    ]
-    for p in paths:
-        if not p: continue
-        print("Trying " +p)
-        p = os.path.abspath(p)
-        if os.path.exists(p):
-            print("Found: "+p)
-            C = ffi.dlopen(p)
-            break
-        else:
-            print("File does not exist: " + p)
-    if not C:
-        print("No library file found")
-        sys.exit(1)
+    library_read_declarations()
+    C = library_open_lib()
 
     C.NK_set_debug_level(int(os.environ.get('LIBNK_DEBUG', 2)))
 
@@ -155,3 +117,63 @@ def get_library(request, allow_offline=False):
 
     return AttrProxy(C, "libnitrokey C")
 
+
+def library_open_lib():
+    C = None
+    path_build = os.path.join("..", "build")
+    paths = [
+        os.environ.get('LIBNK_PATH', None),
+        os.path.join(path_build, "libnitrokey.so"),
+        os.path.join(path_build, "libnitrokey.dylib"),
+        os.path.join(path_build, "libnitrokey.dll"),
+        os.path.join(path_build, "nitrokey.dll"),
+    ]
+    for p in paths:
+        if not p: continue
+        print("Trying " + p)
+        p = os.path.abspath(p)
+        if os.path.exists(p):
+            print("Found: " + p)
+            C = ffi.dlopen(p)
+            break
+        else:
+            print("File does not exist: " + p)
+    if not C:
+        print("No library file found")
+        sys.exit(1)
+    return C
+
+
+def library_read_declarations():
+    fp = '../NK_C_API.h'
+    declarations = []
+    with open(fp, 'r') as f:
+        declarations = f.readlines()
+    cnt = 0
+    a = iter(declarations)
+    for declaration in a:
+        if declaration.strip().startswith('NK_C_API') \
+                or declaration.strip().startswith('struct'):
+            declaration = declaration.replace('NK_C_API', '').strip()
+            while ');' not in declaration and '};' not in declaration:
+                declaration += (next(a)).strip() + '\n'
+            ffi.cdef(declaration, override=True)
+            cnt += 1
+    print('Imported {} declarations'.format(cnt))
+
+
+def pytest_addoption(parser):
+    parser.addoption("--run-skipped", action="store_true",
+                     help="run the tests skipped by default, e.g. adding side effects")
+
+def pytest_runtest_setup(item):
+    if 'skip_by_default' in item.keywords and not item.config.getoption("--run-skipped"):
+        pytest.skip("need --run-skipped option to run this test")
+
+
+def library_device_reconnect(C):
+    C.NK_logout()
+    C = library_open_lib()
+    C.NK_logout()
+    assert C.NK_login_auto() == 1, 'Device not found'
+    return C
\ No newline at end of file
diff --git a/unittest/constants.py b/unittest/constants.py
index 645ef6a..4047f59 100644
--- a/unittest/constants.py
+++ b/unittest/constants.py
@@ -18,12 +18,7 @@ along with libnitrokey. If not, see <http://www.gnu.org/licenses/>.
 
 SPDX-License-Identifier: LGPL-3.0
 """
-
-from misc import to_hex
-
-def bb(x):
-    return bytes(x, encoding='ascii')
-
+from misc import to_hex, bb
 
 RFC_SECRET_HR = '12345678901234567890'
 RFC_SECRET = to_hex(RFC_SECRET_HR)  # '31323334353637383930...'
@@ -39,6 +34,9 @@ class DefaultPasswords:
     USER_TEMP = b'234234234'
     UPDATE = b'12345678'
     UPDATE_TEMP = b'123update123'
+    UPDATE_LONG = b'1234567890'*2
+    UPDATE_TOO_LONG = UPDATE_LONG + b'x'
+    UPDATE_TOO_SHORT = UPDATE_LONG[:7]
 
 
 class DeviceErrorCode:
@@ -49,6 +47,7 @@ class DeviceErrorCode:
     STATUS_NOT_AUTHORIZED = 5
     STATUS_AES_DEC_FAILED = 0xa
     STATUS_UNKNOWN_ERROR = 100
+    STATUS_DISCONNECTED = 255
 
 
 class LibraryErrors:
@@ -59,4 +58,5 @@ class LibraryErrors:
 
 
 HOTP_slot_count = 3
-TOTP_slot_count = 15
\ No newline at end of file
+TOTP_slot_count = 15
+PWS_SLOT_COUNT = 16
diff --git a/unittest/helpers.py b/unittest/helpers.py
new file mode 100644
index 0000000..90c818e
--- /dev/null
+++ b/unittest/helpers.py
@@ -0,0 +1,51 @@
+from constants import DeviceErrorCode, PWS_SLOT_COUNT, DefaultPasswords
+from misc import gs, bb
+
+
+def helper_fill(str_to_fill, target_width):
+    assert target_width >= len(str_to_fill)
+    numbers = '1234567890' * 4
+    str_to_fill += numbers[:target_width - len(str_to_fill)]
+    assert len(str_to_fill) == target_width
+    return bb(str_to_fill)
+
+
+def helper_PWS_get_pass(suffix):
+    return helper_fill('pass' + suffix, 20)
+
+
+def helper_PWS_get_loginname(suffix):
+    return helper_fill('login' + suffix, 32)
+
+
+def helper_PWS_get_slotname(suffix):
+    return helper_fill('slotname' + suffix, 11)
+
+
+def helper_check_device_for_data(C):
+    assert C.NK_lock_device() == DeviceErrorCode.STATUS_OK
+    assert C.NK_enable_password_safe(DefaultPasswords.USER) == DeviceErrorCode.STATUS_OK
+
+    for i in range(0, PWS_SLOT_COUNT):
+        iss = str(i)
+        assert gs(C.NK_get_password_safe_slot_name(i)) == helper_PWS_get_slotname(iss)
+        assert gs(C.NK_get_password_safe_slot_login(i)) == helper_PWS_get_loginname(iss)
+        assert gs(C.NK_get_password_safe_slot_password(i)) == helper_PWS_get_pass(iss)
+    return True
+
+
+def helper_populate_device(C):
+    # FIXME use object with random data, and check against it
+    # FIXME generate OTP as well, and check codes against its secrets
+    assert C.NK_lock_device() == DeviceErrorCode.STATUS_OK
+    res = C.NK_enable_password_safe(DefaultPasswords.USER)
+    if res != DeviceErrorCode.STATUS_OK:
+        assert C.NK_build_aes_key(DefaultPasswords.ADMIN) == DeviceErrorCode.STATUS_OK
+        assert C.NK_enable_password_safe(DefaultPasswords.USER) == DeviceErrorCode.STATUS_OK
+
+    for i in range(0, PWS_SLOT_COUNT):
+        iss = str(i)
+        assert C.NK_write_password_safe_slot(i,
+                                             helper_PWS_get_slotname(iss), helper_PWS_get_loginname(iss),
+                                             helper_PWS_get_pass(iss)) == DeviceErrorCode.STATUS_OK
+    return True
diff --git a/unittest/misc.py b/unittest/misc.py
index e9e1753..6a0d486 100644
--- a/unittest/misc.py
+++ b/unittest/misc.py
@@ -72,3 +72,7 @@ def is_long_OTP_secret_handled(C):
 
 def has_binary_counter(C):
     return (not is_storage(C)) or (is_storage(C) and get_devices_firmware_version(C) >= 54)
+
+
+def bb(x):
+    return bytes(x, encoding='ascii')
\ No newline at end of file
diff --git a/unittest/test_multiple.py b/unittest/test_multiple.py
index 821a3b7..96b23d7 100644
--- a/unittest/test_multiple.py
+++ b/unittest/test_multiple.py
@@ -28,8 +28,8 @@ from collections import defaultdict
 from tqdm import tqdm
 
 from conftest import skip_if_device_version_lower_than
-from constants import DefaultPasswords, DeviceErrorCode, bb
-from misc import gs, wait, ffi
+from constants import DefaultPasswords, DeviceErrorCode
+from misc import gs, wait, ffi, bb
 
 pprint = pprint.PrettyPrinter(indent=4).pprint
 
diff --git a/unittest/test_pro.py b/unittest/test_pro.py
index a8df7cd..99d7b1f 100644
--- a/unittest/test_pro.py
+++ b/unittest/test_pro.py
@@ -22,9 +22,10 @@ SPDX-License-Identifier: LGPL-3.0
 import pytest
 
 from conftest import skip_if_device_version_lower_than
-from constants import DefaultPasswords, DeviceErrorCode, RFC_SECRET, bb, bbRFC_SECRET, LibraryErrors, HOTP_slot_count, \
+from constants import DefaultPasswords, DeviceErrorCode, RFC_SECRET, bbRFC_SECRET, LibraryErrors, HOTP_slot_count, \
     TOTP_slot_count
-from misc import ffi, gs, wait, cast_pointer_to_tuple, has_binary_counter
+from helpers import helper_PWS_get_slotname, helper_PWS_get_loginname, helper_PWS_get_pass
+from misc import ffi, gs, wait, cast_pointer_to_tuple, has_binary_counter, bb
 from misc import is_storage
 
 @pytest.mark.lock_device
@@ -50,37 +51,21 @@ def test_write_password_safe_slot(C):
 @pytest.mark.PWS
 @pytest.mark.slowtest
 def test_write_all_password_safe_slots_and_read_10_times(C):
-    def fill(s, wid):
-        assert wid >= len(s)
-        numbers = '1234567890'*4
-        s += numbers[:wid-len(s)]
-        assert len(s) == wid
-        return bb(s)
-
-    def get_pass(suffix):
-        return fill('pass' + suffix, 20)
-
-    def get_loginname(suffix):
-        return fill('login' + suffix, 32)
-
-    def get_slotname(suffix):
-        return fill('slotname' + suffix, 11)
-
     assert C.NK_lock_device() == DeviceErrorCode.STATUS_OK
     assert C.NK_enable_password_safe(DefaultPasswords.USER) == DeviceErrorCode.STATUS_OK
     PWS_slot_count = 16
     for i in range(0, PWS_slot_count):
         iss = str(i)
         assert C.NK_write_password_safe_slot(i,
-                                             get_slotname(iss), get_loginname(iss),
-                                             get_pass(iss)) == DeviceErrorCode.STATUS_OK
+                                             helper_PWS_get_slotname(iss), helper_PWS_get_loginname(iss),
+                                             helper_PWS_get_pass(iss)) == DeviceErrorCode.STATUS_OK
 
     for j in range(0, 10):
         for i in range(0, PWS_slot_count):
             iss = str(i)
-            assert gs(C.NK_get_password_safe_slot_name(i)) == get_slotname(iss)
-            assert gs(C.NK_get_password_safe_slot_login(i)) == get_loginname(iss)
-            assert gs(C.NK_get_password_safe_slot_password(i)) == get_pass(iss)
+            assert gs(C.NK_get_password_safe_slot_name(i)) == helper_PWS_get_slotname(iss)
+            assert gs(C.NK_get_password_safe_slot_login(i)) == helper_PWS_get_loginname(iss)
+            assert gs(C.NK_get_password_safe_slot_password(i)) == helper_PWS_get_pass(iss)
 
 
 @pytest.mark.lock_device
@@ -88,22 +73,6 @@ def test_write_all_password_safe_slots_and_read_10_times(C):
 @pytest.mark.slowtest
 @pytest.mark.xfail(reason="This test should be run directly after test_write_all_password_safe_slots_and_read_10_times")
 def test_read_all_password_safe_slots_10_times(C):
-    def fill(s, wid):
-        assert wid >= len(s)
-        numbers = '1234567890'*4
-        s += numbers[:wid-len(s)]
-        assert len(s) == wid
-        return bb(s)
-
-    def get_pass(suffix):
-        return fill('pass' + suffix, 20)
-
-    def get_loginname(suffix):
-        return fill('login' + suffix, 32)
-
-    def get_slotname(suffix):
-        return fill('slotname' + suffix, 11)
-
     assert C.NK_lock_device() == DeviceErrorCode.STATUS_OK
     assert C.NK_enable_password_safe(DefaultPasswords.USER) == DeviceErrorCode.STATUS_OK
     PWS_slot_count = 16
@@ -111,9 +80,9 @@ def test_read_all_password_safe_slots_10_times(C):
     for j in range(0, 10):
         for i in range(0, PWS_slot_count):
             iss = str(i)
-            assert gs(C.NK_get_password_safe_slot_name(i)) == get_slotname(iss)
-            assert gs(C.NK_get_password_safe_slot_login(i)) == get_loginname(iss)
-            assert gs(C.NK_get_password_safe_slot_password(i)) == get_pass(iss)
+            assert gs(C.NK_get_password_safe_slot_name(i)) == helper_PWS_get_slotname(iss)
+            assert gs(C.NK_get_password_safe_slot_login(i)) == helper_PWS_get_loginname(iss)
+            assert gs(C.NK_get_password_safe_slot_password(i)) == helper_PWS_get_pass(iss)
 
 
 @pytest.mark.lock_device
@@ -977,31 +946,6 @@ def test_get_device_model(C):
     # assert C.NK_get_device_model() != C.NK_DISCONNECTED
 
 
-@pytest.mark.firmware
-def test_bootloader_password_change_pro(C):
-    skip_if_device_version_lower_than({'P': 11})
-    assert C.NK_change_firmware_password_pro(b'zxcasd', b'zxcasd') == DeviceErrorCode.WRONG_PASSWORD
-
-    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_TEMP) == DeviceErrorCode.STATUS_OK
-    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE_TEMP, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
-
-
-@pytest.mark.firmware
-def test_bootloader_run_pro(C):
-    skip_if_device_version_lower_than({'P': 11})
-    assert C.NK_enable_firmware_update_pro(DefaultPasswords.UPDATE_TEMP) == DeviceErrorCode.WRONG_PASSWORD
-    # Not enabled due to lack of side-effect removal at this point
-    # assert C.NK_enable_firmware_update_pro(DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
-
-
-@pytest.mark.firmware
-def test_bootloader_password_change_pro_too_long(C):
-    skip_if_device_version_lower_than({'P': 11})
-    long_string = b'a' * 100
-    assert C.NK_change_firmware_password_pro(long_string, long_string) == LibraryErrors.TOO_LONG_STRING
-    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, long_string) == LibraryErrors.TOO_LONG_STRING
-
-
 @pytest.mark.otp
 @pytest.mark.parametrize('counter_mid', [10**3-1, 10**4-1, 10**7-1, 10**8-10, 2**16, 2**31-1, 2**32-1, 2**33, 2**50, 2**60, 2**63])  # 2**64-1
 def test_HOTP_counter_getter(C, counter_mid: int):
diff --git a/unittest/test_pro_bootloader.py b/unittest/test_pro_bootloader.py
new file mode 100644
index 0000000..4cb7470
--- /dev/null
+++ b/unittest/test_pro_bootloader.py
@@ -0,0 +1,71 @@
+import pytest
+
+from conftest import skip_if_device_version_lower_than, library_device_reconnect
+from constants import DefaultPasswords, DeviceErrorCode, LibraryErrors
+from helpers import helper_populate_device, helper_check_device_for_data
+
+
+@pytest.mark.firmware
+def test_bootloader_password_change_pro_length(C):
+    skip_if_device_version_lower_than({'P': 11})
+
+    # Test whether the correct password is set
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    # Change to the longest possible password
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_LONG) == DeviceErrorCode.STATUS_OK
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE_LONG, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    # Use longer or shorter passwords than possible
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_TOO_LONG) == LibraryErrors.TOO_LONG_STRING
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_TOO_SHORT) == DeviceErrorCode.WRONG_PASSWORD
+
+
+
+@pytest.mark.firmware
+def test_bootloader_password_change_pro(C):
+    skip_if_device_version_lower_than({'P': 11})
+    assert C.NK_change_firmware_password_pro(b'zxcasd', b'zxcasd') == DeviceErrorCode.WRONG_PASSWORD
+
+    # Revert effects of broken test run, if needed
+    C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE_TEMP, DefaultPasswords.UPDATE)
+
+    # Change to the same password
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    # Change password
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_TEMP) == DeviceErrorCode.STATUS_OK
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE_TEMP, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+
+
+@pytest.mark.firmware
+def test_bootloader_run_pro_wrong_password(C):
+    skip_if_device_version_lower_than({'P': 11})
+    assert C.NK_enable_firmware_update_pro(DefaultPasswords.UPDATE_TEMP) == DeviceErrorCode.WRONG_PASSWORD
+
+
+@pytest.mark.skip_by_default
+@pytest.mark.firmware
+def test_bootloader_run_pro_real(C):
+    skip_if_device_version_lower_than({'P': 11})
+    # Not enabled due to lack of side-effect removal at this point
+    assert C.NK_enable_firmware_update_pro(DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_DISCONNECTED
+
+
+@pytest.mark.firmware
+def test_bootloader_password_change_pro_too_long(C):
+    skip_if_device_version_lower_than({'P': 11})
+    long_string = b'a' * 100
+    assert C.NK_change_firmware_password_pro(long_string, long_string) == LibraryErrors.TOO_LONG_STRING
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, long_string) == LibraryErrors.TOO_LONG_STRING
+
+
+@pytest.mark.skip_by_default
+@pytest.mark.firmware
+def test_bootloader_data_rention(C):
+    skip_if_device_version_lower_than({'P': 11})
+
+    assert helper_populate_device(C)
+    assert C.NK_enable_firmware_update_pro(DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_DISCONNECTED
+    input('Please press ENTER after uploading new firmware to the device')
+    C = library_device_reconnect(C)
+    assert helper_check_device_for_data(C)
+
diff --git a/unittest/test_storage.py b/unittest/test_storage.py
index 0f960cc..a435a15 100644
--- a/unittest/test_storage.py
+++ b/unittest/test_storage.py
@@ -23,8 +23,8 @@ import pprint
 import pytest
 
 from conftest import skip_if_device_version_lower_than
-from constants import DefaultPasswords, DeviceErrorCode, bb
-from misc import gs, wait, ffi
+from constants import DefaultPasswords, DeviceErrorCode
+from misc import gs, wait, ffi, bb
 
 pprint = pprint.PrettyPrinter(indent=4).pprint