whyphi · jinyoungbang · Mar 10, 2024 · Mar 4, 2024 · Mar 4, 2024 · Mar 4, 2024
diff --git a/.chalice/policy-dev.json b/.chalice/policy-dev.json
@@ -34,6 +34,7 @@
 				"arn:*:logs:*:*:*",
 				"arn:aws:s3:::whyphi-zap/dev/*",
 				"arn:aws:ses:us-east-1:280776660572:identity/why-phi.com",
+				"arn:aws:ssm:us-east-1:280776660572:parameter/Zap/AUTH_SECRET",
 				"arn:aws:ssm:us-east-1:280776660572:parameter/Zap/MONGO_ADMIN_USER",
 				"arn:aws:ssm:us-east-1:280776660572:parameter/Zap/MONGO_ADMIN_PASSWORD"
 			]

diff --git a/.chalice/policy-prod.json b/.chalice/policy-prod.json
@@ -34,6 +34,7 @@
 				"arn:*:logs:*:*:*",
 				"arn:aws:s3:::whyphi-zap/prod/*",
 				"arn:aws:ses:us-east-1:280776660572:identity/why-phi.com",
+				"arn:aws:ssm:us-east-1:280776660572:parameter/Zap/AUTH_SECRET",
 				"arn:aws:ssm:us-east-1:280776660572:parameter/Zap/MONGO_ADMIN_USER",
 				"arn:aws:ssm:us-east-1:280776660572:parameter/Zap/MONGO_ADMIN_PASSWORD"
 			]

diff --git a/Pipfile b/Pipfile
@@ -12,7 +12,8 @@ pydantic = "*"
 pytest = "*"
 coverage = "*"
 pytest-cov = "*"
-moto = {extras = ["dynamodb"], version = "*"}
+moto = {extras = ["dynamodb"], version = "4.2.13"}
+pyjwt = "*"
 
 [dev-packages]
 

diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/chalicelib/api/applicants.py b/chalicelib/api/applicants.py
@@ -1,23 +1,29 @@
 from chalice import Blueprint
 from chalicelib.services.ApplicantService import applicant_service
 from chalicelib.handlers.error_handler import handle_exceptions
+from chalicelib.decorators import auth
+from chalicelib.models.roles import Roles
+
 from pydantic import ValidationError
 
 applicants_api = Blueprint(__name__)
 
 
 @applicants_api.route("/applicant/{id}", methods=["GET"], cors=True)
+@auth(applicants_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def get_applicant(id):
     """Get an applicant from <applicant_id>"""
     return applicant_service.get(id)
 
 
 @applicants_api.route("/applicants", methods=["GET"], cors=True)
+@auth(applicants_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def get_applicants():
     return applicant_service.get_all()
 
 
 @applicants_api.route("/applicants/{listing_id}", methods=["GET"], cors=True)
+@auth(applicants_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def get_all_applicants(listing_id):
     """Gets all applicants from <listing_id>"""
     return applicant_service.get_all_from_listing(listing_id)
diff --git a/chalicelib/api/insights.py b/chalicelib/api/insights.py
@@ -2,12 +2,16 @@
 from chalice import Blueprint
 from chalicelib.services.InsightsService import insights_service
 from chalicelib.handlers.error_handler import handle_exceptions
+from chalicelib.decorators import auth
+from chalicelib.models.roles import Roles
+
 from pydantic import ValidationError
 
 insights_api = Blueprint(__name__)
 
 
 @insights_api.route("/insights/listing/{listing_id}", methods=["GET"], cors=True)
+@auth(insights_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def get_listing_insights(listing_id):
     """Get insights from <listing_id>"""
-    return insights_service.get_insights_from_listing(listing_id)
+    return insights_service.get_insights_from_listing(listing_id)
diff --git a/chalicelib/api/listings.py b/chalicelib/api/listings.py
@@ -1,6 +1,9 @@
 from chalice import Blueprint
 from chalicelib.services.ListingService import listing_service
 from chalicelib.handlers.error_handler import handle_exceptions
+from chalicelib.decorators import auth
+from chalicelib.models.roles import Roles
+
 from pydantic import ValidationError
 
 listings_api = Blueprint(__name__)
@@ -13,24 +16,28 @@ def apply_to_listing():
 
 
 @listings_api.route("/create", methods=["POST"], cors=True)
+@auth(listings_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def create_listing():
     """Creates a new listing with given information"""
     return listing_service.create(listings_api.current_request.json_body)
 
 
 @listings_api.route("/listings/{id}", methods=["GET"], cors=True)
+@auth(listings_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def get_listing(id):
     """Gets a listing from id"""
     return listing_service.get(id)
 
 
 @listings_api.route("/listings", methods=["GET"], cors=True)
+@auth(listings_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def get_all_listings():
     """Gets all listings available"""
     return listing_service.get_all()
 
 
 @listings_api.route("/listings/{id}", methods=["DELETE"], cors=True)
+@auth(listings_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def delete_listing(id):
     """Deletes a listing with the given ID."""
     try:
@@ -40,6 +47,7 @@ def delete_listing(id):
 
 
 @listings_api.route("/listings/{id}/toggle/visibility", methods=["PATCH"], cors=True)
+@auth(listings_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def toggle_visibility(id):
     """Toggles visibilility of a given <listing_id>"""
     try:
@@ -50,6 +58,7 @@ def toggle_visibility(id):
 
 @listings_api.route("/listings/{id}/update-field", methods=["PATCH"], cors=True)
 @handle_exceptions
+@auth(listings_api, roles=[Roles.ADMIN, Roles.MEMBER])
 def update_listing_field_route(id):
     try:
         return listing_service.update_field_route(

diff --git a/chalicelib/api/members.py b/chalicelib/api/members.py
@@ -1,10 +1,12 @@
 from chalice import Blueprint
 from chalicelib.services.MemberService import member_service
+from chalicelib.decorators import auth
 
 members_api = Blueprint(__name__)
 
 
 @members_api.route("/members", methods=["GET"], cors=True)
+@auth(members_api, roles=["admin", "member"])
 def get_all_members():
     """Get an applicant from <applicant_id>"""
     return member_service.get_all()
diff --git a/chalicelib/decorators.py b/chalicelib/decorators.py
@@ -1,4 +1,22 @@
+import boto3
+import jwt
+from chalice import UnauthorizedError
+
+
 def add_env_suffix(func):
+    """
+    Decorator for adding an environment suffix to a table name based on the provided 'env' flag.
+
+    Args:
+        func (function): The original function to be decorated.
+
+    Returns:
+        function: A wrapper function that modifies the input 'table_name' based on the 'env' flag.
+
+    - When calling my_function with env=True, the table_name will be suffixed with '-prod'.
+    - When calling my_function with env=False or without the 'env' flag, the table_name will be suffixed with '-dev'.
+    """
+
     def wrapper(self, table_name: str, *args, **kwargs):
         if "env" in kwargs and kwargs["env"]:
             table_name += "-prod"
@@ -7,4 +25,52 @@ def wrapper(self, table_name: str, *args, **kwargs):
 
         return func(self, table_name, *args, **kwargs)
 
-    return wrapper
+    return wrapper
+
+
+def auth(blueprint, roles):
+    """
+    Decorator for authenticating and authorizing access to API routes.
+
+    Args:
+        blueprint (object): The Chalice Blueprint object, providing access to the current request.
+        roles (list[str]): The required role for authorization.
+
+    Returns:
+        function: A decorator function that authenticates and authorizes access based on the provided role.
+
+    Raises:
+        401 Unauthorized: If the Authorization header is missing, the token is invalid, or the token has expired.
+        403 Forbidden: If the decoded role is not part of the given role.
+    """
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            api_request = blueprint.current_request
+            auth_header = api_request.headers.get("Authorization", None)
+            if not auth_header:
+                raise UnauthorizedError("Authorization header is missing.")
+
+            _, token = auth_header.split(" ", 1) if " " in auth_header else (None, None)
+            if not token:
+                raise UnauthorizedError("Token is missing.")
+
+            try:
+                ssm_client = boto3.client("ssm")
+                auth_secret = ssm_client.get_parameter(
+                    Name="/Zap/AUTH_SECRET", WithDecryption=True
+                )["Parameter"]["Value"]
+                decoded = jwt.decode(token, auth_secret, algorithms=["HS256"])
+                print(roles)
+                # TODO: if decoded role is not part of given, reject auth
+
+                return func(*args, **kwargs)
+
+            except jwt.ExpiredSignatureError:
+                raise UnauthorizedError("Token has expired.")
+            except jwt.InvalidTokenError:
+                raise UnauthorizedError("Invalid token.")
+
+        return wrapper
+
+    return decorator
diff --git a/chalicelib/models/roles.py b/chalicelib/models/roles.py
@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class Roles(Enum):
+    ADMIN = "admin"
+    MEMBER = "member"
+    # Add more roles as needed
diff --git a/tests/api/test_applicants.py b/tests/api/test_applicants.py
@@ -15,44 +15,61 @@ def test_get_applicant():
     # Create a Chalice test client
     with Client(app) as client:
         # Mock applicant_service's get method
-        with patch(
-            "chalicelib.services.ApplicantService.applicant_service.get"
-        ) as mock_get:
-            mock_get.return_value = TEST_APPLICANTS[0]
-            response = client.http.get(
-                f"/applicant/{TEST_APPLICANTS[0]['applicantId']}"
-            )
+        with patch("chalicelib.decorators.jwt.decode") as mock_decode:
+            # Assuming the decoded token has the required role
+            mock_decode.return_value = {"role": "admin"}
 
-            # Check the response status code and body
-            assert response.status_code == 200
-            assert response.json_body == TEST_APPLICANTS[0]
+            with patch(
+                "chalicelib.services.ApplicantService.applicant_service.get"
+            ) as mock_get:
+                mock_get.return_value = TEST_APPLICANTS[0]
+                response = client.http.get(
+                    f"/applicant/{TEST_APPLICANTS[0]['applicantId']}",
+                    headers={"Authorization": "Bearer SAMPLE_TOKEN_STRING"},
+                )
+                # Check the response status code and body
+                assert response.status_code == 200
+                assert response.json_body == TEST_APPLICANTS[0]
 
 
 def test_get_all_applicants():
     # Create a Chalice test client
     with Client(app) as client:
         # Mock applicant_service's get method
-        with patch(
-            "chalicelib.services.ApplicantService.applicant_service.get_all"
-        ) as mock_get_all:
-            mock_get_all.return_value = TEST_APPLICANTS
-            response = client.http.get(f"/applicants")
+        with patch("chalicelib.decorators.jwt.decode") as mock_decode:
+            # Assuming the decoded token has the required role
+            mock_decode.return_value = {"role": "admin"}
+            with patch(
+                "chalicelib.services.ApplicantService.applicant_service.get_all"
+            ) as mock_get_all:
 
-            # Check the response status code and body
-            assert response.status_code == 200
-            assert response.json_body == TEST_APPLICANTS
+                mock_get_all.return_value = TEST_APPLICANTS
+                response = client.http.get(
+                    f"/applicants",
+                    headers={"Authorization": "Bearer SAMPLE_TOKEN_STRING"},
+                )
+
+                # Check the response status code and body
+                assert response.status_code == 200
+                assert response.json_body == TEST_APPLICANTS
 
 
 def test_get_all_applicants_from_listing():
     # Create a Chalice test client
     with Client(app) as client:
-        # Mock applicant_service's get method
-        with patch(
-            "chalicelib.services.ApplicantService.applicant_service.get_all_from_listing"
-        ) as mock_get_all_from_listing:
-            mock_get_all_from_listing.return_value = TEST_APPLICANTS
-            response = client.http.get(f"/applicants/test_listing_id")
-
-            # Check the response status code and body
-            assert response.status_code == 200
-            assert response.json_body == TEST_APPLICANTS
+
+        with patch("chalicelib.decorators.jwt.decode") as mock_decode:
+            # Assuming the decoded token has the required role
+            mock_decode.return_value = {"role": "admin"}
+            with patch(
+                "chalicelib.services.ApplicantService.applicant_service.get_all_from_listing"
+            ) as mock_get_all_from_listing:
+                mock_get_all_from_listing.return_value = TEST_APPLICANTS
+                response = client.http.get(
+                    f"/applicants/test_listing_id",
+                    headers={"Authorization": "Bearer SAMPLE_TOKEN_STRING"},
+                )
+
+                # Check the response status code and body
+                assert response.status_code == 200
+                assert response.json_body == TEST_APPLICANTS
diff --git a/tests/api/test_insights.py b/tests/api/test_insights.py
@@ -5,23 +5,32 @@
 
 from app import app
 
-with open('tests/fixtures/insights/general/sample_dashboard.json') as f:
+with open("tests/fixtures/insights/general/sample_dashboard.json") as f:
     SAMPLE_DASHBOARD = json.load(f)
 
-with open('tests/fixtures/insights/general/sample_distribution.json') as f:
+with open("tests/fixtures/insights/general/sample_distribution.json") as f:
     SAMPLE_DISTRIBUTION = json.load(f)
 
 
 def test_get_insights_from_listing():
     # Create a Chalice test client
     with Client(app) as client:
         # Mock applicant_service's get method
-        with patch(
-            "chalicelib.services.InsightsService.insights_service.get_insights_from_listing"
-        ) as mock_get_insights_from_listing:
-            mock_get_insights_from_listing.return_value = [SAMPLE_DASHBOARD, SAMPLE_DISTRIBUTION]
-            response = client.http.get(f"/insights/listing/test_listing_id")
+        with patch("chalicelib.decorators.jwt.decode") as mock_decode:
+            # Assuming the decoded token has the required role
+            mock_decode.return_value = {"role": "admin"}
+            with patch(
+                "chalicelib.services.InsightsService.insights_service.get_insights_from_listing",
+            ) as mock_get_insights_from_listing:
+                mock_get_insights_from_listing.return_value = [
+                    SAMPLE_DASHBOARD,
+                    SAMPLE_DISTRIBUTION,
+                ]
+                response = client.http.get(
+                    f"/insights/listing/test_listing_id",
+                    headers={"Authorization": "Bearer SAMPLE_TOKEN_STRING"},
+                )
 
-            # Check the response status code and body
-            assert response.status_code == 200
-            assert response.json_body == [SAMPLE_DASHBOARD, SAMPLE_DISTRIBUTION]
+                # Check the response status code and body
+                assert response.status_code == 200
+                assert response.json_body == [SAMPLE_DASHBOARD, SAMPLE_DISTRIBUTION]