import io
import tarfile
import unittest

from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto.Hash import HMAC
from Crypto.Protocol.KDF import PBKDF2

import kobackupdec


class PmsSplitTests(unittest.TestCase):

    PASSWORD = b'test-password-1234-AA'
    SALT = bytes(range(32))
    NONCE = bytes(range(16))

    @staticmethod
    def _prf(password, salt):
        return HMAC.new(password, salt, SHA256).digest()

    def _encrypt_part(self, cleartext):
        key = PBKDF2(self.PASSWORD, self.SALT, 32, 10000,
                     self._prf, hmac_hash_module=None)
        cipher = AES.new(key, AES.MODE_GCM, nonce=self.NONCE, mac_len=16)
        ciphertext, tag = cipher.encrypt_and_digest(cleartext)
        return ciphertext + tag

    def _make_tar(self):
        output = io.BytesIO()
        with tarfile.open(fileobj=output, mode='w') as tar:
            payloads = {
                'cache/example.txt': b'cache data',
                'files/games/com.mojang/minecraftWorlds/test/level.dat':
                    b'level data',
            }
            for name, data in payloads.items():
                info = tarfile.TarInfo(name)
                info.size = len(data)
                tar.addfile(info, io.BytesIO(data))
        return output.getvalue()

    def _material(self):
        material = kobackupdec.DecryptMaterial('BackupFileModuleInfo')
        material.name = 'com.example.app'
        material.encMsgV3 = (self.SALT + self.NONCE).hex()
        material.is_pms_split = True
        return material

    def _decryptor(self, password=None):
        decryptor = kobackupdec.Decryptor(password or self.PASSWORD)
        decryptor.type_attch = 3
        decryptor.crypto_init()
        return decryptor

    def test_decrypt_pms_split_parts_reconstructs_tar(self):
        clear_tar = self._make_tar()
        split = len(clear_tar) // 3
        clear_parts = [
            clear_tar[:split],
            clear_tar[split:2 * split],
            clear_tar[2 * split:],
        ]
        encrypted_parts = [self._encrypt_part(part) for part in clear_parts]

        decryptor = self._decryptor()
        material = self._material()
        reconstructed = b''.join(
            decryptor.decrypt_pms_split_part(material, part)
            for part in encrypted_parts)

        self.assertEqual(clear_tar, reconstructed)
        with tarfile.open(fileobj=io.BytesIO(reconstructed), mode='r:*') as tar:
            self.assertIn('cache/example.txt', tar.getnames())

    def test_wrong_password_fails_gcm_authentication(self):
        encrypted = self._encrypt_part(b'test payload')
        decryptor = self._decryptor(b'wrong-password')
        self.assertIsNone(
            decryptor.decrypt_pms_split_part(self._material(), encrypted))

    def test_boolean_xml_value(self):
        document = kobackupdec.xml.dom.minidom.parseString(
            '<column><value Boolean="true" /></column>')
        self.assertTrue(
            kobackupdec.xml_get_column_value(document.documentElement))


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