"""The pak key search on a synthetic main: a key whose halves are two separate 16-byte literals, and one stored whole; no game files or keys."""
from pathlib import Path
import os
import struct
import sys
import tempfile
import unittest
from unittest.mock import patch

REPO = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO / 'dm/gfx/tools'))
import uepak


def adrp(rd, pc, target):
    imm = ((target & ~0xFFF) - (pc & ~0xFFF)) >> 12 & 0x1FFFFF
    return 0x90000000 | (imm & 3) << 29 | (imm >> 2) << 5 | rd


def ldr_q(rt, rn, off):
    return 0x3DC00000 | (off // 16) << 10 | rn << 5 | rt


STP_Q0_Q1_X0, RET = 0xAD000400, 0xD65F03C0


def pak(path, key, mount=b'../../../\0'):
    """A version 11 pak footer whose index starts with the mount point, encrypted with key (the 0xcc footer of Unreal Engine 4.27; key None: in the clear)"""
    index = struct.pack('<i', len(mount)) + mount
    index += bytes(-len(index) % 16)
    enc = aes_encrypt(key, index) if key else index
    foot = bytes(16) + (b'\x01' if key else b'\x00') + struct.pack('<IIQQ', uepak.PAK_MAGIC, 11, 0, len(enc)) + bytes(20) + bytes(32 * 5)
    path.write_bytes(enc + foot)


def aes_encrypt(key, data):
    import ctypes
    cc = ctypes.CDLL('/usr/lib/libSystem.B.dylib').CCCrypt
    out = ctypes.create_string_buffer(len(data)); moved = ctypes.c_size_t()
    assert cc(0, 0, 2, key, len(key), None, data, len(data), out, len(data), ctypes.byref(moved)) == 0
    return out.raw[:moved.value]


@unittest.skipUnless(sys.platform == 'darwin', 'CommonCrypto')
class KeySearch(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory(); self.addCleanup(self.tmp.cleanup)
        self.key = bytes(range(7, 39))

    def image(self, split):
        """text (0x1000 bytes) + read-only data: the callback's code in text, the key's halves 0x2350 bytes apart in it, or the key stored whole"""
        lo = 0x1000
        ro = bytearray(os.urandom(0x4000))
        a, b = lo + 0x100, lo + 0x2450
        if split: ro[a - lo:a - lo + 16], ro[b - lo:b - lo + 16] = self.key[:16], self.key[16:]
        else: ro[a - lo:a - lo + 32] = self.key
        code = [adrp(8, 0x40, a), adrp(9, 0x44, b), ldr_q(0, 8, a & 0xFFF), ldr_q(1, 9, b & 0xFFF), STP_Q0_Q1_X0, RET] if split else [RET]
        text = bytearray(lo); text[0x40:0x40 + 4 * len(code)] = struct.pack(f'<{len(code)}I', *code)
        return bytes(text + ro), lo, lo + len(ro)

    def test_split_literals(self):
        img, lo, hi = self.image(True)
        self.assertEqual(uepak.code_keys(img, lo), [self.key])
        p = Path(self.tmp.name) / 'a.pak'; pak(p, self.key)
        self.assertEqual(uepak.find_key(img, lo, hi, str(p)), self.key)

    def test_whole_window(self):
        img, lo, hi = self.image(False)
        self.assertEqual(uepak.code_keys(img, lo), [])
        p = Path(self.tmp.name) / 'a.pak'; pak(p, self.key)
        self.assertEqual(uepak.find_key(img, lo, hi, str(p)), self.key)

    def test_root_mount_point(self):
        """The empty pak beside an IoStore container mounts at "/" (DRAGON BALL Sparking! ZERO's pakchunk1-17): the key that opens the game's first pak opens it too;
        a mount point that is no path does not count as opened"""
        img, lo, hi = self.image(False)
        d = Path(self.tmp.name); a, b, c = d / 'a.pak', d / 'b.pak', d / 'c.pak'
        pak(a, self.key); pak(b, self.key, mount=b'/\0'); pak(c, self.key, mount=b'x\x01\x02\0')
        self.assertEqual(uepak.find_key(img, lo, hi, str(b)), self.key)
        self.assertEqual(uepak.find_key(img, lo, hi, str(b), first=[self.key]), self.key)
        self.assertIsNone(uepak.find_key(img, lo, hi, str(c)))

    def test_clobbered_base(self):
        """A register written between the adrp and the load holds no page any more"""
        img, lo, hi = self.image(True)
        w = list(struct.unpack_from('<6I', img, 0x40)); w.insert(2, 0xD2800008)   # mov x8, #0
        img = img[:0x40] + struct.pack('<7I', *w) + img[0x40 + 28:]
        self.assertEqual(uepak.code_keys(img, lo), [])

    def test_the_keys_of_a_game(self):
        """uepak.Keys: no main is loaded for a pak in the clear; the base game's main opens the first pak, the update's is never loaded, and the key found opens the next pak
        before any search of main"""
        img, lo, hi = self.image(False)
        d = Path(self.tmp.name); plain, a, b = d / 'plain.pak', d / 'a.pak', d / 'b.pak'
        pak(plain, None); pak(a, self.key); pak(b, self.key)
        keys = uepak.Keys(['base', 'update'])
        with patch.object(uepak, 'nso_image', return_value=(img, lo, hi)) as load:
            self.assertIsNone(keys(str(plain))); load.assert_not_called()
            self.assertEqual(keys(str(a)), self.key)
            with patch.object(uepak, 'code_keys', side_effect=AssertionError('main searched again')):
                self.assertEqual(keys(str(b)), self.key)
        self.assertEqual([c.args for c in load.call_args_list], [('base',)])

    def test_no_main_holds_the_key(self):
        p = Path(self.tmp.name) / 'a.pak'; pak(p, self.key)
        with patch.object(uepak, 'nso_image', return_value=(bytes(0x3000), 0x1000, 0x3000)):
            with self.assertRaisesRegex(ValueError, f'{p}: the index is encrypted and no key in main'): uepak.Keys(['base', 'update'])(str(p))


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