Add type annotations to journalist_app.api

mypy.ini

@@ -9,5 +9,8 @@ disallow_untyped_defs = True
 [mypy-journalist_app.models]
 disallow_untyped_defs = True
 
+[mypy-journalist_app.api]
+disallow_untyped_defs = True
+
 [mypy-config]
 ignore_errors = True

securedrop/journalist_app/api.py

@@ -1,69 +1,69 @@
 import json
 
 from datetime import datetime, timedelta
+from typing import Tuple, Callable, Any
+
+import flask
+import werkzeug
 from flask import abort, Blueprint, current_app, jsonify, request
 from functools import wraps
+
 from sqlalchemy.exc import IntegrityError
 from os import path
 from uuid import UUID
-from werkzeug.exceptions import default_exceptions  # type: ignore
+from werkzeug.exceptions import default_exceptions
 
 from db import db
 from journalist_app import utils
 from models import (Journalist, Reply, Source, Submission,
                     LoginThrottledException, InvalidUsernameException,
                     BadTokenException, WrongPasswordException)
+from sdconfig import SDConfig
 from store import NotEncrypted
 
 
 TOKEN_EXPIRATION_MINS = 60 * 8
 
 
-def get_user_object(request):
-    """Helper function to use in token_required views that need a user
-    object
-    """
-    auth_token = request.headers.get('Authorization').split(" ")[1]
-    user = Journalist.validate_api_token_and_get_user(auth_token)
-    return user
+def _authenticate_user_from_auth_header(request: flask.Request) -> Journalist:
+    try:
+        auth_header = request.headers['Authorization']
+    except KeyError:
+        return abort(403, 'API token not found in Authorization header.')
+
+    if auth_header:
+        split = auth_header.split(" ")
+        if len(split) != 2 or split[0] != 'Token':
+            abort(403, 'Malformed authorization header.')
+        auth_token = split[1]
+    else:
+        auth_token = ''
+    authenticated_user = Journalist.validate_api_token_and_get_user(auth_token)
+    if not authenticated_user:
+        return abort(403, 'API token is invalid or expired.')
+    return authenticated_user
 
 
-def token_required(f):
+def token_required(f: Callable) -> Callable:
     @wraps(f)
-    def decorated_function(*args, **kwargs):
-        try:
-            auth_header = request.headers['Authorization']
-        except KeyError:
-            return abort(403, 'API token not found in Authorization header.')
-
-        if auth_header:
-            split = auth_header.split(" ")
-            if len(split) != 2 or split[0] != 'Token':
-                abort(403, 'Malformed authorization header.')
-            auth_token = split[1]
-        else:
-            auth_token = ''
-        if not Journalist.validate_api_token_and_get_user(auth_token):
-            return abort(403, 'API token is invalid or expired.')
+    def decorated_function(*args: Any, **kwargs: Any) -> Any:
+        _authenticate_user_from_auth_header(request)
         return f(*args, **kwargs)
     return decorated_function
 
 
-def get_or_404(model, object_id, column=''):
-    if column:
-        result = model.query.filter(column == object_id).one_or_none()
-    else:
-        result = model.query.get(object_id)
+def get_or_404(model: db.Model, object_id: str, column: str) -> db.Model:
+    result = model.query.filter(column == object_id).one_or_none()
     if result is None:
         abort(404)
     return result
 
 
-def make_blueprint(config):
+def make_blueprint(config: SDConfig) -> Blueprint:
     api = Blueprint('api', __name__)
 
     @api.route('/')
-    def get_endpoints():
+    def get_endpoints() -> Tuple[flask.Response, int]:
         endpoints = {'sources_url': '/api/v1/sources',
                      'current_user_url': '/api/v1/user',
                      'submissions_url': '/api/v1/submissions',
@@ -73,7 +73,7 @@ def get_endpoints():
 
     # Before every post, we validate the payload before processing the request
     @api.before_request
-    def validate_data():
+    def validate_data() -> None:
         if request.method == 'POST':
             # flag, star, and logout can have empty payloads
             if not request.data:
@@ -86,28 +86,28 @@ def validate_data():
                 for endpoint in dataless_endpoints:
                     if request.endpoint == 'api.' + endpoint:
                         return
-                return abort(400, 'malformed request')
+                abort(400, 'malformed request')
             # other requests must have valid JSON payload
             else:
                 try:
                     json.loads(request.data.decode('utf-8'))
                 except (ValueError):
-                    return abort(400, 'malformed request')
+                    abort(400, 'malformed request')
 
     @api.route('/token', methods=['POST'])
-    def get_token():
+    def get_token() -> Tuple[flask.Response, int]:
         creds = json.loads(request.data.decode('utf-8'))
 
         username = creds.get('username', None)
         passphrase = creds.get('passphrase', None)
         one_time_code = creds.get('one_time_code', None)
 
         if username is None:
-            return abort(400, 'username field is missing')
+            abort(400, 'username field is missing')
         if passphrase is None:
-            return abort(400, 'passphrase field is missing')
+            abort(400, 'passphrase field is missing')
         if one_time_code is None:
-            return abort(400, 'one_time_code field is missing')
+            abort(400, 'one_time_code field is missing')
 
         try:
             journalist = Journalist.login(username, passphrase, one_time_code)
@@ -134,41 +134,43 @@ def get_token():
 
     @api.route('/sources', methods=['GET'])
     @token_required
-    def get_all_sources():
+    def get_all_sources() -> Tuple[flask.Response, int]:
         sources = Source.query.filter_by(pending=False, deleted_at=None).all()
         return jsonify(
             {'sources': [source.to_json() for source in sources]}), 200
 
     @api.route('/sources/<source_uuid>', methods=['GET', 'DELETE'])
     @token_required
-    def single_source(source_uuid):
+    def single_source(source_uuid: str) -> Tuple[flask.Response, int]:
         if request.method == 'GET':
             source = get_or_404(Source, source_uuid, column=Source.uuid)
             return jsonify(source.to_json()), 200
         elif request.method == 'DELETE':
             source = get_or_404(Source, source_uuid, column=Source.uuid)
             utils.delete_collection(source.filesystem_id)
             return jsonify({'message': 'Source and submissions deleted'}), 200
+        else:
+            abort(404)
 
     @api.route('/sources/<source_uuid>/add_star', methods=['POST'])
     @token_required
-    def add_star(source_uuid):
+    def add_star(source_uuid: str) -> Tuple[flask.Response, int]:
         source = get_or_404(Source, source_uuid, column=Source.uuid)
         utils.make_star_true(source.filesystem_id)
         db.session.commit()
         return jsonify({'message': 'Star added'}), 201
 
     @api.route('/sources/<source_uuid>/remove_star', methods=['DELETE'])
     @token_required
-    def remove_star(source_uuid):
+    def remove_star(source_uuid: str) -> Tuple[flask.Response, int]:
         source = get_or_404(Source, source_uuid, column=Source.uuid)
         utils.make_star_false(source.filesystem_id)
         db.session.commit()
         return jsonify({'message': 'Star removed'}), 200
 
     @api.route('/sources/<source_uuid>/flag', methods=['POST'])
     @token_required
-    def flag(source_uuid):
+    def flag(source_uuid: str) -> Tuple[flask.Response, int]:
         source = get_or_404(Source, source_uuid,
                             column=Source.uuid)
         source.flagged = True
@@ -177,7 +179,7 @@ def flag(source_uuid):
 
     @api.route('/sources/<source_uuid>/submissions', methods=['GET'])
     @token_required
-    def all_source_submissions(source_uuid):
+    def all_source_submissions(source_uuid: str) -> Tuple[flask.Response, int]:
         source = get_or_404(Source, source_uuid, column=Source.uuid)
         return jsonify(
             {'submissions': [submission.to_json() for
@@ -186,7 +188,7 @@ def all_source_submissions(source_uuid):
     @api.route('/sources/<source_uuid>/submissions/<submission_uuid>/download',  # noqa
                methods=['GET'])
     @token_required
-    def download_submission(source_uuid, submission_uuid):
+    def download_submission(source_uuid: str, submission_uuid: str) -> flask.Response:
         get_or_404(Source, source_uuid, column=Source.uuid)
         submission = get_or_404(Submission, submission_uuid,
                                 column=Submission.uuid)
@@ -200,7 +202,7 @@ def download_submission(source_uuid, submission_uuid):
     @api.route('/sources/<source_uuid>/replies/<reply_uuid>/download',
                methods=['GET'])
     @token_required
-    def download_reply(source_uuid, reply_uuid):
+    def download_reply(source_uuid: str, reply_uuid: str) -> flask.Response:
         get_or_404(Source, source_uuid, column=Source.uuid)
         reply = get_or_404(Reply, reply_uuid, column=Reply.uuid)
 
@@ -209,7 +211,7 @@ def download_reply(source_uuid, reply_uuid):
     @api.route('/sources/<source_uuid>/submissions/<submission_uuid>',
                methods=['GET', 'DELETE'])
     @token_required
-    def single_submission(source_uuid, submission_uuid):
+    def single_submission(source_uuid: str, submission_uuid: str) -> Tuple[flask.Response, int]:
         if request.method == 'GET':
             get_or_404(Source, source_uuid, column=Source.uuid)
             submission = get_or_404(Submission, submission_uuid, column=Submission.uuid)
@@ -219,10 +221,12 @@ def single_submission(source_uuid, submission_uuid):
             submission = get_or_404(Submission, submission_uuid, column=Submission.uuid)
             utils.delete_file_object(submission)
             return jsonify({'message': 'Submission deleted'}), 200
+        else:
+            abort(404)
 
     @api.route('/sources/<source_uuid>/replies', methods=['GET', 'POST'])
     @token_required
-    def all_source_replies(source_uuid):
+    def all_source_replies(source_uuid: str) -> Tuple[flask.Response, int]:
         if request.method == 'GET':
             source = get_or_404(Source, source_uuid, column=Source.uuid)
             return jsonify(
@@ -237,7 +241,7 @@ def all_source_replies(source_uuid):
             if 'reply' not in request.json:
                 abort(400, 'reply not found in request body')
 
-            user = get_user_object(request)
+            user = _authenticate_user_from_auth_header(request)
 
             data = request.json
             if not data['reply']:
@@ -282,54 +286,60 @@ def all_source_replies(source_uuid):
             return jsonify({'message': 'Your reply has been stored',
                             'uuid': reply.uuid,
                             'filename': reply.filename}), 201
+        else:
+            abort(404)
 
     @api.route('/sources/<source_uuid>/replies/<reply_uuid>',
                methods=['GET', 'DELETE'])
     @token_required
-    def single_reply(source_uuid, reply_uuid):
+    def single_reply(source_uuid: str, reply_uuid: str) -> Tuple[flask.Response, int]:
         get_or_404(Source, source_uuid, column=Source.uuid)
         reply = get_or_404(Reply, reply_uuid, column=Reply.uuid)
         if request.method == 'GET':
             return jsonify(reply.to_json()), 200
         elif request.method == 'DELETE':
             utils.delete_file_object(reply)
             return jsonify({'message': 'Reply deleted'}), 200
+        else:
+            abort(404)
 
     @api.route('/submissions', methods=['GET'])
     @token_required
-    def get_all_submissions():
+    def get_all_submissions() -> Tuple[flask.Response, int]:
         submissions = Submission.query.all()
         return jsonify({'submissions': [submission.to_json() for
                                         submission in submissions if submission.source]}), 200
 
     @api.route('/replies', methods=['GET'])
     @token_required
-    def get_all_replies():
+    def get_all_replies() -> Tuple[flask.Response, int]:
         replies = Reply.query.all()
         return jsonify(
             {'replies': [reply.to_json() for reply in replies if reply.source]}), 200
 
     @api.route('/user', methods=['GET'])
     @token_required
-    def get_current_user():
-        user = get_user_object(request)
+    def get_current_user() -> Tuple[flask.Response, int]:
+        user = _authenticate_user_from_auth_header(request)
         return jsonify(user.to_json()), 200
 
     @api.route('/logout', methods=['POST'])
     @token_required
-    def logout():
-        user = get_user_object(request)
-        auth_token = request.headers.get('Authorization').split(" ")[1]
+    def logout() -> Tuple[flask.Response, int]:
+        user = _authenticate_user_from_auth_header(request)
+        auth_token = request.headers['Authorization'].split(" ")[1]
         utils.revoke_token(user, auth_token)
         return jsonify({'message': 'Your token has been revoked.'}), 200
 
-    def _handle_api_http_exception(error):
+    def _handle_api_http_exception(
+        error: werkzeug.exceptions.HTTPException
+    ) -> Tuple[flask.Response, int]:
         # Workaround for no blueprint-level 404/5 error handlers, see:
         # https://github.com/pallets/flask/issues/503#issuecomment-71383286
         response = jsonify({'error': error.name,
                            'message': error.description})
 
-        return response, error.code
+        return response, error.code  # type: ignore
 
     for code in default_exceptions:
         api.errorhandler(code)(_handle_api_http_exception)

securedrop/journalist_app/utils.py

@@ -364,7 +364,7 @@ def col_download_all(cols_selected: List[str]) -> werkzeug.Response:
     return download("all", submissions)
 
 
-def serve_file_with_etag(db_obj: Union[Reply, Submission]) -> werkzeug.Response:
+def serve_file_with_etag(db_obj: Union[Reply, Submission]) -> flask.Response:
     file_path = current_app.storage.path(db_obj.source.filesystem_id, db_obj.filename)
     response = send_file(file_path,
                          mimetype="application/pgp-encrypted",

securedrop/models.py

@@ -697,7 +697,7 @@ def validate_token_is_not_expired_or_invalid(token):
         return True
 
     @staticmethod
-    def validate_api_token_and_get_user(token: str) -> 'Union[Journalist, None]':
+    def validate_api_token_and_get_user(token: str) -> 'Optional[Journalist]':
         s = TimedJSONWebSignatureSerializer(current_app.config['SECRET_KEY'])
         try:
             data = s.loads(token)