1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
|
#!/usr/bin/env python
import base64
import errno
import mock
import os
import random
from argparse import ArgumentParser
from io import BytesIO
from tempfile import mkdtemp
from twisted.internet import reactor, defer
from leap.soledad.client._db.blobs import BlobManager
DEFAULT_TARGET_DIR = './blob-templates'
def _get_put_function(path):
def _put(_, data, params=None):
with open(path, 'w') as f:
f.write(data.read())
return defer.succeed(mock.Mock(code=200))
return _put
def _get_encrypt_function(user, path):
tempdir = mkdtemp()
manager = BlobManager(tempdir, None, '123', '123', user)
manager._client.put = _get_put_function(path)
return manager._encrypt_and_upload
def payload(size):
random.seed(1337) # same seed to avoid different bench results
payload_bytes = bytearray(random.getrandbits(8) for _ in xrange(size))
# encode as base64 to avoid ascii encode/decode errors
return base64.b64encode(payload_bytes)[:size] # remove b64 overhead
def _encrypt(path, data):
encrypt = _get_encrypt_function('user-0', path)
return encrypt('blob', BytesIO(data))
def mkdir_p(path):
try:
os.makedirs(path)
except OSError as exc: # Python >2.5
if exc.errno == errno.EEXIST and os.path.isdir(path):
pass
else:
raise
def create_blob_templates(target_dir, amount, size):
mkdir_p(target_dir)
data = payload(size)
semaphore = defer.DeferredSemaphore(20)
deferreds = []
for i in xrange(amount):
path = os.path.join(target_dir,
'%dK-%d.blob' % (size / 1000, i))
d = semaphore.run(_encrypt, path, data)
deferreds.append(d)
return defer.gatherResults(deferreds)
def parse_args():
parser = ArgumentParser()
parser.add_argument(
'target_dir',
help='The target directory where templates will be written to.')
parser.add_argument(
'--amount', default=1000, type=int,
help='The number of blob templates to create.')
parser.add_argument(
'--size', default=1000, type=int,
help='The size of each template in KB.')
return parser.parse_args()
if __name__ == '__main__':
args = parse_args()
d = create_blob_templates(args.target_dir, args.amount, args.size * 1000)
d.addCallback(lambda _: reactor.stop())
reactor.run()
|
