from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import (
    Cipher, algorithms, modes
)
from cryptography.exceptions import InvalidTag

import random
import unittest

TAG_LENGTH = 16

class TestAES(unittest.TestCase):
	def test_aesgcm_cav(self):
		pass

	def test_aesgcm(self):
		buf = bytes.fromhex('00000000000000000000000000000000')

		key = bytes.fromhex('00000000000000000000000000000000')

		iv = bytes.fromhex('000000000000000000000000')

		origdata = buf

		# encryption
		encryptor = Cipher(algorithms.AES(key),
		    modes.GCM(iv), backend=default_backend()
		).encryptor()

		last = True
		#if version == 'aes128gcm':
		data = encryptor.update(buf)

		self.assertEqual(data, bytes.fromhex('0388dace60b6a392f328c2b971b2fe78'))
		data += encryptor.finalize()

		self.assertEqual(encryptor.tag, bytes.fromhex('ab6e47d42cec13bdf53a67b21257bddf'))

		data += encryptor.tag

		encdata = data

		# decryption
		content = encdata
		decryptor = Cipher(algorithms.AES(key),
		    modes.GCM(iv, tag=content[-TAG_LENGTH:]),
		    backend=default_backend()
		).decryptor()
		decdata = decryptor.update(content[:-TAG_LENGTH]) + decryptor.finalize()

		self.assertEqual(origdata, decdata)

		# decryption
		content = encdata
		decryptor = Cipher(algorithms.AES(key),
		    modes.GCM(iv, tag=b'\x00' * TAG_LENGTH),
		    backend=default_backend()
		).decryptor()
		decdata = decryptor.update(content[:-TAG_LENGTH])

		self.assertRaises(InvalidTag, decryptor.finalize)