create submission files and seen data

securedrop/create-dev-data.py

@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 
 import datetime
+import io
 import os
 import argparse
 import math
@@ -15,11 +16,11 @@
 
 from sdconfig import config
 from db import db
-from models import Journalist, Reply, SeenReply, Source, Submission
+from models import Journalist, Reply, SeenFile, SeenMessage, SeenReply, Source, Submission
 from specialstrings import strings
 
 
-submissions = cycle(strings)
+messages = cycle(strings)
 replies = cycle(strings)
 
 
@@ -30,18 +31,22 @@ def main(staging=False):
         test_password = "correct horse battery staple profanity oil chewy"
         test_otp_secret = "JHCOGO7VCER3EJ4L"
 
-        add_test_user("journalist",
-                      test_password,
-                      test_otp_secret,
-                      is_admin=True)
+        journalist_who_saw = add_test_user(
+            "journalist",
+            test_password,
+            test_otp_secret,
+            is_admin=True
+        )
 
         if staging:
             return
 
-        add_test_user("dellsberg",
-                      test_password,
-                      test_otp_secret,
-                      is_admin=False)
+        dellsberg = add_test_user(
+            "dellsberg",
+            test_password,
+            test_otp_secret,
+            is_admin=False
+        )
 
         journalist_tobe_deleted = add_test_user("clarkkent",
                                                 test_password,
@@ -50,21 +55,24 @@ def main(staging=False):
                                                 first_name="Clark",
                                                 last_name="Kent")
 
-        NUM_SOURCES = os.getenv('NUM_SOURCES', 2)
+        NUM_SOURCES = os.getenv('NUM_SOURCES', 3)
         if NUM_SOURCES == "ALL":
             # We ingest two strings per source, so this will create the required
             # number of sources to include all special strings
             NUM_SOURCES = math.ceil(len(strings) / 2)
-        # Add test sources and submissions
+
+        # Create source data
         num_sources = int(NUM_SOURCES)
-        for i in range(1, num_sources + 1):
-            if i == 1:
-                # For the first source, the journalist who replied will be deleted
-                create_source_and_submissions(
-                    i, num_sources, journalist_who_replied=journalist_tobe_deleted
-                )
-                continue
-            create_source_and_submissions(i, num_sources)
+        for i in range(num_sources):
+            # For the first source, the journalist who replied will be deleted, otherwise dellsberg
+            journalist_who_replied = journalist_tobe_deleted if i == 0 else dellsberg
+
+            create_source_data(
+                i,
+                num_sources,
+                journalist_who_replied,
+                journalist_who_saw
+            )
 
         # Now let us delete one journalist
         db.session.delete(journalist_tobe_deleted)
@@ -91,8 +99,14 @@ def add_test_user(username, password, otp_secret, is_admin=False,
         db.session.rollback()
 
 
-def create_source_and_submissions(
-    source_index, source_count, num_submissions=2, num_replies=2, journalist_who_replied=None  # noqa: W605, E501
+def create_source_data(
+    source_index,
+    source_count,
+    journalist_who_replied,
+    journalist_who_saw,
+    num_files=2,
+    num_messages=2,
+    num_replies=2,
 ):
     # Store source in database
     codename = current_app.crypto_util.genrandomid()
@@ -107,10 +121,16 @@ def create_source_and_submissions(
     os.mkdir(current_app.storage.path(source.filesystem_id))
     current_app.crypto_util.genkeypair(source.filesystem_id, codename)
 
-    # Generate some test submissions
-    for _ in range(num_submissions):
+    # Mark a third of sources as seen, a third as partially-seen, and a third as unseen
+    seen_files = 0 if source_index % 3 == 0 else math.floor(num_files / (source_index % 3))
+    seen_messages = 0 if source_index % 3 == 0 else math.floor(num_messages / (source_index % 3))
+    seen_replies = 0 if source_index % 3 == 0 else math.floor(num_replies / (source_index % 3))
+
+    # Generate some test messages
+    seen_messages_count = 0
+    for _ in range(num_messages):
         source.interaction_count += 1
-        submission_text = next(submissions)
+        submission_text = next(messages)
         fpath = current_app.storage.save_message_submission(
             source.filesystem_id,
             source.interaction_count,
@@ -120,8 +140,40 @@ def create_source_and_submissions(
         source.last_updated = datetime.datetime.utcnow()
         submission = Submission(source, fpath)
         db.session.add(submission)
+        if seen_messages_count < seen_messages:
+            seen_messages_count = seen_messages_count + 1
+            db.session.flush()
+            seen_message = SeenMessage(
+                message_id=submission.id,
+                journalist_id=journalist_who_saw.id
+            )
+            db.session.add(seen_message)
+
+    # Generate some test files
+    seen_files_count = 0
+    for _ in range(num_files):
+        source.interaction_count += 1
+        fpath = current_app.storage.save_file_submission(
+            source.filesystem_id,
+            source.interaction_count,
+            source.journalist_filename,
+            "memo.txt",
+            io.BytesIO(b"This is an example of a plain text file upload.")
+        )
+        source.last_updated = datetime.datetime.utcnow()
+        submission = Submission(source, fpath)
+        db.session.add(submission)
+        if seen_files_count < seen_files:
+            seen_files_count = seen_files_count + 1
+            db.session.flush()
+            seen_file = SeenFile(
+                file_id=submission.id,
+                journalist_id=journalist_who_saw.id
+            )
+            db.session.add(seen_file)
 
     # Generate some test replies
+    seen_replies_count = 0
     for _ in range(num_replies):
         source.interaction_count += 1
         fname = "{}-{}-reply.gpg".format(source.interaction_count,
@@ -132,23 +184,35 @@ def create_source_and_submissions(
              config.JOURNALIST_KEY],
             current_app.storage.path(source.filesystem_id, fname))
 
-        if not journalist_who_replied:
-            journalist = Journalist.query.first()
-        else:
-            journalist = journalist_who_replied
-        reply = Reply(journalist, source, fname)
+        reply = Reply(journalist_who_replied, source, fname)
         db.session.add(reply)
         db.session.flush()
-        seen_reply = SeenReply(reply_id=reply.id, journalist_id=journalist.id)
+        # Journalist who replied has seen the reply
+        seen_reply = SeenReply(
+            reply_id=reply.id,
+            journalist_id=journalist_who_replied.id
+        )
         db.session.add(seen_reply)
+        if seen_replies_count < seen_replies:
+            seen_replies_count = seen_replies_count + 1
+            seen_reply = SeenReply(
+                reply_id=reply.id,
+                journalist_id=journalist_who_saw.id
+            )
+            db.session.add(seen_reply)
 
     db.session.commit()
 
     print(
         "Test source {}/{} (codename: '{}', journalist designation '{}') "
-        "added with {} submissions and {} replies".format(
-            source_index, source_count, codename, journalist_designation,
-            num_submissions, num_replies
+        "added with {} files, {} messages, and {} replies".format(
+            source_index + 1,
+            source_count,
+            codename,
+            journalist_designation,
+            num_files,
+            num_messages,
+            num_replies
         )
     )