sqlcrypt.py

import os
import apsw


from Crypto.Cipher import AES
from Crypto.Protocol.KDF import scrypt
from Crypto.Hash import HMAC, SHA256

# PASSWORD SALT (16)
HEADER_SIZE = 16
# IV (16) | MAC (32) | CIPHERTEXT
BLOCK_HEADER_SIZE = 16 + 32
BLOCK_SIZE = 4096 + BLOCK_HEADER_SIZE


def _parse_header(file):
    file.seek(0)
    password_salt = file.read(HEADER_SIZE)
    return password_salt


def _decrypt(key: bytes, data: bytes) -> bytearray:
    assert len(key) == 32, "Must decrypt in AES-256"
    assert len(data) == BLOCK_SIZE
    iv = data[:16]
    mac = data[16:48]
    ciphertext = data[BLOCK_HEADER_SIZE:]

    plaintext = AES.new(key, AES.MODE_CBC, iv=iv).decrypt(ciphertext)
    if mac != HMAC.new(key, msg=plaintext, digestmod=SHA256).digest():
        return b""

    assert len(plaintext) == BLOCK_SIZE - BLOCK_HEADER_SIZE
    return bytearray(plaintext)


def _encrypt(key: bytes, plaintext: bytes) -> bytearray:
    assert len(key) == 32, "Must encrypt in AES-256"
    assert len(plaintext) == BLOCK_SIZE - BLOCK_HEADER_SIZE

    iv = os.urandom(16)
    mac = HMAC.new(key, msg=plaintext, digestmod=SHA256).digest()

    ciphertext = AES.new(key, AES.MODE_CBC, iv=iv).encrypt(plaintext)

    assert len(ciphertext) == BLOCK_SIZE - BLOCK_HEADER_SIZE
    assert len(iv + mac) == BLOCK_HEADER_SIZE

    return bytearray(iv + mac + ciphertext)


def _derive_password(password: str, salt: bytes) -> bytes:
    """Generate the AES key from the master password."""
    return scrypt(password, salt=salt, key_len=32, N=2 ** 15, r=8, p=1)


def decrypt_database(password: str, filename: str) -> bytes:
    """For testing purpose only.

    Decrypt the encrypt database. To use for testing purpose, to check that the decrypted
    database is the same as the not-encrypted one.
    """
    with open(filename, "rb") as file:
        password_salt = _parse_header(file)
        key = _derive_password(password, password_salt)

        file.seek(HEADER_SIZE)
        data = file.read()

        plaindata = b""
        for i in range(0, len(data), BLOCK_SIZE):
            plaindata += _decrypt(key, data[i : i + BLOCK_SIZE])

    return plaindata


class EncryptedVFSFile(apsw.VFSFile):
    def __init__(self, key: bytes, *args):
        """Encrypt the data when writing on the disk and decrypt when reading.

        :param key: AES key used for encryption/decryption
        """
        assert len(key) == 32

        self.key = key
        super().__init__(*args)

        # check that the SQLite page size match with the encryption block size
        sector_size = self.xSectorSize()
        assert sector_size <= BLOCK_SIZE - BLOCK_HEADER_SIZE
        assert not (BLOCK_SIZE - BLOCK_HEADER_SIZE) % sector_size

    def xRead(self, amount: int, offset: int) -> bytes:
        assert amount <= BLOCK_SIZE

        offset += (offset // (BLOCK_SIZE - BLOCK_HEADER_SIZE)) * BLOCK_HEADER_SIZE

        start_block_offset = offset - (offset % BLOCK_SIZE)
        end_offset = offset + amount
        end_block_offset = end_offset - (end_offset % -BLOCK_SIZE)

        assert not (end_block_offset - start_block_offset) % BLOCK_SIZE

        data = super().xRead(
            end_block_offset - start_block_offset, start_block_offset + HEADER_SIZE
        )

        if not data:
            return data

        data = _decrypt(self.key, data)
        return data[offset % BLOCK_SIZE : offset % BLOCK_SIZE + amount]

    def xWrite(self, data: bytes, offset: int):
        offset += (offset // (BLOCK_SIZE - BLOCK_HEADER_SIZE)) * BLOCK_HEADER_SIZE

        start_block_offset = offset - (offset % BLOCK_SIZE)

        assert not start_block_offset % BLOCK_SIZE

        if (offset + len(data)) > (start_block_offset + BLOCK_SIZE - BLOCK_HEADER_SIZE):
            end_data_pos = (
                start_block_offset + BLOCK_SIZE - BLOCK_HEADER_SIZE
            ) - offset
            self.xWrite(data[:end_data_pos], start_block_offset + BLOCK_SIZE)
            data = data[:end_data_pos]

        assert offset + len(data) <= start_block_offset + BLOCK_SIZE

        # the data doesn't fill into blocks
        # so we need to read the entire block to add data
        blocks_data = super().xRead(BLOCK_SIZE, start_block_offset + HEADER_SIZE)

        assert len(blocks_data) in [0, BLOCK_SIZE]

        if not len(blocks_data):
            # we add new block at the end of the file
            blocks_data = bytearray(os.urandom(BLOCK_SIZE - BLOCK_HEADER_SIZE))
        else:
            blocks_data = _decrypt(self.key, blocks_data)

        blocks_data[offset % BLOCK_SIZE : offset % BLOCK_SIZE + len(data)] = data
        blocks_data = _encrypt(self.key, blocks_data)

        super().xWrite(blocks_data, start_block_offset + HEADER_SIZE)

    def xFileSize(self) -> int:
        file_size = super().xFileSize()
        return file_size - HEADER_SIZE - (file_size // BLOCK_SIZE) - BLOCK_HEADER_SIZE


class EncryptedVFS(apsw.VFS):
    def __init__(self, password: str, vfsname: str = "encrypted", basevfs: str = ""):
        self.vfsname = vfsname
        self.basevfs = basevfs
        self.password = password

        # SQLite open and close many time the same file
        # so we store the derived keys and the IVs salt
        # to not re-generate them each time the file is opened
        self.files_key = {}

        super().__init__(self.vfsname, self.basevfs)

    def xOpen(self, name, flags):
        """Open the file and read the header (store necessary data)."""
        filename = name if isinstance(name, str) else name.filename()

        if filename not in self.files_key:
            with open(filename, "a+b") as file:
                password_salt = _parse_header(file)

                if not password_salt:
                    # init the header of the file
                    password_salt = os.urandom(16)

                    file.seek(0)
                    file.write(password_salt)

            self.files_key[filename] = _derive_password(self.password, password_salt)

        return EncryptedVFSFile(self.files_key[filename], self.basevfs, name, flags)


class Connection(apsw.Connection):
    def __init__(self, filename: str, password: str) -> "Connection":
        self.password = password
        self.encrypted_vfs = EncryptedVFS(password)
        super().__init__(filename, vfs=self.encrypted_vfs.vfsname)

    def change_password(self, new_password: str):
        new_password_salt = os.urandom(16)
        new_key = _derive_password(new_password, new_password_salt)

        if not self.filename:
            # ':memory:' database (so not encrypted)
            return

        with open(self.filename, "r+b") as file:
            file.seek(0)

            header = file.read(HEADER_SIZE)

            old_password_salt = header[:16]
            old_key = _derive_password(self.password, old_password_salt)

            if len(header) != HEADER_SIZE:
                self.__init__(self.filename, new_password)
                return

            file.seek(0)
            file.write(new_password_salt)

            data = file.read(BLOCK_SIZE)
            offset = 0
            while data:
                data = _decrypt(old_key, data)
                data = _encrypt(new_key, data)

                file.seek(-BLOCK_SIZE, 1)
                file.write(data)

                offset += BLOCK_SIZE
                data = file.read(BLOCK_SIZE)

        self.__init__(self.filename, new_password)