import sys
import time
import unittest
import serveronstick
import serveronstick.util

import test_common

from StringIO import StringIO

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()

        # clear memory and configure as HSM - has a few prompts we have to get past
        self.config_do("hsm ffffffff\r\r\ryes")

        self.config_do("sysinfo")

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

        self.config_do("keylist")

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

        self.hsm.drain()

        self.hsm.reset()

    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
            if recv.endswith("NO_CFG> "):
                break
            if recv.endswith("HSM> "):
                break
        return recv

    def add_keys(self, iterator):
        for num in iterator:
            flag = 1 << num
            key = ("%02x" % (num + 1)) * 32
            keyline = "%d,%s\r" % (num + 1, key)
            self.config_do("flags %04x" % (flag))
            escape_char = chr(27)
            self.config_do("keyload\r" + keyline + escape_char, add_cr = False)

@@ -1969,16 +1969,66 @@
 rator):%0A
+        # Set up one key for every available flag%0A
         
@@ -2064,16 +2064,17 @@
     flag
+s
  = 1 %3C%3C 
@@ -2137,21 +2137,379 @@
     
-keyline = %22%25d
+self.add_key(flags, num + 1, key)%0A%0A        # Set up some extra keys with the same key as the flag-keys, but other flags -%0A%0A        # flags SOS_OTP_BLOB_VALIDATE (0x200) matching key 0x09 (with flags 0x100, SOS_BLOB_STORE)%0A        flags = 0x200%0A        key = %2209%22 * 32%0A        self.add_key(flags, 0x1000, key)%0A%0A    def add_key(self, flags, num, key):%0A        keyline = %22%25x
 ,%25s%5C
@@ -2517,20 +2517,16 @@
 %22 %25 (num
- + 1
 , key)%0A 
@@ -2524,36 +2524,32 @@
 m, key)%0A        
-    
 self.config_do(%22
@@ -2567,16 +2567,17 @@
  %25 (flag
+s
 ))%0A     
@@ -2579,20 +2579,16 @@
         
-  
-  
 escape_c
@@ -2601,20 +2601,16 @@
 chr(27)%0A
-    
         

Set up an OTP validation key.