Merge pull request #418 from dod-ccpo/invited-user-#160301892

Invited user #160301892

atst/routes/__init__.py

@@ -2,6 +2,7 @@ import urllib.parse as url
 from flask import Blueprint, render_template, g, redirect, session, url_for, request
 
 from flask import current_app as app
+from jinja2.exceptions import TemplateNotFound
 import pendulum
 import os
 
@@ -10,6 +11,7 @@ from atst.domain.users import Users
 from atst.domain.authnid import AuthenticationContext
 from atst.domain.audit_log import AuditLog
 from atst.domain.auth import logout as _logout
+from werkzeug.exceptions import NotFound
 
 
 bp = Blueprint("atst", __name__)
@@ -77,7 +79,10 @@ def styleguide():
 
 @bp.route("/<path:path>")
 def catch_all(path):
-    return render_template("{}.html".format(path))
+    try:
+        return render_template("{}.html".format(path))
+    except TemplateNotFound:
+        raise NotFound()
 
 
 def _make_authentication_context():

atst/routes/errors.py

@@ -3,27 +3,35 @@ from flask_wtf.csrf import CSRFError
 import werkzeug.exceptions as werkzeug_exceptions
 
 import atst.domain.exceptions as exceptions
-from atst.domain.invitations import InvitationError
+from atst.domain.invitations import (
+    InvitationError,
+    ExpiredError as InvitationExpiredError,
+    WrongUserError as InvitationWrongUserError,
+)
+
+
+def log_error(e):
+    error_message = e.message if hasattr(e, "message") else str(e)
+    current_app.logger.error(error_message)
+
+
+def handle_error(e, message="Not Found", code=404):
+    log_error(e)
+    return render_template("error.html", message=message), code
 
 
 def make_error_pages(app):
-    def log_error(e):
-        error_message = e.message if hasattr(e, "message") else str(e)
-        app.logger.error(error_message)
-
     @app.errorhandler(werkzeug_exceptions.NotFound)
     @app.errorhandler(exceptions.NotFoundError)
     @app.errorhandler(exceptions.UnauthorizedError)
     # pylint: disable=unused-variable
     def not_found(e):
-        log_error(e)
-        return render_template("error.html", message="Not Found"), 404
+        return handle_error(e)
 
     @app.errorhandler(exceptions.UnauthenticatedError)
     # pylint: disable=unused-variable
     def unauthorized(e):
-        log_error(e)
-        return render_template("error.html", message="Log in Failed"), 401
+        return handle_error(e, message="Log in Failed", code=401)
 
     @app.errorhandler(CSRFError)
     # pylint: disable=unused-variable
@@ -43,14 +51,16 @@ def make_error_pages(app):
         )
 
     @app.errorhandler(InvitationError)
+    @app.errorhandler(InvitationWrongUserError)
     # pylint: disable=unused-variable
     def invalid_invitation(e):
-        log_error(e)
-        return (
-            render_template(
-                "error.html", message="The invitation link you clicked is invalid."
-            ),
-            404,
+        return handle_error(e, message="The link you followed is invalid.", code=404)
+
+    @app.errorhandler(InvitationExpiredError)
+    # pylint: disable=unused-variable
+    def invalid_invitation(e):
+        return handle_error(
+            e, message="The invitation you followed has expired.", code=404
         )
 
     return app

atst/routes/workspaces.py

@@ -361,9 +361,7 @@ def update_member(workspace_id, member_id):
 
 @bp.route("/workspaces/invitation/<token>", methods=["GET"])
 def accept_invitation(token):
-    # TODO: check that the current_user DOD ID matches the user associated with
-    # the invitation
-    invite = Invitations.accept(token)
+    invite = Invitations.accept(g.current_user, token)
 
     return redirect(
         url_for("workspaces.show_workspace", workspace_id=invite.workspace.id)