peewee/tests/sqlcipher_ext.py

import datetime
import os
from hashlib import sha1

from peewee import DatabaseError
from playhouse.sqlcipher_ext import *
from playhouse.sqlite_ext import *

from .base import ModelTestCase
from .base import TestModel


PASSPHRASE = 'testing sqlcipher'
PRAGMAS = {
    'kdf_iter': 10,   # Much faster for testing. Totally unsafe.
    'cipher_log_level': 'none',
}
db = SqlCipherDatabase('peewee_test.dbc', passphrase=PASSPHRASE,
                       pragmas=PRAGMAS, rank_functions=True)


@db.func('shazam')
def shazam(s):
    return sha1((s or '').encode('utf-8')).hexdigest()[:5]


class Thing(TestModel):
    name = CharField()


class FTSNote(FTSModel, TestModel):
    content = TextField()


class Note(TestModel):
    content = TextField()
    timestamp = DateTimeField(default=datetime.datetime.now)


class CleanUpModelTestCase(ModelTestCase):
    def tearDown(self):
        super(CleanUpModelTestCase, self).tearDown()
        if os.path.exists(self.database.database):
            os.unlink(self.database.database)


class SqlCipherTestCase(CleanUpModelTestCase):
    database = db
    requires = [Thing]

    def test_good_and_bad_passphrases(self):
        things = ('t1', 't2', 't3')
        for thing in things:
            Thing.create(name=thing)

        # Try to open db with wrong passphrase
        bad_db = SqlCipherDatabase(db.database, passphrase='wrong passphrase')
        self.assertRaises(DatabaseError, bad_db.get_tables)

        # Assert that we can still access the data with the good passphrase.
        query = Thing.select().order_by(Thing.name)
        self.assertEqual([t.name for t in query], ['t1', 't2', 't3'])

    def test_rekey(self):
        things = ('t1', 't2', 't3')
        for thing in things:
            Thing.create(name=thing)

        self.database.rekey('a new passphrase')

        db2 = SqlCipherDatabase(db.database, passphrase='a new passphrase',
                                pragmas=PRAGMAS)
        cursor = db2.execute_sql('select name from thing order by name;')
        self.assertEqual([name for name, in cursor], ['t1', 't2', 't3'])

        query = Thing.select().order_by(Thing.name)
        self.assertEqual([t.name for t in query], ['t1', 't2', 't3'])

        self.database.close()
        self.database.connect()

        query = Thing.select().order_by(Thing.name)
        self.assertEqual([t.name for t in query], ['t1', 't2', 't3'])

        # Re-set to the original passphrase.
        self.database.rekey(PASSPHRASE)

    def test_empty_passphrase(self):
        db = SqlCipherDatabase(':memory:')

        class CM(TestModel):
            data = TextField()
            class Meta:
                database = db

        db.connect()
        db.create_tables([CM])
        cm = CM.create(data='foo')
        cm_db = CM.get(CM.data == 'foo')
        self.assertEqual(cm_db.id, cm.id)
        self.assertEqual(cm_db.data, 'foo')


config_db = SqlCipherDatabase('peewee_test.dbc', pragmas={
    'kdf_iter': 1234,
    'cipher_page_size': 8192}, passphrase=PASSPHRASE)

class TestSqlCipherConfiguration(CleanUpModelTestCase):
    database = config_db

    def test_configuration_via_pragma(self):
        # Write some data so the database file is created.
        self.database.execute_sql('create table foo (data TEXT)')
        self.database.close()

        self.database.connect()
        self.assertEqual(int(self.database.pragma('kdf_iter')), 1234)
        self.assertEqual(int(self.database.pragma('cipher_page_size')), 8192)
        self.assertTrue('foo' in self.database.get_tables())


class SqlCipherExtTestCase(CleanUpModelTestCase):
    database = db
    requires = [Note]

    def setUp(self):
        super(SqlCipherExtTestCase, self).setUp()
        FTSNote._meta.database = db
        FTSNote.drop_table(True)
        FTSNote.create_table(tokenize='porter', content=Note.content)

    def tearDown(self):
        FTSNote.drop_table(True)
        super(SqlCipherExtTestCase, self).tearDown()

    def test_fts(self):
        strings = [
            'python and peewee for working with databases',
            'relational databases are the best',
            'sqlite is the best relational database',
            'sqlcipher is a cool database extension']
        for s in strings:
            Note.create(content=s)
        FTSNote.rebuild()

        query = (FTSNote
                 .select(FTSNote, FTSNote.rank().alias('score'))
                 .where(FTSNote.match('relational databases'))
                 .order_by(SQL('score').desc()))
        notes = [note.content for note in query]
        self.assertEqual(notes, [
            'relational databases are the best',
            'sqlite is the best relational database'])

        alt_conn = SqliteDatabase(db.database)
        self.assertRaises(
            DatabaseError,
            alt_conn.execute_sql,
            'SELECT * FROM "%s"' % (FTSNote._meta.table_name))

    def test_func(self):
        Note.create(content='hello')
        Note.create(content='baz')
        Note.create(content='nug')

        query = (Note
                 .select(Note.content, fn.shazam(Note.content).alias('shz'))
                 .order_by(Note.id)
                 .dicts())
        results = list(query)
        self.assertEqual(results, [
            {'content': 'hello', 'shz': 'aaf4c'},
            {'content': 'baz', 'shz': 'bbe96'},
            {'content': 'nug', 'shz': '52616'},
        ])