# Copyright (c) 2011 Yubico AB
# See the file COPYING for licence statement.

import re
import sys
import os
import time
import unittest
import pyhsm
import pyhsm.util

import test_common

from StringIO import StringIO
from test_common import CfgPassphrase, AdminYubiKeys, HsmPassphrase, PrimaryAdminYubiKey


if __name__ == '__main__':
    class ConfigureYubiHSMforTest(test_common.YHSM_TestCase):

        def test_aaa_echo(self):
            """ Test echo before reconfiguration. """
            self.assertTrue(self.hsm.echo('test'))

        def test_configure_YHSM(self):
            """
            Reconfiguring YubiHSM for tests.
            """
            self.ser = self.hsm.get_raw_device()

            # get the YubiHSM to exit to configuration mode.
            #self.assertTrue(self.hsm.monitor_exit())
            self.hsm.monitor_exit()

            # get the first prompt without sending anything
            self.config_do("", add_cr = False)

            self.config_do("sysinfo")

            self.config_do("help")

            # clear memory and configure as HSM - has a few prompts we have to get past
            #
            if not self.hsm.version.have_key_store_decrypt():
                self.config_do ("hsm ffffffff\r%s\r%s\ryes" % (CfgPassphrase, HsmPassphrase))
            else:
                # HSM> < hsm ffffffff
                # Enabled flags ffffffff = ...
                # Enter cfg password (g to generate)
                # Enter admin Yubikey public id (enter when done)
                # Enter master key (g to generate) yes
                # Confirm current config being erased (type yes)
                AdminYubiKeysStr = '\r'.join(AdminYubiKeys)
                AdminYubiKeysStr += '\r'
                self.config_do ("hsm ffffffff\r%s\r%s\r%s\ryes" % (CfgPassphrase, AdminYubiKeysStr, HsmPassphrase))

            self.hsm.drain()
            self.add_keys(xrange(31))
            self.hsm.drain()

            self.config_do("keylist")

            if self.hsm.version.have_key_store_decrypt():
                self.config_do("keycommit")

            # load a YubiKey (the first Admin YubiKey) into the internal database
            escape_char = chr(27)
            self.config_do("dbload\r00001,%s,%s,%s,\r" % (PrimaryAdminYubiKey) + escape_char, add_cr = False)

            self.config_do("dblist")

            # get back into HSM mode
            sys.stderr.write("exit")
            self.ser.write("exit\r")

            self.hsm.drain()

            self.hsm.reset()

        def test_zzz_unlock(self):
            """ Test unlock of keystore after reconfiguration. """
            if self.hsm.version.have_unlock():
                Params = PrimaryAdminYubiKey
                YK = test_common.FakeYubiKey(pyhsm.yubikey.modhex_decode(Params[0]).decode('hex'),
                                            Params[1].decode('hex'), Params[2].decode('hex')
                                            )
                # After reconfigure, we know the counter values for PrimaryAdminYubiKey is zero
                # in the internal db. However, the test suite initialization will unlock the keystore
                # (in test_common.YHSM_TestCase.setUp) so a value of 0/1 should result in a replayed OTP.
                YK.use_ctr = 0
                YK.session_ctr = 1
                # first verify counters 1/0 gives the expected YSM_OTP_REPLAY
                try:
                    self.hsm.unlock(otp = YK.from_key())
                except pyhsm.exception.YHSM_CommandFailed, e:
                    if e.status != pyhsm.defines.YSM_OTP_REPLAY:
                        raise
                # now do real unlock with values 2/1 (there is an extra unlock done somewhere...)
                YK.use_ctr = 2
                self.assertTrue(self.hsm.unlock(password = HsmPassphrase.decode("hex"), otp = YK.from_key()))
            else:
                self.assertTrue(self.hsm.unlock(password = HsmPassphrase.decode("hex")))

        def test_zzz_echo(self):
            """ Test echo after reconfiguration. """
            self.assertTrue(self.hsm.echo('test'))

        def config_do(self, cmd, add_cr = True):
            # Don't have to output command - it is echoed
            #sys.__stderr__.write("> " + cmd + "\n")
            if add_cr:
                self.ser.write(cmd + "\r")
            else:
                self.ser.write(cmd)
            #time.sleep(0.5)
            recv = ''
            fail_count = 0
            sys.stderr.write("< ")
            while True:
                b = self.ser.read(1)
                if not b:
                    fail_count += 1
                    if fail_count == 5:
                        raise Exception("Did not get the next prompt", recv)
                sys.stderr.write(b)

                recv += b
                lines = recv.split('\n')
                if re.match('^(NO_CFG|WSAPI|HSM).*> .*', lines[-1]):
                    break
            return recv

        def add_keys(self, iterator):
            # Set up one key for every available flag
            for num in iterator:
                flags = 1 << num
                key = ("%02x" % (num + 1)) * 32
                self.add_key(flags, num + 1, key)

            # Set up some extra keys with the same key as the flag-keys, but other flags

            # flags YHSM_OTP_BLOB_VALIDATE (0x200) matching key 0x06 (with flags 0x20, YHSM_BLOB_GENERATE)
            flags = 0x200
            key = "06" * 32
            self.add_key(flags, 0x1000, key)

            # Key with full AES ECB capabilities
            # Enabled flags 0000e000 = YHSM_ECB_BLOCK_ENCRYPT,YHSM_ECB_BLOCK_DECRYPT,YHSM_ECB_BLOCK_DECRYPT_CMP
            flags = 0xe000
            key = "1001" * 16
            self.add_key(flags, 0x1001, key)

            # Key allowed to generate AEAD from known data (loaded into buffer), with user specified noncey
            flags = 0x4 | 0x40000000 | 0x20000000
            key = "1002" * 16
            self.add_key(flags, 0x1002, key)

            # Key with everything enabled at once
            flags = 0xffffffff
            key = "2000" * 16
            self.add_key(flags, 0x2000, key)

            # Key with everything enabled at once, and then revoked
            flags = 0xffffffff
            key = "2001" * 16
            self.add_key(flags, 0x2001, key)
            self.config_do("keydis 2001")

            # Key with NIST test vector for HMAC SHA1
            # Enabled flags 00010000 = YSM_HMAC_SHA1_GENERATE
            flags = 0x10000
            key = "303132333435363738393a3b3c3d3e3f40414243".ljust(64, '0')
            self.add_key(flags, 0x3031, key)

            # Key permitting AEAD generate with user specified nonce
            flags = 0x20000002
            key = "20000002" * 8
            self.add_key(flags, 0x20000002, key)

            # Key permitting random AEAD generate with user specified nonce
            flags = 0x20000008
            key = "20000008" * 8
            self.add_key(flags, 0x20000008, key)

        def add_key(self, flags, num, key):
            keyline = "%08x,%s\r" % (num, key)
            self.config_do("flags %04x" % (flags))
            escape_char = chr(27)
            self.config_do("keyload\r" + keyline + escape_char, add_cr = False)


    unittest.main()
