#!/usr/bin/env python

from __future__ import unicode_literals

# Allow direct execution
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from haruhi_dl.aes import aes_decrypt, aes_encrypt, aes_cbc_decrypt, aes_cbc_encrypt, aes_decrypt_text
from haruhi_dl.utils import bytes_to_intlist, intlist_to_bytes
import base64

# the encrypted data can be generate with 'devscripts/generate_aes_testdata.py'


class TestAES(unittest.TestCase):
    def setUp(self):
        self.key = self.iv = [0x20, 0x15] + 14 * [0]
        self.secret_msg = b'Secret message goes here'

    def test_encrypt(self):
        msg = b'message'
        key = list(range(16))
        encrypted = aes_encrypt(bytes_to_intlist(msg), key)
        decrypted = intlist_to_bytes(aes_decrypt(encrypted, key))
        self.assertEqual(decrypted, msg)

    def test_cbc_decrypt(self):
        data = bytes_to_intlist(
            b"\x97\x92+\xe5\x0b\xc3\x18\x91ky9m&\xb3\xb5@\xe6'\xc2\x96.\xc8u\x88\xab9-[\x9e|\xf1\xcd"
        )
        decrypted = intlist_to_bytes(aes_cbc_decrypt(data, self.key, self.iv))
        self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)

    def test_cbc_encrypt(self):
        data = bytes_to_intlist(self.secret_msg)
        encrypted = intlist_to_bytes(aes_cbc_encrypt(data, self.key, self.iv))
        self.assertEqual(
            encrypted,
            b"\x97\x92+\xe5\x0b\xc3\x18\x91ky9m&\xb3\xb5@\xe6'\xc2\x96.\xc8u\x88\xab9-[\x9e|\xf1\xcd")

    def test_decrypt_text(self):
        password = intlist_to_bytes(self.key).decode('utf-8')
        encrypted = base64.b64encode(
            intlist_to_bytes(self.iv[:8])
            + b'\x17\x15\x93\xab\x8d\x80V\xcdV\xe0\t\xcdo\xc2\xa5\xd8ksM\r\xe27N\xae'
        ).decode('utf-8')
        decrypted = (aes_decrypt_text(encrypted, password, 16))
        self.assertEqual(decrypted, self.secret_msg)

        password = intlist_to_bytes(self.key).decode('utf-8')
        encrypted = base64.b64encode(
            intlist_to_bytes(self.iv[:8])
            + b'\x0b\xe6\xa4\xd9z\x0e\xb8\xb9\xd0\xd4i_\x85\x1d\x99\x98_\xe5\x80\xe7.\xbf\xa5\x83'
        ).decode('utf-8')
        decrypted = (aes_decrypt_text(encrypted, password, 32))
        self.assertEqual(decrypted, self.secret_msg)


if __name__ == '__main__':
    unittest.main()

[((63, 4, 63, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((31, 15, 33, 9), 'haruhi_dl.utils.bytes_to_intlist', 'bytes_to_intlist', ({(32, 12, 32, 101): 'b"\\x97\\x92+\\xe5\\x0b\\xc3\\x18\\x91ky9m&\\xb3\\xb5@\\xe6\'\\xc2\\x96.\\xc8u\\x88\\xab9-[\\x9e|\\xf1\\xcd"'}, {}), '(\n    b"\\x97\\x92+\\xe5\\x0b\\xc3\\x18\\x91ky9m&\\xb3\\xb5@\\xe6\'\\xc2\\x96.\\xc8u\\x88\\xab9-[\\x9e|\\xf1\\xcd"\n    )', False, 'from haruhi_dl.utils import bytes_to_intlist, intlist_to_bytes\n'), ((38, 15, 38, 48), 'haruhi_dl.utils.bytes_to_intlist', 'bytes_to_intlist', ({(38, 32, 38, 47): 'self.secret_msg'}, {}), '(self.secret_msg)', False, 'from haruhi_dl.utils import bytes_to_intlist, intlist_to_bytes\n'), ((50, 21, 50, 62), 'haruhi_dl.aes.aes_decrypt_text', 'aes_decrypt_text', ({(50, 38, 50, 47): 'encrypted', (50, 49, 50, 57): 'password', (50, 59, 50, 61): '16'}, {}), '(encrypted, password, 16)', False, 'from haruhi_dl.aes import aes_decrypt, aes_encrypt, aes_cbc_decrypt, aes_cbc_encrypt, aes_decrypt_text\n'), ((58, 21, 58, 62), 'haruhi_dl.aes.aes_decrypt_text', 'aes_decrypt_text', ({(58, 38, 58, 47): 'encrypted', (58, 49, 58, 57): 'password', (58, 59, 58, 61): '32'}, {}), '(encrypted, password, 32)', False, 'from haruhi_dl.aes import aes_decrypt, aes_encrypt, aes_cbc_decrypt, aes_cbc_encrypt, aes_decrypt_text\n'), ((9, 51, 9, 76), 'os.path.abspath', 'os.path.abspath', ({(9, 67, 9, 75): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((26, 32, 26, 53), 'haruhi_dl.utils.bytes_to_intlist', 'bytes_to_intlist', ({(26, 49, 26, 52): 'msg'}, {}), '(msg)', False, 'from haruhi_dl.utils import bytes_to_intlist, intlist_to_bytes\n'), ((27, 37, 27, 64), 'haruhi_dl.aes.aes_decrypt', 'aes_decrypt', ({(27, 49, 27, 58): 'encrypted', (27, 60, 27, 63): 'key'}, {}), '(encrypted, key)', False, 'from haruhi_dl.aes import aes_decrypt, aes_encrypt, aes_cbc_decrypt, aes_cbc_encrypt, aes_decrypt_text\n'), ((34, 37, 34, 77), 'haruhi_dl.aes.aes_cbc_decrypt', 'aes_cbc_decrypt', ({(34, 53, 34, 57): 'data', (34, 59, 34, 67): 'self.key', (34, 69, 34, 76): 'self.iv'}, {}), '(data, self.key, self.iv)', False, 'from haruhi_dl.aes import aes_decrypt, aes_encrypt, aes_cbc_decrypt, aes_cbc_encrypt, aes_decrypt_text\n'), ((39, 37, 39, 77), 'haruhi_dl.aes.aes_cbc_encrypt', 'aes_cbc_encrypt', ({(39, 53, 39, 57): 'data', (39, 59, 39, 67): 'self.key', (39, 69, 39, 76): 'self.iv'}, {}), '(data, self.key, self.iv)', False, 'from haruhi_dl.aes import aes_decrypt, aes_encrypt, aes_cbc_decrypt, aes_cbc_encrypt, aes_decrypt_text\n'), ((45, 19, 45, 45), 'haruhi_dl.utils.intlist_to_bytes', 'intlist_to_bytes', ({(45, 36, 45, 44): 'self.key'}, {}), '(self.key)', False, 'from haruhi_dl.utils import bytes_to_intlist, intlist_to_bytes\n'), ((53, 19, 53, 45), 'haruhi_dl.utils.intlist_to_bytes', 'intlist_to_bytes', ({(53, 36, 53, 44): 'self.key'}, {}), '(self.key)', False, 'from haruhi_dl.utils import bytes_to_intlist, intlist_to_bytes\n'), ((47, 12, 47, 41), 'haruhi_dl.utils.intlist_to_bytes', 'intlist_to_bytes', ({(47, 29, 47, 40): 'self.iv[:8]'}, {}), '(self.iv[:8])', False, 'from haruhi_dl.utils import bytes_to_intlist, intlist_to_bytes\n'), ((55, 12, 55, 41), 'haruhi_dl.utils.intlist_to_bytes', 'intlist_to_bytes', ({(55, 29, 55, 40): 'self.iv[:8]'}, {}), '(self.iv[:8])', False, 'from haruhi_dl.utils import bytes_to_intlist, intlist_to_bytes\n')]