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

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)