from unittest import TestCase, main
import aes, gcm

def bytes(hexbytes):
    return [int(hexbyte, 16) for hexbyte in hexbytes.split()]

def chars(hexbytes):
    return ''.join(map(chr, bytes(hexbytes)))

class TestGCM(TestCase):
    def test_gcm(self):
        # Appendix B of http://csrc.nist.gov/
        # CryptoToolkit/modes/proposedmodes/gcm/gcm-revised-spec.pdf
        equal = self.assertEqual

        K = 0
        P = ''
        C = ''
        T = chars('58 e2 fc ce fa 7e 30 61 36 7f 1d 57 a4 e7 45 5a')
        equal(gcm.GCM(aes.AES(128, K), 0).encrypt(P), (C, T))
        equal(gcm.GCM(aes.AES(128, K), 0).decrypt(C, T), P)

        K = 0
        P = chars('00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00')
        C = chars('03 88 da ce 60 b6 a3 92 f3 28 c2 b9 71 b2 fe 78')
        T = chars('ab 6e 47 d4 2c ec 13 bd f5 3a 67 b2 12 57 bd df')
        equal(gcm.GCM(aes.AES(128, K), 0).encrypt(P), (C, T))
        equal(gcm.GCM(aes.AES(128, K), 0).decrypt(C, T), P)

        K = 0xfeffe9928665731c6d6a8f9467308308
        iv = 0xcafebabefacedbaddecaf888
        P = chars('''d9 31 32 25 f8 84 06 e5 a5 59 09 c5 af f5 26 9a
                     86 a7 a9 53 15 34 f7 da 2e 4c 30 3d 8a 31 8a 72
                     1c 3c 0c 95 95 68 09 53 2f cf 0e 24 49 a6 b5 25 
                     b1 6a ed f5 aa 0d e6 57 ba 63 7b 39 1a af d2 55''')
        C = chars('''42 83 1e c2 21 77 74 24 4b 72 21 b7 84 d0 d4 9c 
                     e3 aa 21 2f 2c 02 a4 e0 35 c1 7e 23 29 ac a1 2e 
                     21 d5 14 b2 54 66 93 1c 7d 8f 6a 5a ac 84 aa 05 
                     1b a3 0b 39 6a 0a ac 97 3d 58 e0 91 47 3f 59 85''')
        T = chars('4d 5c 2a f3 27 cd 64 a6 2c f3 5a bd 2b a6 fa b4')
        equal(gcm.GCM(aes.AES(128, K), iv).encrypt(P), (C, T))
        equal(gcm.GCM(aes.AES(128, K), iv).decrypt(C, T), P)

if __name__ == '__main__':
    main()