diff --git a/securedrop/qa_loader.py b/securedrop/qa_loader.py index 832f33d16f..ca8b66b4f8 100755 --- a/securedrop/qa_loader.py +++ b/securedrop/qa_loader.py @@ -22,9 +22,6 @@ random.seed('~(=^–^)') # mrow? -JOURNALIST_COUNT = 10 -SOURCE_COUNT = 50 - def random_bool(): return bool(random.getrandbits(1)) @@ -56,128 +53,6 @@ def random_datetime(nullable): ) -def new_journalist(): - # Make a diceware-like password - pw = ' '.join([random_chars(3, nullable=False, chars=DICEWARE_SAFE_CHARS) - for _ in range(7)]) - journalist = Journalist(random_chars(random.randint(3, 32), - nullable=False), - pw, - random_bool()) - if random_bool(): - # to add legacy passwords back in - journalist.passphrase_hash = None - journalist.pw_salt = random_chars(32, nullable=False) - journalist.pw_hash = random_chars(64, nullable=False) - - journalist.is_admin = bool_or_none() - - journalist.is_totp = bool_or_none() - journalist.hotp_counter = (random.randint(-1000, 1000) - if random_bool() else None) - journalist.created_on = random_datetime(nullable=True) - journalist.last_access = random_datetime(nullable=True) - - db.session.add(journalist) - - -def new_source(): - fid_len = random.randint(4, 32) - designation_len = random.randint(4, 32) - source = Source(random_chars(fid_len, nullable=False, - chars=string.ascii_lowercase), - random_chars(designation_len, nullable=False)) - source.flagged = bool_or_none() - source.last_updated = random_datetime(nullable=True) - source.pending = False - - db.session.add(source) - - -def new_submission(config, source_id): - source = Source.query.get(source_id) - - # A source may have a null fid according to the DB, but this will - # break storage.path. - if source.filesystem_id is None: - return - - filename = fake_file(config, source.filesystem_id) - submission = Submission(source, filename) - - # For issue #1189 - if random_bool(): - submission.source_id = None - - submission.downloaded = bool_or_none() - - db.session.add(submission) - - -def fake_file(config, source_fid): - source_dir = path.join(config.STORE_DIR, source_fid) - if not path.exists(source_dir): - os.mkdir(source_dir) - - filename = random_chars(20, nullable=False, chars=string.ascii_lowercase) - num = random.randint(0, 100) - msg_type = 'msg' if random_bool() else 'doc.gz' - filename = '{}-{}-{}.gpg'.format(num, filename, msg_type) - f_len = int(math.floor(random.expovariate(100000) * 1024 * 1024 * 500)) - sub_path = current_app.storage.path(source_fid, filename) - with open(sub_path, 'w') as f: - f.write('x' * f_len) - - return filename - - -def new_source_star(source_id): - source = Source.query.get(source_id) - star = SourceStar(source, bool_or_none()) - db.session.add(star) - - -def new_reply(config, journalist_id, source_id): - source = Source.query.get(source_id) - - # A source may have a null fid according to the DB, but this will - # break storage.path. - if source.filesystem_id is None: - return - - journalist = Journalist.query.get(journalist_id) - filename = fake_file(config, source.filesystem_id) - reply = Reply(journalist, source, filename) - db.session.add(reply) - - -def new_journalist_login_attempt(journalist_id): - journalist = Journalist.query.get(journalist_id) - attempt = JournalistLoginAttempt(journalist) - attempt.timestamp = random_datetime(nullable=True) - db.session.add(attempt) - - -def new_abandoned_submission(config, source_id): - '''For issue #1189''' - - source = Source.query.filter(Source.filesystem_id.isnot(None)).all()[0] - filename = fake_file(config, source.filesystem_id) - - # Use this as hack to create a real submission then swap out the source_id - submission = Submission(source, filename) - submission.source_id = source_id - db.session.add(submission) - db.session.commit() - delete_source(source_id) - - -def delete_source(source_id): - '''For issue #1189''' - db.session.execute(text('DELETE FROM sources WHERE id = :source_id'), - {'source_id': source_id}) - - def positive_int(s): i = int(s) if i < 1: @@ -185,40 +60,166 @@ def positive_int(s): return i -def load_data(config, multiplier): - app = create_app(config) - - with app.app_context(): - for _ in range(JOURNALIST_COUNT * multiplier): - new_journalist() +class QaLoader(object): + + JOURNALIST_COUNT = 10 + SOURCE_COUNT = 50 + + def __init__(self, config, multiplier): + self.config = config + self.app = create_app(config) + self.multiplier = multiplier + + def new_journalist(self): + # Make a diceware-like password + pw = ' '.join( + [random_chars(3, nullable=False, chars=DICEWARE_SAFE_CHARS) + for _ in range(7)]) + journalist = Journalist(random_chars(random.randint(3, 32), + nullable=False), + pw, + random_bool()) + if random_bool(): + # to add legacy passwords back in + journalist.passphrase_hash = None + journalist.pw_salt = random_chars(32, nullable=False) + journalist.pw_hash = random_chars(64, nullable=False) + + journalist.is_admin = bool_or_none() + + journalist.is_totp = bool_or_none() + journalist.hotp_counter = (random.randint(-1000, 1000) + if random_bool() else None) + journalist.created_on = random_datetime(nullable=True) + journalist.last_access = random_datetime(nullable=True) + + db.session.add(journalist) + + def new_source(self): + fid_len = random.randint(4, 32) + designation_len = random.randint(4, 32) + source = Source(random_chars(fid_len, nullable=False, + chars=string.ascii_lowercase), + random_chars(designation_len, nullable=False)) + source.flagged = bool_or_none() + source.last_updated = random_datetime(nullable=True) + source.pending = False + + db.session.add(source) + + def new_submission(self, source_id): + source = Source.query.get(source_id) + + # A source may have a null fid according to the DB, but this will + # break storage.path. + if source.filesystem_id is None: + return + + filename = self.fake_file(source.filesystem_id) + submission = Submission(source, filename) + + # For issue #1189 + if random_bool(): + submission.source_id = None + + submission.downloaded = bool_or_none() + + db.session.add(submission) + + def fake_file(self, source_fid): + source_dir = path.join(self.config.STORE_DIR, source_fid) + if not path.exists(source_dir): + os.mkdir(source_dir) + + filename = random_chars(20, + nullable=False, + chars=string.ascii_lowercase) + num = random.randint(0, 100) + msg_type = 'msg' if random_bool() else 'doc.gz' + filename = '{}-{}-{}.gpg'.format(num, filename, msg_type) + f_len = int(math.floor(random.expovariate(100000) * 1024 * 1024 * 500)) + sub_path = current_app.storage.path(source_fid, filename) + with open(sub_path, 'w') as f: + f.write('x' * f_len) + + return filename + + def new_source_star(self, source_id): + source = Source.query.get(source_id) + star = SourceStar(source, bool_or_none()) + db.session.add(star) + + def new_reply(self, journalist_id, source_id): + source = Source.query.get(source_id) + + # A source may have a null fid according to the DB, but this will + # break storage.path. + if source.filesystem_id is None: + return + + journalist = Journalist.query.get(journalist_id) + filename = self.fake_file(source.filesystem_id) + reply = Reply(journalist, source, filename) + db.session.add(reply) + + def new_journalist_login_attempt(self, journalist_id): + journalist = Journalist.query.get(journalist_id) + attempt = JournalistLoginAttempt(journalist) + attempt.timestamp = random_datetime(nullable=True) + db.session.add(attempt) + + def new_abandoned_submission(self, source_id): + '''For issue #1189''' + + source = Source.query.filter(Source.filesystem_id.isnot(None)).all()[0] + filename = self.fake_file(source.filesystem_id) + + # Use this as hack to create a real submission then swap out the + # source_id + submission = Submission(source, filename) + submission.source_id = source_id + db.session.add(submission) db.session.commit() + self.delete_source(source_id) - for _ in range(SOURCE_COUNT * multiplier): - new_source() - db.session.commit() + def delete_source(self, source_id): + '''For issue #1189''' + db.session.execute(text('DELETE FROM sources WHERE id = :source_id'), + {'source_id': source_id}) - for sid in range(1, SOURCE_COUNT * multiplier, 5): - for _ in range(1, multiplier + 1): - new_submission(config, sid) - db.session.commit() + def load(self): + with self.app.app_context(): + for _ in range(self.JOURNALIST_COUNT * self.multiplier): + self.new_journalist() + db.session.commit() - for sid in range(1, SOURCE_COUNT * multiplier, 5): - new_source_star(sid) - db.session.commit() + for _ in range(self.SOURCE_COUNT * self.multiplier): + self.new_source() + db.session.commit() - for jid in range(1, JOURNALIST_COUNT * multiplier, 10): - for sid in range(1, SOURCE_COUNT * multiplier, 10): - for _ in range(1, 3): - new_reply(config, jid, sid) - db.session.commit() + for sid in range(1, self.SOURCE_COUNT * self.multiplier, 5): + for _ in range(1, self.multiplier + 1): + self.new_submission(sid) + db.session.commit() - for jid in range(1, JOURNALIST_COUNT * multiplier, 10): - new_journalist_login_attempt(jid) - db.session.commit() + for sid in range(1, self.SOURCE_COUNT * self.multiplier, 5): + self.new_source_star(sid) + db.session.commit() + + for jid in range(1, self.JOURNALIST_COUNT * self.multiplier, 10): + for sid in range(1, self.SOURCE_COUNT * self.multiplier, 10): + for _ in range(1, 3): + self.new_reply(jid, sid) + db.session.commit() + + for jid in range(1, self.JOURNALIST_COUNT * self.multiplier, 10): + self.new_journalist_login_attempt(jid) + db.session.commit() - for sid in range(SOURCE_COUNT * multiplier, - SOURCE_COUNT * multiplier + multiplier): - new_abandoned_submission(config, sid) + for sid in range( + self.SOURCE_COUNT * self.multiplier, + self.SOURCE_COUNT * self.multiplier + self.multiplier): + self.new_abandoned_submission(sid) def arg_parser(): @@ -234,7 +235,7 @@ def arg_parser(): def main(): args = arg_parser().parse_args() print('Loading data. This make take a while.') - load_data(sdconfig, args.multiplier) + QaLoader(sdconfig, args.multiplier).load() if __name__ == '__main__': diff --git a/securedrop/tests/test_qa_loader.py b/securedrop/tests/test_qa_loader.py index cfaa400499..cb48bd35a8 100644 --- a/securedrop/tests/test_qa_loader.py +++ b/securedrop/tests/test_qa_loader.py @@ -1,8 +1,8 @@ # -*- coding: utf-8 -*- -from qa_loader import load_data +from qa_loader import QaLoader def test_load_data(journalist_app, config): # Use the journalist_app fixture to init the DB - load_data(config, multiplier=1) + QaLoader(config, multiplier=1).load()