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
|
#!/usr/bin/env python
import mock
import os
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
from test_controller.utils import mkdir_p, payload
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 _encrypt(path, data):
encrypt = _get_encrypt_function('user-0', path)
return encrypt('blob', BytesIO(data))
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()
|
