GsTkt99e

1import os.path
2from io import BytesIO
3from functools import partial
4
5from twisted.enterprise import adbapi
6from twisted.internet import defer, reactor
7
8from leap.soledad.client.sqlcipher import SQLCipherOptions
9from leap.soledad.client import pragmas
10
11
12def _init_blob_table(conn):
13    maybe_create = (
14        "CREATE TABLE IF NOT EXISTS "
15        "blobs ("
16        "blob_id PRIMARY KEY, "
17        "payload BLOB)")
18    conn.execute(maybe_create)
19
20
21def _sqlcipherInitFactory(fun):
22    def _initialize(conn):
23        fun(conn)
24        _init_blob_table(conn)
25    return _initialize
26
27
28class SQLiteBlobBackend(object):
29
30    def __init__(self, path, encrypted=False, key=None):
31
32        self.path = os.path.abspath(
33            os.path.join(path, 'soledad_blob.db'))
34
35        if encrypted:
36            if not key:
37                raise ValueError('key cannot be None')
38            backend = 'pysqlcipher.dbapi2'
39            opts = SQLCipherOptions('/tmp/ignored', KEY)
40            pragmafun = partial(pragmas.set_init_pragmas, opts=opts)
41            openfun = _sqlcipherInitFactory(pragmafun)
42        else:
43            backend = 'sqlite3'
44            openfun = _init_blob_table
45
46
47        self.dbpool = dbpool = adbapi.ConnectionPool(
48            backend, self.path, 
49            check_same_thread=False, timeout=5,
50            cp_openfun=openfun,
51            cp_min=1, cp_max=2, cp_name='blob_pool')
52
53
54    def put(self, blob_id, blob_fd):
55        insert = 'INSERT INTO blobs VALUES (?, ?)'
56        raw = blob_fd.getvalue()
57        return self.dbpool.runQuery(insert, (blob_id, raw))
58        
59
60    def get(self, blob_id):
61        select = 'SELECT payload FROM blobs WHERE blob_id = ?'
62        return self.dbpool.runQuery(select, blob_id)
63
64
65
66if __name__ == '__main__':
67
68    import datetime
69    import uuid
70
71    TIMES = 10.0
72    SIZE = 1E6
73    KEY = 'supersikret'
74
75    t1 = 0
76
77    def print_stats(_):
78        t2 = datetime.datetime.now()
79        delta = (t2 - t1).total_seconds()
80
81        took = "%s seconds (%s per %s bytes blob)" % (
82            delta, delta/TIMES, int(SIZE))
83        print(took)
84        rate = int(SIZE) / delta / TIMES / 1024
85        print('%s KB/s' % rate)
86        reactor.stop()
87
88    def test():
89        global t1
90        t1 = datetime.datetime.now()
91
92        bb = SQLiteBlobBackend('/tmp', encrypted=True, key=KEY)
93        payload = BytesIO()
94        payload.write('A' * int(SIZE))
95
96        ds = []
97        for i in range(int(TIMES)):
98            _id = uuid.uuid4().get_hex()
99            ds.append(bb.put(_id, payload))
100        d = defer.gatherResults(ds)
101        d.addCallback(print_stats)
102
103
104    reactor.callWhenRunning(test)
105    reactor.run()