1# This file is part of the TREZOR project.
2#
3# Copyright (C) 2012-2016 Marek Palatinus <slush@satoshilabs.com>
4# Copyright (C) 2012-2016 Pavol Rusnak <stick@satoshilabs.com>
5#
6# This library is free software: you can redistribute it and/or modify
7# it under the terms of the GNU Lesser General Public License as published by
8# the Free Software Foundation, either version 3 of the License, or
9# (at your option) any later version.
10#
11# This library is distributed in the hope that it will be useful,
12# but WITHOUT ANY WARRANTY; without even the implied warranty of
13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14# GNU Lesser General Public License for more details.
15#
16# You should have received a copy of the GNU Lesser General Public License
17# along with this library.  If not, see <http://www.gnu.org/licenses/>.
18#
19# The script has been modified for KeepKey Device.
20
21from __future__ import print_function
22
23import unittest
24import config
25import time
26import semver
27
28from keepkeylib.client import KeepKeyClient, KeepKeyDebuglinkClient, KeepKeyDebuglinkClientVerbose
29from keepkeylib import tx_api
30
31tx_api.cache_dir = 'txcache'
32VERBOSE = False
33
34class KeepKeyTest(unittest.TestCase):
35    def setUp(self):
36        transport = config.TRANSPORT(*config.TRANSPORT_ARGS, **config.TRANSPORT_KWARGS)
37        if hasattr(config, 'DEBUG_TRANSPORT'):
38            debug_transport = config.DEBUG_TRANSPORT(*config.DEBUG_TRANSPORT_ARGS, **config.DEBUG_TRANSPORT_KWARGS)
39            if VERBOSE:
40                self.client = KeepKeyDebuglinkClientVerbose(transport)
41            else:
42                self.client = KeepKeyDebuglinkClient(transport)
43            self.client.set_debuglink(debug_transport)
44        else:
45            self.client = KeepKeyClient(transport)
46        self.client.set_tx_api(tx_api.TxApiBitcoin)
47        # self.client.set_buttonwait(3)
48
49        #                     1      2     3    4      5      6      7     8      9    10    11    12
50        self.mnemonic12 = 'alcohol woman abuse must during monitor noble actual mixed trade anger aisle'
51        self.mnemonic18 = 'owner little vague addict embark decide pink prosper true fork panda embody mixture exchange choose canoe electric jewel'
52        self.mnemonic24 = 'dignity pass list indicate nasty swamp pool script soccer toe leaf photo multiply desk host tomato cradle drill spread actor shine dismiss champion exotic'
53        self.mnemonic20007 = 'fix spot clown mobile oven eagle pond arrest opera buyer muffin myself'
54        self.mnemonic_all = ' '.join(['all'] * 12)
55        self.mnemonic_abandon = ' '.join(['abandon'] * 11) + ' about'
56
57        self.pin4 = '1234'
58        self.pin6 = '789456'
59        self.pin8 = '45678978'
60
61        self.client.wipe_device()
62
63        if VERBOSE:
64            print("Setup finished")
65            print("--------------")
66
67    def setup_mnemonic_allallall(self):
68        self.client.load_device_by_mnemonic(mnemonic=self.mnemonic_all, pin='', passphrase_protection=False, label='test', language='english')
69
70    def setup_mnemonic_abandon(self):
71        self.client.load_device_by_mnemonic(mnemonic=self.mnemonic_abandon, pin='', passphrase_protection=False, label='test', language='english')
72
73    def setup_mnemonic_nopin_nopassphrase(self):
74        self.client.load_device_by_mnemonic(mnemonic=self.mnemonic12, pin='', passphrase_protection=False, label='test', language='english')
75
76    def setup_mnemonic_vuln20007(self):
77        self.client.load_device_by_mnemonic(mnemonic=self.mnemonic20007, pin='', passphrase_protection=False, label='test', language='english')
78
79    def setup_mnemonic_pin_nopassphrase(self):
80        self.client.load_device_by_mnemonic(mnemonic=self.mnemonic12, pin=self.pin4, passphrase_protection=False, label='test', language='english')
81
82    def setup_mnemonic_pin_passphrase(self):
83        self.client.load_device_by_mnemonic(mnemonic=self.mnemonic12, pin=self.pin4, passphrase_protection=True, label='test', language='english')
84
85    def tearDown(self):
86        self.client.close()
87
88    def assertEqual(self, lhs, rhs):
89        if type(lhs) == type(b'') and type(rhs) == type(''):
90            super(KeepKeyTest, self).assertEqual(lhs, rhs.encode('utf-8'))
91        else:
92            super(KeepKeyTest, self).assertEqual(lhs, rhs)
93
94    def assertEndsWith(self, s, suffix):
95        self.assertTrue(s.endswith(suffix), "'{}'.endswith('{}')".format(s, suffix))
96
97    def requires_firmware(self, ver_required):
98        self.client.init_device()
99        features = self.client.features
100        version = "%s.%s.%s" % (features.major_version, features.minor_version, features.patch_version)
101        if semver.compare(version, ver_required) < 0:
102            self.skipTest("Firmware version " + ver_required + " or higher is required to run this test")
103
104