import cryptography.hazmat.primitives.ciphers as ciphers
import cryptography.hazmat.backends as backends
import pathlib
import struct
import uuid
import zlib
import os

def pad(b, i):
    n = (i - (len(b) % i))
    return (b + bytes(([n] * n)))

def unpad(s):
    return s[:-s[-1]]

def read(path):
    with pathlib.Path(path).open('rb') as f:
        return f.read()

class Buffer():
    def __init__(self, s):
        self.s = s
        self.i = 0

    def take(self, n):
        s = self.peek(n)
        self.i += n
        return s

    def peek(self, n):
        return self.s[self.i:self.i + n]

    def unpack(self, fmt, n):
        return struct.unpack(fmt, self.take(n))

    def unpack_I(self):
        return self.unpack('I', 4)

class ExtractedKey():
    def __init__(self, cipher, mapping):
        self.cipher = cipher
        self.mapping = mapping

class ExtractedFile():
    def __init__(self, meta, encrypted):
        self.meta = meta
        self.encrypted = encrypted
        self.block_count = len(self.encrypted) // 128

    def decrypt(self, key):
        decryptor = key.cipher.decryptor()

        blocks = [self.encrypted[(n * 128):((n + 1) * 128)]
                  for n in range(self.block_count)]

        remapped_blocks = {}
        for i in range(self.block_count):
            remapped_blocks[key.mapping[i]] = blocks[i]

        s = b''
        for i in range(self.block_count):
            s += remapped_blocks[i]

        s = unpad(s)
        s = decryptor.update(s) + decryptor.finalize()
        s = unpad(s)
        s = zlib.decompress(s)

        return s

class Archiver():
    def __init__(self):
        print('Reading archive and keystore from disk')
        self.a = Buffer(read('a'))
        self.k = Buffer(read('k'))
        self.files = []
        self.keys = []

    def check_file_headers(self):
        if self.a.take(8) != b'L0LARCH\x00':
            raise 'Archive invalid'
        if self.k.take(8) != b'L0LKSTR\x00':
            raise 'Keystore invalid'

    def extract(self):
        self.check_file_headers()

        (file_count,) = self.a.unpack_I()
        (key_count,) = self.k.unpack_I()

        print('Found %d files' % file_count)

        if file_count != key_count:
            raise 'File and key count mismatch'

        for i in range(file_count):
            file = self.extract_next_file()
            key = self.extract_next_key()

            print("Extracting %s: bytes=%d uuid=%s" % (
                file.meta['filename'],
                file.meta['size'],
                file.meta['uuid'])
            )

            decrypted = file.decrypt(key)
            filename = file.meta['filename']
            path = os.path.dirname(filename)

            if not os.path.exists(path):
                os.makedirs(path)

            with open(filename, 'wb+') as f:
                f.write(decrypted)

            os.chmod(filename, file.meta['mode'])

    def extract_next_key(self):
        file_uuid = uuid.UUID(bytes=self.k.take(16))
        cipher_info = list(self.k.take(32))
        cipher_info.reverse()
        iv = bytes(list(cipher_info[0:16]))
        key = bytes(list(cipher_info[16:32]))
        default_backend = backends.default_backend()

        cipher = ciphers.Cipher(
            ciphers.algorithms.AES(key),
            ciphers.modes.CBC(iv),
            backend=default_backend
        )

        (block_count,) = self.k.unpack_I()

        mapping = {}
        for i in range(block_count):
            (left, right) = self.k.unpack('2I', 8)
            mapping[left] = right

        return ExtractedKey(cipher, mapping)

    def extract_next_file(self):
        meta = self.extract_meta()
        (encrypted_len,) = self.a.unpack_I()
        encrypted = self.a.take(encrypted_len)
        return ExtractedFile(meta, encrypted)

    def extract_meta(self):
        (filename_length,) = self.a.unpack_I()
        filename = self.a.take(filename_length)
        filename = filename[:-1].decode('utf-8')
        file_uuid = uuid.UUID(bytes=self.a.take(16))
        (size, mode) = self.a.unpack('2I', 2 * 4)
        (st_atime, st_mtime, st_ctime) = self.a.unpack('3d', 3 * 8)

        return {
            'filename': filename,
            'uuid': file_uuid,
            'size': size,
            'mode': mode,
            'st_atime': st_atime,
            'st_mtime': st_mtime,
            'st_ctime': st_ctime,
        }

Archiver().extract()