From e62fb66e24d0d5d28e77d8ec5834072517c1d038 Mon Sep 17 00:00:00 2001
From: heartsucker <heartsucker@autistici.org>
Date: Thu, 28 Feb 2019 12:38:32 +0100
Subject: [PATCH 1/6] updated schema to split submission into file/message

---
 securedrop_client/db.py | 71 ++++++++++++++++++++++++++++++-----------
 1 file changed, 52 insertions(+), 19 deletions(-)

diff --git a/securedrop_client/db.py b/securedrop_client/db.py
index 40d82070e..81a3044e9 100644
--- a/securedrop_client/db.py
+++ b/securedrop_client/db.py
@@ -1,8 +1,8 @@
 import os
 
 from sqlalchemy import Boolean, Column, create_engine, DateTime, ForeignKey, Integer, String, \
-    Text, MetaData
-from sqlalchemy.ext.declarative import declarative_base
+    Text, MetaData, CheckConstraint
+from sqlalchemy.ext.declarative import declarative_base, AbstractConcreteBase
 from sqlalchemy.orm import relationship, backref
 
 
@@ -60,15 +60,24 @@ def collection(self):
         """Return the list of submissions and replies for this source, sorted
         in ascending order by the filename/interaction count."""
         collection = []
-        collection.extend(self.submissions)
+        collection.extend(self.messages)
+        collection.extend(self.files)
         collection.extend(self.replies)
         collection.sort(key=lambda x: int(x.filename.split('-')[0]))
         return collection
 
 
-class Submission(Base):
+class Submission(AbstractConcreteBase, Base):
+    pass
 
-    __tablename__ = 'submissions'
+
+class Message(Submission):
+
+    __tablename__ = 'messages'
+    __mapper_args__ = {
+        'polymorphic_identity': 'message',
+        'concrete': True,
+    }
 
     id = Column(Integer, primary_key=True)
     uuid = Column(String(36), unique=True, nullable=False)
@@ -77,30 +86,54 @@ class Submission(Base):
     download_url = Column(String(255), nullable=False)
 
     # This is whether the submission has been downloaded in the local database.
-    is_downloaded = Column(Boolean(name='is_downloaded'),
-                           default=False)
+    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default='false')
 
     # This reflects read status stored on the server.
-    is_read = Column(Boolean(name='is_read'),
-                     default=False)
+    is_read = Column(Boolean(name='is_read'), nullable=False, server_default='false')
+
+    content = Column(
+        Text,
+        # this check contraint ensures the state of the DB is what one would expect
+        CheckConstraint('CASE WHEN is_downloaded = 0 THEN content IS NULL ELSE 1 END',
+                        name='compare_download_vs_content')
+    )
 
     source_id = Column(Integer, ForeignKey('sources.id'))
     source = relationship("Source",
-                          backref=backref("submissions", order_by=id,
+                          backref=backref("messages", order_by=id,
                                           cascade="delete"))
 
-    def __init__(self, source, uuid, size, filename, download_url):
-        # ORM event catching _should_ have already initialized `self.data`
+    def __repr__(self):
+        return '<Message {}>'.format(self.filename)
+
 
-        self.source_id = source.id
-        self.uuid = uuid
-        self.size = size
-        self.filename = filename
-        self.download_url = download_url
-        self.is_download = False
+class File(Submission):
+
+    __tablename__ = 'files'
+    __mapper_args__ = {
+        'polymorphic_identity': 'file',
+        'concrete': True,
+    }
+
+    id = Column(Integer, primary_key=True)
+    uuid = Column(String(36), unique=True, nullable=False)
+    filename = Column(String(255), nullable=False)
+    size = Column(Integer, nullable=False)
+    download_url = Column(String(255), nullable=False)
+
+    # This is whether the submission has been downloaded in the local database.
+    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default='false')
+
+    # This reflects read status stored on the server.
+    is_read = Column(Boolean(name='is_read'), nullable=False, server_default='false')
+
+    source_id = Column(Integer, ForeignKey('sources.id'))
+    source = relationship("Source",
+                          backref=backref("files", order_by=id,
+                                          cascade="delete"))
 
     def __repr__(self):
-        return '<Submission {}>'.format(self.filename)
+        return '<File {}>'.format(self.filename)
 
 
 class Reply(Base):

From 1cffbdf57491d09f96d180cd69c5469dffc1c7b5 Mon Sep 17 00:00:00 2001
From: heartsucker <heartsucker@autistici.org>
Date: Tue, 5 Mar 2019 17:03:00 +0100
Subject: [PATCH 2/6] correctly set default values

---
 securedrop_client/db.py | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/securedrop_client/db.py b/securedrop_client/db.py
index 81a3044e9..02ec1325a 100644
--- a/securedrop_client/db.py
+++ b/securedrop_client/db.py
@@ -32,13 +32,11 @@ class Source(Base):
     uuid = Column(String(36), unique=True, nullable=False)
     journalist_designation = Column(String(255), nullable=False)
     document_count = Column(Integer, server_default="0", nullable=False)
-    is_flagged = Column(Boolean(name='is_flagged'),
-                        server_default="false")
+    is_flagged = Column(Boolean(name='is_flagged'), server_default="0")
     public_key = Column(Text, nullable=True)
     fingerprint = Column(String(64))
     interaction_count = Column(Integer, server_default="0", nullable=False)
-    is_starred = Column(Boolean(name='is_starred'),
-                        server_default="false")
+    is_starred = Column(Boolean(name='is_starred'), server_default="0")
     last_updated = Column(DateTime)
 
     def __init__(self, uuid, journalist_designation, is_flagged, public_key,
@@ -86,10 +84,10 @@ class Message(Submission):
     download_url = Column(String(255), nullable=False)
 
     # This is whether the submission has been downloaded in the local database.
-    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default='false')
+    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default="0")
 
     # This reflects read status stored on the server.
-    is_read = Column(Boolean(name='is_read'), nullable=False, server_default='false')
+    is_read = Column(Boolean(name='is_read'), nullable=False, server_default="0")
 
     content = Column(
         Text,
@@ -122,10 +120,10 @@ class File(Submission):
     download_url = Column(String(255), nullable=False)
 
     # This is whether the submission has been downloaded in the local database.
-    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default='false')
+    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default="0")
 
     # This reflects read status stored on the server.
-    is_read = Column(Boolean(name='is_read'), nullable=False, server_default='false')
+    is_read = Column(Boolean(name='is_read'), nullable=False, server_default="0")
 
     source_id = Column(Integer, ForeignKey('sources.id'))
     source = relationship("Source",

From 0f6f1bd227f5e20a063e641bdc7d9b2a3630b90f Mon Sep 17 00:00:00 2001
From: heartsucker <heartsucker@autistici.org>
Date: Tue, 5 Mar 2019 17:03:14 +0100
Subject: [PATCH 3/6] update readme for creating migrations

---
 README.md | 1 -
 1 file changed, 1 deletion(-)

diff --git a/README.md b/README.md
index 14e53ca17..60cac8aef 100644
--- a/README.md
+++ b/README.md
@@ -100,7 +100,6 @@ rm -f svs.sqlite
 sqlite3 svs.sqlite .databases > /dev/null
 alembic upgrade head
 alembic revision --autogenerate -m "describe your revision here"
-make test-alembic
 ```
 
 ### Merging Migrations

From 8bd7dcb5e06a328870a964a2f91a5d006b76db71 Mon Sep 17 00:00:00 2001
From: heartsucker <heartsucker@autistici.org>
Date: Tue, 5 Mar 2019 17:03:29 +0100
Subject: [PATCH 4/6] update alembic migration

---
 alembic/versions/2f363b3d680e_init.py | 63 ++++++++++++++++++---------
 1 file changed, 43 insertions(+), 20 deletions(-)

diff --git a/alembic/versions/2f363b3d680e_init.py b/alembic/versions/2f363b3d680e_init.py
index 6189dbec5..a4bcdc8cd 100644
--- a/alembic/versions/2f363b3d680e_init.py
+++ b/alembic/versions/2f363b3d680e_init.py
@@ -23,17 +23,16 @@ def upgrade():
         sa.Column('uuid', sa.String(length=36), nullable=False),
         sa.Column('journalist_designation', sa.String(length=255), nullable=False),
         sa.Column('document_count', sa.Integer(), server_default='0', nullable=False),
-        sa.Column('is_flagged', sa.Boolean(name='is_flagged'), server_default='false',
-                  nullable=True),
+        sa.Column('is_flagged', sa.Boolean(name='is_flagged'), server_default='0', nullable=True),
         sa.Column('public_key', sa.Text(), nullable=True),
         sa.Column('fingerprint', sa.String(length=64), nullable=True),
         sa.Column('interaction_count', sa.Integer(), server_default='0', nullable=False),
-        sa.Column('is_starred', sa.Boolean(name='is_starred'), server_default='false',
-                  nullable=True),
+        sa.Column('is_starred', sa.Boolean(name='is_starred'), server_default='0', nullable=True),
         sa.Column('last_updated', sa.DateTime(), nullable=True),
         sa.PrimaryKeyConstraint('id', name=op.f('pk_sources')),
         sa.UniqueConstraint('uuid', name=op.f('uq_sources_uuid'))
     )
+
     op.create_table(
         'users',
         sa.Column('id', sa.Integer(), nullable=False),
@@ -42,6 +41,44 @@ def upgrade():
         sa.PrimaryKeyConstraint('id', name=op.f('pk_users')),
         sa.UniqueConstraint('uuid', name=op.f('uq_users_uuid'))
     )
+
+    op.create_table(
+        'files',
+        sa.Column('id', sa.Integer(), nullable=False),
+        sa.Column('uuid', sa.String(length=36), nullable=False),
+        sa.Column('filename', sa.String(length=255), nullable=False),
+        sa.Column('size', sa.Integer(), nullable=False),
+        sa.Column('download_url', sa.String(length=255), nullable=False),
+        sa.Column('is_downloaded', sa.Boolean(name='is_downloaded'), server_default='0',
+                  nullable=False),
+        sa.Column('is_read', sa.Boolean(name='is_read'), server_default='0', nullable=False),
+        sa.Column('source_id', sa.Integer(), nullable=True),
+        sa.ForeignKeyConstraint(['source_id'], ['sources.id'],
+                                name=op.f('fk_files_source_id_sources')),
+        sa.PrimaryKeyConstraint('id', name=op.f('pk_files')),
+        sa.UniqueConstraint('uuid', name=op.f('uq_files_uuid')),
+    )
+
+    op.create_table(
+        'messages',
+        sa.Column('id', sa.Integer(), nullable=False),
+        sa.Column('uuid', sa.String(length=36), nullable=False),
+        sa.Column('filename', sa.String(length=255), nullable=False),
+        sa.Column('size', sa.Integer(), nullable=False),
+        sa.Column('download_url', sa.String(length=255), nullable=False),
+        sa.Column('is_downloaded', sa.Boolean(name='is_downloaded'), server_default='0',
+                  nullable=False),
+        sa.Column('is_read', sa.Boolean(name='is_read'), server_default='0', nullable=False),
+        sa.Column('content', sa.Text(), nullable=True),
+        sa.Column('source_id', sa.Integer(), nullable=True),
+        sa.ForeignKeyConstraint(['source_id'], ['sources.id'],
+                                name=op.f('fk_messages_source_id_sources')),
+        sa.PrimaryKeyConstraint('id', name=op.f('pk_messages')),
+        sa.UniqueConstraint('uuid', name=op.f('uq_messages_uuid')),
+        sa.CheckConstraint('CASE WHEN is_downloaded = 0 THEN content IS NULL ELSE 1 END',
+                           name='compare_download_vs_content'),
+    )
+
     op.create_table(
         'replies',
         sa.Column('id', sa.Integer(), nullable=False),
@@ -58,25 +95,11 @@ def upgrade():
         sa.PrimaryKeyConstraint('id', name=op.f('pk_replies')),
         sa.UniqueConstraint('uuid', name=op.f('uq_replies_uuid'))
     )
-    op.create_table(
-        'submissions',
-        sa.Column('id', sa.Integer(), nullable=False),
-        sa.Column('uuid', sa.String(length=36), nullable=False),
-        sa.Column('filename', sa.String(length=255), nullable=False),
-        sa.Column('size', sa.Integer(), nullable=False),
-        sa.Column('download_url', sa.String(length=255), nullable=False),
-        sa.Column('is_downloaded', sa.Boolean(name='is_downloaded'), nullable=True),
-        sa.Column('is_read', sa.Boolean(name='is_read'), nullable=True),
-        sa.Column('source_id', sa.Integer(), nullable=True),
-        sa.ForeignKeyConstraint(['source_id'], ['sources.id'],
-                                name=op.f('fk_submissions_source_id_sources')),
-        sa.PrimaryKeyConstraint('id', name=op.f('pk_submissions')),
-        sa.UniqueConstraint('uuid', name=op.f('uq_submissions_uuid'))
-    )
 
 
 def downgrade():
-    op.drop_table('submissions')
     op.drop_table('replies')
+    op.drop_table('messages')
+    op.drop_table('files')
     op.drop_table('users')
     op.drop_table('sources')

From 29fe11f0c06545faff72daa4f67e88a29fbe9059 Mon Sep 17 00:00:00 2001
From: heartsucker <heartsucker@autistici.org>
Date: Tue, 5 Mar 2019 17:03:48 +0100
Subject: [PATCH 5/6] update storage to split submissions into files/messages

---
 securedrop_client/db.py           | 18 ++-----
 securedrop_client/logic.py        | 10 ++--
 securedrop_client/message_sync.py | 12 +++--
 securedrop_client/storage.py      | 90 ++++++++++++++++++++-----------
 4 files changed, 74 insertions(+), 56 deletions(-)

diff --git a/securedrop_client/db.py b/securedrop_client/db.py
index 02ec1325a..f2b993f5b 100644
--- a/securedrop_client/db.py
+++ b/securedrop_client/db.py
@@ -2,7 +2,7 @@
 
 from sqlalchemy import Boolean, Column, create_engine, DateTime, ForeignKey, Integer, String, \
     Text, MetaData, CheckConstraint
-from sqlalchemy.ext.declarative import declarative_base, AbstractConcreteBase
+from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import relationship, backref
 
 
@@ -65,17 +65,9 @@ def collection(self):
         return collection
 
 
-class Submission(AbstractConcreteBase, Base):
-    pass
-
-
-class Message(Submission):
+class Message(Base):
 
     __tablename__ = 'messages'
-    __mapper_args__ = {
-        'polymorphic_identity': 'message',
-        'concrete': True,
-    }
 
     id = Column(Integer, primary_key=True)
     uuid = Column(String(36), unique=True, nullable=False)
@@ -105,13 +97,9 @@ def __repr__(self):
         return '<Message {}>'.format(self.filename)
 
 
-class File(Submission):
+class File(Base):
 
     __tablename__ = 'files'
-    __mapper_args__ = {
-        'polymorphic_identity': 'file',
-        'concrete': True,
-    }
 
     id = Column(Integer, primary_key=True)
     uuid = Column(String(36), unique=True, nullable=False)
diff --git a/securedrop_client/logic.py b/securedrop_client/logic.py
index 70fd04a4c..967b8d1b4 100644
--- a/securedrop_client/logic.py
+++ b/securedrop_client/logic.py
@@ -526,16 +526,14 @@ def set_status(self, message, duration=5000):
         """
         self.gui.set_status(message, duration)
 
-    def on_file_open(self, submission_db_object):
+    def on_file_open(self, file_db_object):
         """
-        Open the already downloaded file associated with the message (which
-        is a Submission).
+        Open the already downloaded file associated with the message (which is a `File`).
         """
-
         # Once downloaded, submissions are stored in the data directory
         # with the same filename as the server, except with the .gz.gpg
         # stripped off.
-        server_filename = submission_db_object.filename
+        server_filename = file_db_object.filename
         fn_no_ext, _ = os.path.splitext(os.path.splitext(server_filename)[0])
         submission_filepath = os.path.join(self.data_dir, fn_no_ext)
 
@@ -561,7 +559,7 @@ def on_file_download(self, source_db_object, message):
             self.on_action_requiring_login()
             return
 
-        if isinstance(message, db.Submission):
+        if isinstance(message, db.File) or isinstance(message, db.Message):
             # Handle submissions.
             func = self.api.download_submission
             sdk_object = sdclientapi.Submission(uuid=message.uuid)
diff --git a/securedrop_client/message_sync.py b/securedrop_client/message_sync.py
index 15f89e7fc..a798543ac 100644
--- a/securedrop_client/message_sync.py
+++ b/securedrop_client/message_sync.py
@@ -26,7 +26,7 @@
 from PyQt5.QtCore import QObject, pyqtSignal
 from securedrop_client import storage
 from securedrop_client.crypto import GpgHelper
-from securedrop_client.db import make_engine
+from securedrop_client.db import make_engine, Message
 from securedrop_client.storage import get_data
 from sqlalchemy.orm import sessionmaker
 
@@ -70,7 +70,8 @@ def __init__(self, api, home, is_qubes):
 
     def run(self, loop=True):
         while True:
-            submissions = storage.find_new_submissions(self.session)
+            submissions = storage.find_new_messages(self.session)
+            submissions.extend(storage.find_new_files(self.session))
 
             for db_submission in submissions:
                 try:
@@ -81,11 +82,16 @@ def run(self, loop=True):
                     # Need to set filename on non-Qubes platforms
                     sdk_submission.filename = db_submission.filename
 
+                    if isinstance(db_submission, Message):
+                        callback = storage.mark_message_as_downloaded
+                    else:
+                        callback = storage.mark_file_as_downloaded
+
                     if self.api:
                         self.fetch_the_thing(sdk_submission,
                                              db_submission,
                                              self.api.download_submission,
-                                             storage.mark_file_as_downloaded)
+                                             callback)
                         self.message_downloaded.emit(db_submission.uuid,
                                                      get_data(self.home, db_submission.filename))
                 except Exception:
diff --git a/securedrop_client/storage.py b/securedrop_client/storage.py
index 508659943..20ea52cb8 100644
--- a/securedrop_client/storage.py
+++ b/securedrop_client/storage.py
@@ -23,7 +23,7 @@
 from dateutil.parser import parse
 import glob
 import os
-from securedrop_client.db import Source, Submission, Reply, User
+from securedrop_client.db import Source, Message, File, Reply, User
 
 
 logger = logging.getLogger(__name__)
@@ -36,11 +36,18 @@ def get_local_sources(session):
     return session.query(Source)
 
 
-def get_local_submissions(session):
+def get_local_messages(session):
     """
     Return all submission objects from the local database.
     """
-    return session.query(Submission)
+    return session.query(Message)
+
+
+def get_local_files(session):
+    """
+    Return all file (a submitted file) objects from the local database.
+    """
+    return session.query(File)
 
 
 def get_local_replies(session):
@@ -85,11 +92,16 @@ def update_local_storage(session, remote_sources, remote_submissions,
     with this data.
     """
     local_sources = get_local_sources(session)
-    local_submissions = get_local_submissions(session)
+    local_files = get_local_files(session)
+    local_messages = get_local_messages(session)
     local_replies = get_local_replies(session)
 
+    remote_messages = [x for x in remote_submissions if x.filename.endswith('.msg.gpg')]
+    remote_files = [x for x in remote_submissions if not x.filename.endswith('.msg.gpg')]
+
     update_sources(remote_sources, local_sources, session, data_dir)
-    update_submissions(remote_submissions, local_submissions, session, data_dir)
+    update_files(remote_files, local_files, session, data_dir)
+    update_messages(remote_messages, local_messages, session, data_dir)
     update_replies(remote_replies, local_replies, session, data_dir)
 
 
@@ -147,8 +159,19 @@ def update_sources(remote_sources, local_sources, session, data_dir):
     session.commit()
 
 
-def update_submissions(remote_submissions, local_submissions, session, data_dir):
+def update_files(remote_submissions, local_submissions, session, data_dir):
+    __update_submissions(File, remote_submissions, local_submissions, session, data_dir)
+
+
+def update_messages(remote_submissions, local_submissions, session, data_dir):
+    __update_submissions(Message, remote_submissions, local_submissions, session, data_dir)
+
+
+def __update_submissions(model, remote_submissions, local_submissions, session, data_dir):
     """
+    The logic for updating files and messages is effectively the same, so this function is somewhat
+    overloaded to allow us to do both in a DRY way.
+
     * Existing submissions are updated in the local database.
     * New submissions have an entry created in the local database.
     * Local submissions not returned in the remote submissions are deleted
@@ -159,10 +182,12 @@ def update_submissions(remote_submissions, local_submissions, session, data_dir)
         if submission.uuid in local_uuids:
             local_submission = [s for s in local_submissions
                                 if s.uuid == submission.uuid][0]
+
             # Update files on disk to match new filename.
             if (local_submission.filename != submission.filename):
                 rename_file(data_dir, local_submission.filename,
                             submission.filename)
+
             # Update an existing record.
             local_submission.filename = submission.filename
             local_submission.size = submission.size
@@ -177,10 +202,8 @@ def update_submissions(remote_submissions, local_submissions, session, data_dir)
             # A new submission to be added to the database.
             _, source_uuid = submission.source_url.rsplit('/', 1)
             source = session.query(Source).filter_by(uuid=source_uuid)[0]
-            ns = Submission(source=source, uuid=submission.uuid,
-                            size=submission.size,
-                            filename=submission.filename,
-                            download_url=submission.download_url)
+            ns = model(source_id=source.id, uuid=submission.uuid, size=submission.size,
+                       filename=submission.filename, download_url=submission.download_url)
             session.add(ns)
             logger.debug('Added new submission {}'.format(submission.uuid))
 
@@ -270,40 +293,43 @@ def find_or_create_user(uuid, username, session):
         return new_user
 
 
-def find_new_submissions(session):
-    submissions = session.query(Submission) \
-                         .filter_by(is_downloaded=False) \
-                         .filter(Submission.filename.like('%-msg.gpg')) \
-                         .all()
-    return submissions
+def find_new_messages(session):
+    return session.query(Message).filter_by(is_downloaded=False).all()
+
+
+def find_new_files(session):
+    return session.query(File).filter_by(is_downloaded=False).all()
 
 
 def find_new_replies(session):
-    replies = session.query(Reply) \
-                     .filter_by(is_downloaded=False) \
-                     .all()
-    return replies
+    return session.query(Reply).filter_by(is_downloaded=False).all()
 
 
 def mark_file_as_downloaded(uuid, session):
-    """Mark file as downloaded in the database. The file itself will be
-    stored in the data directory.
     """
-    submission_db_object = session.query(Submission) \
-                                  .filter_by(uuid=uuid) \
-                                  .one_or_none()
-    submission_db_object.is_downloaded = True
-    session.add(submission_db_object)
+    Mark file as downloaded in the database.
+    """
+    file_db_object = session.query(File).filter_by(uuid=uuid).one()
+    file_db_object.is_downloaded = True
+    session.add(file_db_object)
+    session.commit()
+
+
+def mark_message_as_downloaded(uuid, session):
+    """
+    Mark message as downloaded in the database.
+    """
+    message_db_object = session.query(Message).filter_by(uuid=uuid).one()
+    message_db_object.is_downloaded = True
+    session.add(message_db_object)
     session.commit()
 
 
 def mark_reply_as_downloaded(uuid, session):
-    """Mark reply as downloaded in the database. The file itself will be
-    stored in the data directory.
     """
-    reply_db_object = session.query(Reply) \
-                             .filter_by(uuid=uuid) \
-                             .one_or_none()
+    Mark reply as downloaded in the database.
+    """
+    reply_db_object = session.query(Reply).filter_by(uuid=uuid).one()
     reply_db_object.is_downloaded = True
     session.add(reply_db_object)
     session.commit()

From 9d37899f44fa403e6f78a36455593cbf2f318036 Mon Sep 17 00:00:00 2001
From: heartsucker <heartsucker@autistici.org>
Date: Tue, 5 Mar 2019 17:04:05 +0100
Subject: [PATCH 6/6] update tests

---
 tests/gui/test_main.py     |  11 ++--
 tests/gui/test_widgets.py  |  36 +++++-------
 tests/test_logic.py        |  14 ++---
 tests/test_message_sync.py |  24 +++++---
 tests/test_models.py       |  31 ++++++----
 tests/test_storage.py      | 112 +++++++++++++++++++++++--------------
 6 files changed, 129 insertions(+), 99 deletions(-)

diff --git a/tests/gui/test_main.py b/tests/gui/test_main.py
index 5ae6ab7a8..c6788a522 100644
--- a/tests/gui/test_main.py
+++ b/tests/gui/test_main.py
@@ -4,7 +4,7 @@
 from PyQt5.QtWidgets import QApplication, QVBoxLayout
 from securedrop_client.gui.main import Window
 from securedrop_client.resources import load_icon
-from securedrop_client.db import Submission
+from securedrop_client.db import Message
 from uuid import uuid4
 
 
@@ -245,13 +245,10 @@ def test_conversation_pending_message(mocker):
     mock_source.journalistic_designation = 'Testy McTestface'
 
     msg_uuid = str(uuid4())
-    submission = Submission(source=mock_source, uuid=msg_uuid, size=123,
-                            filename="test.msg.gpg",
-                            download_url='http://test/test')
+    message = Message(source=mock_source, uuid=msg_uuid, size=123, filename="test.msg.gpg",
+                      download_url='http://test/test', is_downloaded=False)
 
-    submission.is_downloaded = False
-
-    mock_source.collection = [submission]
+    mock_source.collection = [message]
 
     mocked_add_message = mocker.patch('securedrop_client.gui.widgets.ConversationView.add_message')
     mocker.patch('securedrop_client.gui.main.QVBoxLayout')
diff --git a/tests/gui/test_widgets.py b/tests/gui/test_widgets.py
index db939dfe3..c6f1aa00d 100644
--- a/tests/gui/test_widgets.py
+++ b/tests/gui/test_widgets.py
@@ -682,12 +682,10 @@ def test_FileWidget_init_left(mocker):
     """
     mock_controller = mocker.MagicMock()
     source = factory.Source()
-    submission = db.Submission(source, 'submission-uuid', 123,
-                               'mah-reply.gpg',
-                               'http://mah-server/mah-reply-url')
-    submission.is_downloaded = True
+    message = db.Message(source=source, uuid='uuid', size=123, filename='mah-reply.gpg',
+                         download_url='http://mah-server/mah-reply-url', is_downloaded=True)
 
-    fw = FileWidget(source, submission, mock_controller, align='left')
+    fw = FileWidget(source, message, mock_controller, align='left')
 
     layout = fw.layout()
     assert isinstance(layout.takeAt(0), QWidgetItem)
@@ -702,12 +700,10 @@ def test_FileWidget_init_right(mocker):
     """
     mock_controller = mocker.MagicMock()
     source = factory.Source()
-    submission = db.Submission(source, 'submission-uuid', 123,
-                               'mah-reply.gpg',
-                               'http://mah-server/mah-reply-url')
-    submission.is_downloaded = True
+    message = db.Message(source=source, uuid='uuid', size=123, filename='mah-reply.gpg',
+                         download_url='http://mah-server/mah-reply-url', is_downloaded=True)
 
-    fw = FileWidget(source, submission, mock_controller, align='right')
+    fw = FileWidget(source, message, mock_controller, align='right')
     layout = fw.layout()
     assert isinstance(layout.takeAt(0), QSpacerItem)
     assert isinstance(layout.takeAt(0), QWidgetItem)
@@ -721,14 +717,12 @@ def test_FileWidget_mousePressEvent_download(mocker):
     """
     mock_controller = mocker.MagicMock()
     source = factory.Source()
-    submission = db.Submission(source, 'submission-uuid', 123,
-                               'mah-reply.gpg',
-                               'http://mah-server/mah-reply-url')
-    submission.is_downloaded = False
+    file_ = db.File(source=source, uuid='uuid', size=123, filename='mah-reply.gpg',
+                    download_url='http://mah-server/mah-reply-url', is_downloaded=False)
 
-    fw = FileWidget(source, submission, mock_controller)
+    fw = FileWidget(source, file_, mock_controller)
     fw.mouseReleaseEvent(None)
-    fw.controller.on_file_download.assert_called_once_with(source, submission)
+    fw.controller.on_file_download.assert_called_once_with(source, file_)
 
 
 def test_FileWidget_mousePressEvent_open(mocker):
@@ -737,14 +731,12 @@ def test_FileWidget_mousePressEvent_open(mocker):
     """
     mock_controller = mocker.MagicMock()
     source = factory.Source()
-    submission = db.Submission(source, 'submission-uuid', 123,
-                               'mah-reply.gpg',
-                               'http://mah-server/mah-reply-url')
-    submission.is_downloaded = True
+    file_ = db.File(source=source, uuid='uuid', size=123, filename='mah-reply.gpg',
+                    download_url='http://mah-server/mah-reply-url', is_downloaded=True)
 
-    fw = FileWidget(source, submission, mock_controller)
+    fw = FileWidget(source, file_, mock_controller)
     fw.mouseReleaseEvent(None)
-    fw.controller.on_file_open.assert_called_once_with(submission)
+    fw.controller.on_file_open.assert_called_once_with(file_)
 
 
 def test_ConversationView_init(mocker, homedir):
diff --git a/tests/test_logic.py b/tests/test_logic.py
index 5f08d8db1..98dbcfcb2 100644
--- a/tests/test_logic.py
+++ b/tests/test_logic.py
@@ -951,18 +951,18 @@ def test_Client_on_file_download_Submission(homedir, config, mocker):
     mock_session = mocker.MagicMock()
     cl = Client('http://localhost', mock_gui, mock_session, homedir)
     source = factory.Source()
-    submission = db.Submission(source, 'submission-uuid', 1234,
-                               'myfile.doc.gpg', 'http://myserver/myfile')
+    file_ = db.File(source=source, uuid='uuid', size=1234, filename='myfile.doc.gpg',
+                    download_url='http://myserver/myfile', is_downloaded=False)
     cl.call_api = mocker.MagicMock()
     cl.api = mocker.MagicMock()
     submission_sdk_object = mocker.MagicMock()
     mock_submission = mocker.patch('sdclientapi.Submission')
     mock_submission.return_value = submission_sdk_object
-    cl.on_file_download(source, submission)
+    cl.on_file_download(source, file_)
     cl.call_api.assert_called_once_with(
         cl.api.download_submission, cl.on_file_downloaded,
         cl.on_download_timeout, submission_sdk_object,
-        cl.data_dir, current_object=submission)
+        cl.data_dir, current_object=file_)
 
 
 def test_Client_on_file_downloaded_success(homedir, config, mocker):
@@ -1048,14 +1048,14 @@ def test_Client_on_file_download_user_not_signed_in(homedir, config, mocker):
     mock_session = mocker.MagicMock()
     cl = Client('http://localhost', mock_gui, mock_session, homedir)
     source = factory.Source()
-    submission = db.Submission(source, 'submission-uuid', 1234,
-                               'myfile.doc.gpg', 'http://myserver/myfile')
+    file_ = db.File(source=source, uuid='uuid', size=1234, filename='myfile.doc.gpg',
+                    download_url='http://myserver/myfile', is_downloaded=False)
     cl.on_action_requiring_login = mocker.MagicMock()
     cl.api = None
     submission_sdk_object = mocker.MagicMock()
     mock_submission = mocker.patch('sdclientapi.Submission')
     mock_submission.return_value = submission_sdk_object
-    cl.on_file_download(source, submission)
+    cl.on_file_download(source, file_)
     cl.on_action_requiring_login.assert_called_once_with()
 
 
diff --git a/tests/test_message_sync.py b/tests/test_message_sync.py
index b7159add6..5fe613e20 100644
--- a/tests/test_message_sync.py
+++ b/tests/test_message_sync.py
@@ -2,6 +2,7 @@
 Make sure the message sync object behaves as expected.
 """
 from securedrop_client.crypto import CryptoError
+from securedrop_client.db import Message, File
 from securedrop_client.message_sync import MessageSync, ReplySync
 
 
@@ -28,15 +29,17 @@ def test_MessageSync_init(mocker):
     assert ms.session == mock_session_class()
 
 
-def test_MessageSync_run_success(mocker):
-    submission = mocker.MagicMock()
-    submission.uuid = 'mock id'
-    submission.download_url = "http://foo"
-    submission.filename = "foo.gpg"
+def test_MessageSync_run_success(mocker, source):
+    submission1 = File(source=source['source'], uuid='uuid1', filename='filename',
+                       download_url='http://test.net')
+    submission2 = Message(source=source['source'], uuid='uuid2', filename='filename',
+                          download_url='http://test.net')
 
     # mock the fetching of submissions
-    mocker.patch('securedrop_client.storage.find_new_submissions', return_value=[submission])
+    mocker.patch('securedrop_client.storage.find_new_messages', return_value=[submission1])
+    mocker.patch('securedrop_client.storage.find_new_files', return_value=[submission2])
     mocker.patch('securedrop_client.message_sync.storage.mark_file_as_downloaded')
+    mocker.patch('securedrop_client.message_sync.storage.mark_message_as_downloaded')
     # don't create the signal
     mocker.patch('securedrop_client.message_sync.pyqtSignal')
     # mock the GpgHelper creation since we don't have directories/keys setup
@@ -72,7 +75,8 @@ def test_MessageSync_exception(homedir, config, mocker):
     is_qubes = False
 
     # mock to return the submission we want
-    mocker.patch('securedrop_client.storage.find_new_submissions', return_value=[submission])
+    mocker.patch('securedrop_client.storage.find_new_messages', return_value=[submission])
+    mocker.patch('securedrop_client.storage.find_new_files', return_value=[])
     # mock to prevent GpgHelper from raising errors on init
     mocker.patch('securedrop_client.crypto.safe_mkdir')
 
@@ -89,9 +93,11 @@ def test_MessageSync_run_failure(mocker):
     submission.filename = "foo.gpg"
 
     # mock the fetching of submissions
-    mocker.patch('securedrop_client.storage.find_new_submissions', return_value=[submission])
+    mocker.patch('securedrop_client.storage.find_new_messages', return_value=[submission])
+    mocker.patch('securedrop_client.storage.find_new_files', return_value=[])
     # mock the handling of the downloaded files
     mocker.patch('securedrop_client.message_sync.storage.mark_file_as_downloaded')
+    mocker.patch('securedrop_client.message_sync.storage.mark_message_as_downloaded')
     # mock the GpgHelper creation since we don't have directories/keys setup
     mocker.patch('securedrop_client.message_sync.GpgHelper')
 
@@ -171,7 +177,7 @@ def test_ReplySync_run_failure(mocker):
     # mock finding new replies
     mocker.patch('securedrop_client.storage.find_new_replies', return_value=[reply])
     # mock handling the new reply
-    mocker.patch('securedrop_client.message_sync.storage.mark_file_as_downloaded')
+    mocker.patch('securedrop_client.message_sync.storage.mark_reply_as_downloaded')
     mocker.patch('securedrop_client.message_sync.GpgHelper')
 
     api = mocker.MagicMock()
diff --git a/tests/test_models.py b/tests/test_models.py
index 99af72288..1be8a4d33 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,5 +1,5 @@
 from tests import factory
-from securedrop_client.db import Reply, Submission, User
+from securedrop_client.db import Reply, File, Message, User
 
 
 def test_string_representation_of_user():
@@ -12,12 +12,18 @@ def test_string_representation_of_source():
     source.__repr__()
 
 
-def test_string_representation_of_submission():
+def test_string_representation_of_message():
     source = factory.Source()
-    submission = Submission(source=source, uuid="test", size=123,
-                            filename="test.docx",
-                            download_url='http://test/test')
-    submission.__repr__()
+    msg = Message(source=source, uuid="test", size=123, filename="test.docx",
+                  download_url='http://test/test')
+    msg.__repr__()
+
+
+def test_string_representation_of_file():
+    source = factory.Source()
+    file_ = File(source=source, uuid="test", size=123, filename="test.docx",
+                 download_url='http://test/test')
+    file_.__repr__()
 
 
 def test_string_representation_of_reply():
@@ -31,15 +37,18 @@ def test_string_representation_of_reply():
 def test_source_collection():
     # Create some test submissions and replies
     source = factory.Source()
-    submission = Submission(source=source, uuid="test", size=123,
-                            filename="2-test.doc.gpg",
-                            download_url='http://test/test')
+    file_ = File(source=source, uuid="test", size=123, filename="2-test.doc.gpg",
+                 download_url='http://test/test')
+    message = Message(source=source, uuid="test", size=123, filename="3-test.doc.gpg",
+                      download_url='http://test/test')
     user = User('hehe')
     reply = Reply(source=source, journalist=user, filename="1-reply.gpg",
                   size=1234, uuid='test')
-    source.submissions = [submission]
+    source.files = [file_]
+    source.messages = [message]
     source.replies = [reply]
 
     # Now these items should be in the source collection in the proper order
     assert source.collection[0] == reply
-    assert source.collection[1] == submission
+    assert source.collection[1] == file_
+    assert source.collection[2] == message
diff --git a/tests/test_storage.py b/tests/test_storage.py
index 64b26c39b..8fea2ff4b 100644
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -6,10 +6,11 @@
 import uuid
 import securedrop_client.db
 from dateutil.parser import parse
-from securedrop_client.storage import get_local_sources, get_local_submissions, get_local_replies, \
-    get_remote_data, update_local_storage, update_sources, update_submissions, update_replies, \
-    find_or_create_user, find_new_submissions, find_new_replies, mark_file_as_downloaded, \
-    mark_reply_as_downloaded, delete_single_submission_or_reply_on_disk, get_data, rename_file
+from securedrop_client.storage import get_local_sources, get_local_messages, get_local_replies, \
+    get_remote_data, update_local_storage, update_sources, update_files, update_replies, \
+    find_or_create_user, find_new_messages, find_new_replies, mark_file_as_downloaded, \
+    mark_reply_as_downloaded, delete_single_submission_or_reply_on_disk, get_data, rename_file, \
+    get_local_files, find_new_files, mark_message_as_downloaded
 from securedrop_client import db
 from sdclientapi import Source, Submission, Reply
 
@@ -65,14 +66,24 @@ def test_get_local_sources(mocker):
     mock_session.query.assert_called_once_with(securedrop_client.db.Source)
 
 
-def test_get_local_submissions(mocker):
+def test_get_local_messages(mocker):
     """
-    At this moment, just return all submissions.
+    At this moment, just return all messages.
     """
     mock_session = mocker.MagicMock()
-    get_local_submissions(mock_session)
+    get_local_messages(mock_session)
     mock_session.query.\
-        assert_called_once_with(securedrop_client.db.Submission)
+        assert_called_once_with(securedrop_client.db.Message)
+
+
+def test_get_local_files(mocker):
+    """
+    At this moment, just return all files (submissions)
+    """
+    mock_session = mocker.MagicMock()
+    get_local_files(mock_session)
+    mock_session.query.\
+        assert_called_once_with(securedrop_client.db.File)
 
 
 def test_get_local_replies(mocker):
@@ -118,31 +129,29 @@ def test_update_local_storage(homedir, mocker):
     Assuming no errors getting data, check the expected functions to update
     the state of the local database are called with the necessary data.
     """
-    source = make_remote_source()
-    submission = mocker.MagicMock()
-    reply = mocker.MagicMock()
-    sources = [source, ]
-    submissions = [submission, ]
-    replies = [reply, ]
+    remote_source = make_remote_source()
+    remote_message = mocker.Mock(filename='foo.msg.gpg')
+    remote_file = mocker.Mock(filename='foo.gpg')
+    remote_submissions = [remote_message, remote_file]
+    remote_reply = mocker.MagicMock()
     # Some local source, submission and reply objects from the local database.
     mock_session = mocker.MagicMock()
     local_source = mocker.MagicMock()
-    local_submission = mocker.MagicMock()
-    local_replies = mocker.MagicMock()
-    mock_session.query.side_effect = [[local_source, ], [local_submission, ],
-                                      [local_replies, ], ]
+    local_file = mocker.MagicMock()
+    local_message = mocker.MagicMock()
+    local_reply = mocker.MagicMock()
+    mock_session.query.side_effect = [
+        [local_source], [local_file], [local_message], [local_reply]]
     src_fn = mocker.patch('securedrop_client.storage.update_sources')
     rpl_fn = mocker.patch('securedrop_client.storage.update_replies')
-    sub_fn = mocker.patch('securedrop_client.storage.update_submissions')
+    file_fn = mocker.patch('securedrop_client.storage.update_files')
+    msg_fn = mocker.patch('securedrop_client.storage.update_messages')
 
-    update_local_storage(mock_session, sources, submissions, replies,
-                         homedir)
-    src_fn.assert_called_once_with([source, ], [local_source, ],
-                                   mock_session, homedir)
-    rpl_fn.assert_called_once_with([reply, ], [local_replies, ],
-                                   mock_session, homedir)
-    sub_fn.assert_called_once_with([submission, ], [local_submission, ],
-                                   mock_session, homedir)
+    update_local_storage(mock_session, [remote_source], remote_submissions, [remote_reply], homedir)
+    src_fn.assert_called_once_with([remote_source], [local_source], mock_session, homedir)
+    rpl_fn.assert_called_once_with([remote_reply], [local_reply], mock_session, homedir)
+    file_fn.assert_called_once_with([remote_file], [local_file], mock_session, homedir)
+    msg_fn.assert_called_once_with([remote_message], [local_message], mock_session, homedir)
 
 
 def test_update_sources(homedir, mocker):
@@ -198,7 +207,7 @@ def test_update_sources(homedir, mocker):
     assert mock_session.commit.call_count == 1
 
 
-def test_update_submissions_renames_file_on_disk(homedir, mocker):
+def test_update_files_renames_file_on_disk(homedir, mocker):
     """
     Check that:
 
@@ -227,8 +236,7 @@ def test_update_submissions_renames_file_on_disk(homedir, mocker):
     local_source.id = 123
     mock_session.query().filter_by.return_value = [local_source, ]
 
-    update_submissions(remote_submissions, local_submissions, mock_session,
-                       data_dir)
+    update_files(remote_submissions, local_submissions, mock_session, data_dir)
 
     updated_local_filename = '1-spotted-potato-msg'
     assert local_submission.filename == remote_submission.filename
@@ -319,8 +327,7 @@ def test_update_submissions_deletes_files_associated_with_the_submission(
     local_source.uuid = 'test-source-uuid'
     local_source.id = 666
     mock_session.query().filter_by.return_value = [local_source, ]
-    update_submissions(remote_submissions, local_submissions, mock_session,
-                       homedir)
+    update_files(remote_submissions, local_submissions, mock_session, homedir)
 
     # Ensure the files associated with the submission are deleted on disk.
     assert not os.path.exists(abs_server_filename)
@@ -407,10 +414,10 @@ def test_update_sources_deletes_files_associated_with_the_source(
 
     # Here we're not mocking out the models use so that we can use the collection attribute.
     local_source = factory.Source()
-    file_submission = db.Submission(
+    file_submission = db.File(
         source=local_source, uuid="test", size=123, filename=file_server_filename,
         download_url='http://test/test')
-    msg_submission = db.Submission(
+    msg_submission = db.File(
         source=local_source, uuid="test", size=123, filename=msg_server_filename,
         download_url='http://test/test')
     user = db.User('hehe')
@@ -445,7 +452,7 @@ def test_update_sources_deletes_files_associated_with_the_source(
     assert mock_session.commit.call_count == 1
 
 
-def test_update_submissions(homedir, mocker):
+def test_update_files(homedir, mocker):
     """
     Check that:
 
@@ -486,8 +493,7 @@ def test_update_submissions(homedir, mocker):
     mock_session.query().filter_by.return_value = [local_source, ]
     patch_rename_file = mocker.patch('securedrop_client.storage.rename_file')
 
-    update_submissions(remote_submissions, local_submissions, mock_session,
-                       data_dir)
+    update_files(remote_submissions, local_submissions, mock_session, data_dir)
 
     # Check the expected local submission object has been updated with values
     # from the API.
@@ -627,14 +633,23 @@ def test_find_or_create_user_new(mocker):
     mock_session.commit.assert_called_once_with()
 
 
-def test_find_new_submissions(mocker):
+def test_find_new_messages(mocker):
     mock_session = mocker.MagicMock()
     mock_submission = mocker.MagicMock()
     mock_submission.is_downloaded = False
     mock_submissions = [mock_submission]
-    mock_session.query().filter_by() \
-                        .filter().all.return_value = mock_submissions
-    submissions = find_new_submissions(mock_session)
+    mock_session.query().filter_by().all.return_value = mock_submissions
+    submissions = find_new_messages(mock_session)
+    assert submissions[0].is_downloaded is False
+
+
+def test_find_new_files(mocker):
+    mock_session = mocker.MagicMock()
+    mock_submission = mocker.MagicMock()
+    mock_submission.is_downloaded = False
+    mock_submissions = [mock_submission]
+    mock_session.query().filter_by().all.return_value = mock_submissions
+    submissions = find_new_files(mock_session)
     assert submissions[0].is_downloaded is False
 
 
@@ -653,18 +668,29 @@ def test_mark_file_as_downloaded(mocker):
     mock_session = mocker.MagicMock()
     mock_submission = mocker.MagicMock()
     mock_submission.is_downloaded is False
-    mock_session.query().filter_by().one_or_none.return_value = mock_submission
+    mock_session.query().filter_by().one.return_value = mock_submission
     mark_file_as_downloaded('test-filename', mock_session)
     assert mock_submission.is_downloaded is True
     mock_session.add.assert_called_once_with(mock_submission)
     mock_session.commit.assert_called_once_with()
 
 
+def test_mark_message_as_downloaded(mocker):
+    mock_session = mocker.MagicMock()
+    mock_submission = mocker.MagicMock()
+    mock_submission.is_downloaded is False
+    mock_session.query().filter_by().one.return_value = mock_submission
+    mark_message_as_downloaded('test-messagename', mock_session)
+    assert mock_submission.is_downloaded is True
+    mock_session.add.assert_called_once_with(mock_submission)
+    mock_session.commit.assert_called_once_with()
+
+
 def test_mark_reply_as_downloaded(mocker):
     mock_session = mocker.MagicMock()
     mock_reply = mocker.MagicMock()
     mock_reply.is_downloaded is False
-    mock_session.query().filter_by().one_or_none.return_value = mock_reply
+    mock_session.query().filter_by().one.return_value = mock_reply
     mark_reply_as_downloaded('test-filename', mock_session)
     assert mock_reply.is_downloaded is True
     mock_session.add.assert_called_once_with(mock_reply)