[Issue #2809] Handle parsing the jwt we created, and connect to a user (

#2959)

## Summary
Fixes #2809 

### Time to review: __15 mins__

## Changes proposed
Setup logic to process the jwt we created in
#2898

Setup a method to automatically generate a key for local development in
a secure way via an override file

## Context for reviewers
The core part of this PR is pretty straightforward, we parse the JWT, do
some validation, raise specific error messages for certain scenarios,
and have tests for that behavior.

For the auth token in the API request header, instead of using `Bearer
..` I left it as a dedicated header field. The bearer format doesn't let
you specify the header name and if we ever need multiple tokens
supported in an endpoint will lead to more headache.

---

Where things got complex was setting up the private/public key for the
API. These just need to be stored in env vars, but putting them directly
in our local.env file isn't ideal - even though the key will be
distinctly local-only, it will always be something flagged in security
scans and just generally look problematic.

To work around this fun problem, I realized I could solve another
annoyance at the same time, Docker as of January 2024 allows you to
specify multiple env files + make them optional. So - I used that. I
setup a script that creates an `override.env` file that you can freely
modify and won't be checked in, and more importantly, automatically
contains secrets like those public/private keys we didn't want to check
in. (Note - if you're wondering why I didn't use Docker secrets, they're
far more complex and this PR would've been 20+ files to make that
half-work).


## Additional information
Locally I confirmed we can set tokens in the swagger docs and they work
- we don't yet have an endpoint that uses this outside of the unit test
I setup, but I temporarily modified the healthcheck endpoint to validate
things work outside of tests as well.

<img width="641" alt="Screenshot 2024-11-20 at 4 06 48 PM"
src="https://github.com/user-attachments/assets/7f4b6d7e-0c2b-4a5b-a057-5695672ec31f">


The override file we generate looks like this (with the relevant key
info removed):
```
# override.env
#
# Any environment variables written to this file
# will take precedence over those defined in local.env
#
# This file will not be checked into github and it is safe
# to store secrets here, however you should still follow caution
# with using any secrets locally if they cause the app to interact
# with external systems.
#
# This file was generated by running:
#    make setup-env-overrides
#
# Which runs as part of our "make init" flow.
#
# If you would like to re-generate this file, please run:
#    make setup-env-overrides --recreate
#
# Note that this will completely erase any existing configuration you may have


############################
# Authentication
############################

API_JWT_PRIVATE_KEY="-----BEGIN RSA PRIVATE KEY-----
...
-----END RSA PRIVATE KEY-----"

API_JWT_PUBLIC_KEY="-----BEGIN PUBLIC KEY-----
...
-----END PUBLIC KEY-----"
```

---------

Co-authored-by: nava-platform-bot <[email protected]>

api/.gitignore

@@ -19,6 +19,7 @@ coverage.*
 # Environment variables
 .env
 .envrc
+override.env
 
 # mypy
 .mypy_cache
@@ -31,4 +32,9 @@ coverage.*
 /test-results/
 
 # localstack
-/volume
+/volume
+
+# All pem/pub/secret keys
+*.key
+*.pub
+*.pem

api/Makefile

@@ -84,6 +84,9 @@ setup-local:
 	# Install dependencies
 	poetry install --no-root --all-extras --with dev
 
+setup-env-override-file:
+	./bin/setup-env-override-file.sh $(args)
+
 ##################################################
 # API Build & Run
 ##################################################
@@ -100,7 +103,7 @@ start-debug:
 run-logs: start ## Start the API and follow the logs
 	docker compose logs --follow --no-color $(APP_NAME)
 
-init: build init-db init-opensearch init-localstack
+init: setup-env-override-file build init-db init-opensearch init-localstack
 
 clean-volumes: ## Remove project docker volumes - which includes the DB, and OpenSearch state
 	docker compose down --volumes

api/bin/setup-env-override-file.sh

@@ -0,0 +1,123 @@
+#!/usr/bin/env bash
+# setup-env-override-file.sh
+#
+# Generate an override.env file
+# with secrets pre-populated for local development.
+#
+# Examples:
+#   ./setup-env-override-file.sh
+#   ./setup-env-override-file.sh --recreate
+#
+
+set -o errexit -o pipefail
+
+PROGRAM_NAME=$(basename "$0")
+
+CYAN='\033[96m'
+GREEN='\033[92m'
+RED='\033[01;31m'
+END='\033[0m'
+
+USAGE="Usage: $PROGRAM_NAME [OPTION]
+
+  --recreate         Recreate the override.env file, fully overwriting any existing file
+"
+
+main() {
+  print_log "Running $PROGRAM_NAME"
+
+  for arg in "$@"
+  do
+    if [ "$arg" == "--recreate" ]; then
+      recreate=1
+    else
+      echo "$USAGE"
+      exit 1
+    fi
+  done
+
+  OVERRIDE_FILE="override.env"
+
+  if [ -f "$OVERRIDE_FILE" ] ; then
+    if [ $recreate ] ; then
+      print_log "Recreating existing override.env file"
+    else
+      print_log "override.env already exists, not recreating"
+      exit 0
+    fi
+  fi
+
+  # Delete any key files that may be leftover from a prior run
+  cleanup_files
+
+  # Generate RSA keys
+  # note ssh-keygen generates a different format for
+  # the public key so we run it through openssl to fix it
+  ssh-keygen -t rsa -b 2048 -m PEM -N '' -f tmp_jwk.key 2>&1 >/dev/null
+  openssl rsa -in tmp_jwk.key -pubout -outform PEM -out tmp_jwk.pub
+
+  PUBLIC_KEY=`cat tmp_jwk.pub`
+  PRIVATE_KEY=`cat tmp_jwk.key`
+
+  cat > $OVERRIDE_FILE <<EOF
+# override.env
+#
+# Any environment variables written to this file
+# will take precedence over those defined in local.env
+#
+# This file will not be checked into github and it is safe
+# to store secrets here, however you should still follow caution
+# with using any secrets locally if they cause the app to interact
+# with external systems.
+#
+# This file was generated by running:
+#    make setup-env-override-file
+#
+# Which runs as part of our "make init" flow.
+#
+# If you would like to re-generate this file, please run:
+#    make setup-env-override-file args="--recreate"
+#
+# Note that this will completely erase any existing configuration you may have
+
+############################
+# Authentication
+############################
+
+API_JWT_PRIVATE_KEY="$PRIVATE_KEY"
+
+API_JWT_PUBLIC_KEY="$PUBLIC_KEY"
+EOF
+
+
+  print_log "Created new override.env"
+
+  # Cleanup all keys generated in this run
+  cleanup_files
+}
+
+# Cleanup a single file if it exists
+cleanup_file()
+{
+  FILE=$1
+  shift;
+
+  if [ -f "$FILE" ] ; then
+    rm "$FILE"
+  fi
+}
+
+# Cleanup all miscellaneous keys generated
+cleanup_files()
+{
+  cleanup_file tmp_jwk.key
+  cleanup_file tmp_jwk.pub
+  cleanup_file tmp_jwk.key.pub
+}
+
+print_log() {
+  printf "$CYAN%s $GREEN%s: $END%s\\n" "$(date "+%Y-%m-%d %H:%M:%S")" "$PROGRAM_NAME" "$*"
+}
+
+# Entry point
+main "$@"

api/docker-compose.yml

@@ -90,7 +90,11 @@ services:
         "--reload",
       ]
     container_name: grants-api
-    env_file: ./local.env
+    env_file:
+      - path: ./local.env
+        required: true
+      - path: ./override.env
+        required: false
     ports:
       - 8080:8080
     volumes:

api/local.env

@@ -45,7 +45,13 @@ LOGIN_GOV_JWK_ENDPOINT=http://localhost:5001/issuer1/jwks
 LOGIN_GOV_ENDPOINT=http://localhost:5001
 LOGIN_GOV_CLIENT_ID=TODO
 
+# These should be set in your override.env file
+# which can be created by running `make setup-env-override-file`
+API_JWT_PRIVATE_KEY=
+API_JWT_PUBLIC_KEY=
+
 ENABLE_AUTH_ENDPOINT=TRUE
+
 ############################
 # DB Environment Variables
 ############################

api/openapi.generated.yml

@@ -1423,4 +1423,8 @@ components:
       type: apiKey
       in: header
       name: X-Auth
+    ApiJwtAuth:
+      type: apiKey
+      in: header
+      name: X-SGG-Token
 

api/src/app.py

@@ -21,7 +21,8 @@
 from src.api.schemas import response_schema
 from src.api.users.user_blueprint import user_blueprint
 from src.app_config import AppConfig
-from src.auth.api_key_auth import get_app_security_scheme
+from src.auth.api_jwt_auth import initialize_jwt_auth
+from src.auth.auth_utils import get_app_security_scheme
 from src.data_migration.data_migration_blueprint import data_migration_blueprint
 from src.search.backend.load_search_data_blueprint import load_search_data_blueprint
 from src.task import task_blueprint
@@ -58,6 +59,12 @@ def create_app() -> APIFlask:
     register_index(app)
     register_search_client(app)
 
+    # TODO - once we merge the auth changes for setting up the initial route
+    # will reuse the config from it, for now we'll do this a bit hacky
+    # This cannot be removed non-locally until we setup RSA keys for non-local envs
+    if os.getenv("ENVIRONMENT") == "local":
+        initialize_jwt_auth()
+
     return app
 
 

api/src/auth/api_jwt_auth.py

@@ -1,19 +1,27 @@
 import logging
 import uuid
 from datetime import timedelta
+from typing import Tuple
 
 import jwt
+from apiflask import HTTPTokenAuth
 from pydantic import Field
 from sqlalchemy import select
+from sqlalchemy.orm import selectinload
 
 import src.util.datetime_util as datetime_util
 from src.adapters import db
+from src.adapters.db import flask_db
+from src.api.route_utils import raise_flask_error
 from src.auth.auth_errors import JwtValidationError
 from src.db.models.user_models import User, UserTokenSession
+from src.logging.flask_logger import add_extra_data_to_current_request_logs
 from src.util.env_config import PydanticBaseEnvConfig
 
 logger = logging.getLogger(__name__)
 
+api_jwt_auth = HTTPTokenAuth("ApiKey", header="X-SGG-Token", security_scheme_name="ApiJwtAuth")
+
 
 class ApiJwtConfig(PydanticBaseEnvConfig):
 
@@ -32,7 +40,7 @@ class ApiJwtConfig(PydanticBaseEnvConfig):
 _config: ApiJwtConfig | None = None
 
 
-def initialize() -> None:
+def initialize_jwt_auth() -> None:
     global _config
     if not _config:
         _config = ApiJwtConfig()
@@ -53,14 +61,14 @@ def get_config() -> ApiJwtConfig:
     global _config
 
     if _config is None:
-        raise Exception("No JWT configuration - initialize() must be run first")
+        raise Exception("No JWT configuration - initialize_jwt_auth() must be run first")
 
     return _config
 
 
 def create_jwt_for_user(
     user: User, db_session: db.Session, config: ApiJwtConfig | None = None
-) -> str:
+) -> Tuple[str, UserTokenSession]:
     if config is None:
         config = get_config()
 
@@ -72,13 +80,8 @@ def create_jwt_for_user(
     expiration_time = current_time + timedelta(minutes=config.token_expiration_minutes)
 
     # Create the session in the DB
-    db_session.add(
-        UserTokenSession(
-            user=user,
-            token_id=token_id,
-            expires_at=expiration_time,
-        )
-    )
+    user_token_session = UserTokenSession(user=user, token_id=token_id, expires_at=expiration_time)
+    db_session.add(user_token_session)
 
     # Create the JWT with information we'll want to receive back
     payload = {
@@ -89,13 +92,21 @@ def create_jwt_for_user(
         "iss": config.issuer,
     }
 
-    return jwt.encode(payload, config.private_key, algorithm="RS256")
+    logger.info(
+        "Created JWT token",
+        extra={
+            "auth.user_id": str(user_token_session.user_id),
+            "auth.token_id": str(user_token_session.token_id),
+        },
+    )
+
+    return jwt.encode(payload, config.private_key, algorithm="RS256"), user_token_session
 
 
 def parse_jwt_for_user(
     token: str, db_session: db.Session, config: ApiJwtConfig | None = None
-) -> User:
-    # TODO - more implementation/validation to come in https://github.com/HHS/simpler-grants-gov/issues/2809
+) -> UserTokenSession:
+    """Handle processing a jwt token, and connecting it to a user token session in our DB"""
     if config is None:
         config = get_config()
 
@@ -135,8 +146,10 @@ def parse_jwt_for_user(
         raise JwtValidationError("Token missing sub field")
 
     token_session: UserTokenSession | None = db_session.execute(
-        select(UserTokenSession).join(User).where(UserTokenSession.token_id == sub_id)
-    ).scalar_one_or_none()
+        select(UserTokenSession)
+        .where(UserTokenSession.token_id == sub_id)
+        .options(selectinload("*"))
+    ).scalar()
 
     # We check both the token expires_at timestamp as well as an
     # is_valid flag to make sure the token is still valid.
@@ -147,4 +160,48 @@ def parse_jwt_for_user(
     if token_session.is_valid is False:
         raise JwtValidationError("Token is no longer valid")
 
-    return token_session.user
+    return token_session
+
+
+@api_jwt_auth.verify_token
+@flask_db.with_db_session()
+def decode_token(db_session: db.Session, token: str) -> UserTokenSession:
+    """
+    Process an internal jwt token as created by the above create_jwt_for_user method.
+
+    To add this auth to an endpoint, simply put::
+
+        from src.auth.api_jwt_auth import api_jwt_auth
+
+        @example_blueprint.get("/example")
+        @example_blueprint.auth_required(api_jwt_auth)
+        @flask_db.with_db_session()
+        def example_method(db_session: db.Session) -> response.ApiResponse:
+            # The token session object can be fetched from the auth object
+            token_session: UserTokenSession = api_jwt_auth.current_user
+
+            # If you want to modify the token_session or user, you will
+            # need to add it to the DB session otherwise it won't do anything
+            db_session.add(token_session)
+            token_session.expires_at = ...
+            ...
+    """
+
+    try:
+        user_token_session = parse_jwt_for_user(token, db_session)
+
+        add_extra_data_to_current_request_logs(
+            {
+                "auth.user_id": str(user_token_session.user_id),
+                "auth.token_id": str(user_token_session.token_id),
+            }
+        )
+        logger.info("JWT Authentication Successful")
+
+        # Return the user token session object
+        return user_token_session
+    except JwtValidationError as e:
+        # If validation of the jwt fails, pass the error message back to the user
+        # The message is just the value we set when constructing the JwtValidationError
+        logger.info("JWT Authentication Failed for provided token", extra={"auth.issue": e.message})
+        raise_flask_error(401, e.message)