import unittest
from axolotl.ratchet.chainkey import ChainKey
from axolotl.kdf.hkdf import HKDF
class ChainKeyTest(unittest.TestCase):
    def test_chainKeyDerivationV2(self):
        seed = bytearray([0x8a, 0xb7, 0x2d, 0x6f, 0x4c,
                           0xc5, 0xac, 0x0d, 0x38, 0x7e,
                           0xaf, 0x46, 0x33, 0x78, 0xdd,
                           0xb2, 0x8e, 0xdd, 0x07, 0x38,
                           0x5b, 0x1c, 0xb0, 0x12, 0x50,
                           0xc7, 0x15, 0x98, 0x2e, 0x7a,
                           0xd4, 0x8f])

        messageKey   = bytearray([0x02, 0xa9, 0xaa, 0x6c, 0x7d,
                           0xbd, 0x64, 0xf9, 0xd3, 0xaa,
                           0x92, 0xf9, 0x2a, 0x27, 0x7b,
                           0xf5, 0x46, 0x09, 0xda, 0xdf,
                           0x0b, 0x00, 0x82, 0x8a, 0xcf,
                           0xc6, 0x1e, 0x3c, 0x72, 0x4b,
                           0x84, 0xa7])

        macKey       = bytearray([0xbf, 0xbe, 0x5e, 0xfb, 0x60,
                           0x30, 0x30, 0x52, 0x67, 0x42,
                           0xe3, 0xee, 0x89, 0xc7, 0x02,
                           0x4e, 0x88, 0x4e, 0x44, 0x0f,
                           0x1f, 0xf3, 0x76, 0xbb, 0x23,
                           0x17, 0xb2, 0xd6, 0x4d, 0xeb,
                           0x7c, 0x83])

        nextChainKey = bytearray([0x28, 0xe8, 0xf8, 0xfe, 0xe5,
                           0x4b, 0x80, 0x1e, 0xef, 0x7c,
                           0x5c, 0xfb, 0x2f, 0x17, 0xf3,
                           0x2c, 0x7b, 0x33, 0x44, 0x85,
                           0xbb, 0xb7, 0x0f, 0xac, 0x6e,
                           0xc1, 0x03, 0x42, 0xa2, 0x46,
                           0xd1, 0x5d])

        chainKey =  ChainKey(HKDF.createFor(2), seed, 0)
        self.assertEqual(chainKey.getKey(), seed)
        self.assertEqual(chainKey.getMessageKeys().getCipherKey(), messageKey)
        self.assertEqual(chainKey.getMessageKeys().getMacKey(), macKey)
        self.assertEqual(chainKey.getNextChainKey().getKey(), nextChainKey)
        self.assertEqual(chainKey.getIndex(), 0)
        self.assertEqual(chainKey.getMessageKeys().getCounter(), 0)
        self.assertEqual(chainKey.getNextChainKey().getIndex(), 1)
        self.assertEqual(chainKey.getNextChainKey().getMessageKeys().getCounter(), 1)

