26556 - Add in user / membership row on API key creation (#3327)

Author: seeker25

auth-api/src/auth_api/models/membership.py

@@ -24,7 +24,7 @@
 from sqlalchemy import Column, ForeignKey, Integer, and_, desc, func
 from sqlalchemy.orm import relationship
 
-from auth_api.utils.enums import OrgType, Status
+from auth_api.utils.enums import LoginSource, OrgType, Status
 from auth_api.utils.roles import ADMIN, COORDINATOR, USER, VALID_ORG_STATUSES, VALID_STATUSES
 
 from .base_model import BaseModel
@@ -101,6 +101,7 @@ def find_members_by_org_id_by_status_by_roles(
             .filter(and_(Membership.status == status, Membership.membership_type_code.in_(roles)))
             .join(OrgModel)
             .filter(OrgModel.id == int(org_id or -1))
+            .filter(~Membership.user.has(login_source=LoginSource.API_GW.value))
             .all()
         )
 

auth-api/src/auth_api/models/user.py

@@ -144,6 +144,25 @@ def create_from_jwt_token(cls, first_name: str, last_name: str, **kwargs):
             return user
         return None
 
+    @staticmethod
+    def create_user_for_api_user(username, keycloak_guid):
+        """Create user for API user."""
+        api_user = User(
+            username=username,
+            firstname=None,
+            lastname=username,
+            email=None,
+            keycloak_guid=keycloak_guid,
+            created=datetime.datetime.now(tz=datetime.timezone.utc),
+            login_source=LoginSource.API_GW.value,
+            status=UserStatusCode.get_default_type(),
+            idp_userid=username,
+            login_time=datetime.datetime.now(tz=datetime.timezone.utc),
+            type=Role.PUBLIC_USER.name,
+            verified=True,
+        ).save()
+        return api_user
+
     @classmethod
     @user_context
     def update_from_jwt_token(

auth-api/src/auth_api/services/api_gateway.py

@@ -21,12 +21,17 @@
 from structured_logging import StructuredLogging
 
 from auth_api.exceptions import BusinessException, Error
+from auth_api.models.membership import Membership as MembershipModel
 from auth_api.models.org import Org as OrgModel
+from auth_api.models.user import User as UserModel
 from auth_api.services.authorization import check_auth
+from auth_api.services.flags import flags
 from auth_api.services.keycloak import KeycloakService
+from auth_api.services.membership import Membership as MembershipService
 from auth_api.services.rest_service import RestService
 from auth_api.utils.api_gateway import generate_client_representation
 from auth_api.utils.constants import GROUP_ACCOUNT_HOLDERS, GROUP_API_GW_SANDBOX_USERS, GROUP_API_GW_USERS
+from auth_api.utils.enums import Status
 from auth_api.utils.roles import ADMIN, STAFF
 from auth_api.utils.user_context import UserContext, user_context
 
@@ -60,6 +65,8 @@ def create_key(cls, org_id: int, request_json: Dict[str, str]):
             # If env is sandbox; then create a sandbox payment account.
             if env != "prod":
                 cls._create_payment_account(org)
+                # Future - if PROD and target is SANDBOX - Call into AUTH-API to create an org, this will call PAY-API
+                # to create payment account
             cls._create_consumer(name, org, env=env)
             org.has_api_access = True
             org.save()
@@ -74,20 +81,41 @@ def create_key(cls, org_id: int, request_json: Dict[str, str]):
             )
             response = api_key_response.json()
 
+        cls._create_user_and_membership_for_api_user(org_id, env)
         return response
 
+    @classmethod
+    def _create_user_and_membership_for_api_user(cls, org_id: int, env: str):
+        """Create a user and membership for the api user."""
+        if flags.is_on("enable-api-gw-user-membership-creation", True) is True:
+            client_name = ApiGateway.get_api_client_id(org_id, env)
+            client = KeycloakService.get_service_account_by_client_name(client_name)
+            if (api_user := UserModel.find_by_username(client_name)) is None:
+                api_user = UserModel.create_user_for_api_user(client_name, client.get("id"))
+            if MembershipModel.find_membership_by_user_and_org(api_user.id, org_id) is None:
+                MembershipService.create_admin_membership_for_api_user(org_id, api_user.id)
+
     @classmethod
     def _get_api_gw_key(cls, env):
+        """Get the api gateway key."""
         logger.info("_get_api_gw_key %s", env)
         return current_app.config.get("API_GW_KEY") if env == "prod" else current_app.config.get("API_GW_NON_PROD_KEY")
 
+    @staticmethod
+    def get_api_client_id(org_id, env):
+        """Get the client id for the org."""
+        client_id_pattern = current_app.config.get("API_GW_KC_CLIENT_ID_PATTERN")
+        suffix = "-sandbox" if env != "prod" else ""
+        client_id = f"{client_id_pattern}{suffix}".format(account_id=org_id)
+        return client_id
+
     @classmethod
     def _create_consumer(cls, name, org, env):
         """Create an API Gateway consumer."""
         consumer_endpoint: str = cls._get_api_consumer_endpoint(env)
         gw_api_key = cls._get_api_gw_key(env)
         email = cls._get_email_id(org.id, env)
-        client_rep = generate_client_representation(org.id, current_app.config.get("API_GW_KC_CLIENT_ID_PATTERN"), env)
+        client_rep = generate_client_representation(org.id, ApiGateway.get_api_client_id(org.id, env))
         KeycloakService.create_client(client_rep)
         service_account = KeycloakService.get_service_account_user(client_rep.get("id"))
 
@@ -184,6 +212,7 @@ def _add_key_to_response(_key):
 
     @classmethod
     def _get_email_id(cls, org_id, env) -> str:
+        """Get the email id for the org."""
         if current_app.config.get("API_GW_CONSUMER_EMAIL", None) is not None:
             return current_app.config.get("API_GW_CONSUMER_EMAIL")
 
@@ -243,6 +272,7 @@ def _create_payment_account(cls, org: OrgModel, **kwargs):
 
     @classmethod
     def _create_sandbox_pay_account(cls, pay_request, user):
+        """Create a sandbox payment account."""
         logger.info("Creating Sandbox Payload %s", pay_request)
         pay_sandbox_accounts_endpoint = f"{current_app.config.get('PAY_API_SANDBOX_URL')}/accounts?sandbox=true"
         RestService.post(
@@ -251,12 +281,14 @@ def _create_sandbox_pay_account(cls, pay_request, user):
 
     @classmethod
     def _get_pay_account(cls, org, user):
+        """Get the payment account for the org."""
         pay_accounts_endpoint = f"{current_app.config.get('PAY_API_URL')}/accounts/{org.id}"
         pay_account = RestService.get(endpoint=pay_accounts_endpoint, token=user.bearer_token).json()
         return pay_account
 
     @classmethod
     def _get_api_consumer_endpoint(cls, env):
+        """Get the consumer endpoint for the environment."""
         logger.info("_get_api_consumer_endpoint %s", env)
         return (
             current_app.config.get("API_GW_CONSUMERS_API_URL")

auth-api/src/auth_api/services/keycloak.py

@@ -446,6 +446,24 @@ def create_client(client_representation: Dict[str, any]):
         )
         response.raise_for_status()
 
+    @staticmethod
+    def get_service_account_by_client_name(client_name: str):
+        """Get client by name."""
+        config = current_app.config
+        base_url = config.get("KEYCLOAK_BASE_URL")
+        realm = config.get("KEYCLOAK_REALMNAME")
+        timeout = config.get("CONNECT_TIMEOUT", 60)
+        admin_token = KeycloakService._get_admin_token()
+        headers = {"Content-Type": ContentType.JSON.value, "Authorization": f"Bearer {admin_token}"}
+        response = requests.get(
+            f"{base_url}/auth/admin/realms/{realm}/clients?clientId={client_name}",
+            headers=headers,
+            timeout=timeout,
+        )
+        response.raise_for_status()
+        client_id = response.json()[0]["id"]
+        return KeycloakService.get_service_account_user(client_id)
+
     @staticmethod
     def get_service_account_user(client_identifier: str):
         """Return service account user."""
@@ -454,7 +472,6 @@ def get_service_account_user(client_identifier: str):
         realm = config.get("KEYCLOAK_REALMNAME")
         timeout = config.get("CONNECT_TIMEOUT", 60)
         admin_token = KeycloakService._get_admin_token()
-
         headers = {"Content-Type": ContentType.JSON.value, "Authorization": f"Bearer {admin_token}"}
         response = requests.get(
             f"{base_url}/auth/admin/realms/{realm}/clients/{client_identifier}/service-account-user",

auth-api/src/auth_api/services/membership.py

@@ -18,6 +18,7 @@
 
 import json
 
+from flask import current_app
 from jinja2 import Environment, FileSystemLoader
 from sbc_common_components.utils.enums import QueueMessageTypes
 from structured_logging import StructuredLogging
@@ -338,6 +339,7 @@ def _add_or_remove_group(model: MembershipModel):
 
     @staticmethod
     def add_or_remove_group_for_staff(model: MembershipModel):
+        """Add or remove the user from/to various staff keycloak groups."""
         mapping_group = org_type_to_group_mapping.get(model.org.type_code)
         if not mapping_group:
             return
@@ -383,3 +385,11 @@ def add_staff_membership(user_id):
     def remove_staff_membership(user_id):
         """Remove staff membership for the specified user."""
         MembershipModel.remove_membership_for_staff(user_id)
+
+    @staticmethod
+    def create_admin_membership_for_api_user(org_id, user_id):
+        """Create a membership for an api user."""
+        current_app.logger.info(f"Creating membership in {org_id} for API user {user_id}")
+        return MembershipModel(
+            org_id=org_id, user_id=user_id, membership_type_code=ADMIN, status=Status.ACTIVE.value
+        ).save()

auth-api/src/auth_api/services/user.py

@@ -17,9 +17,11 @@
 """
 
 import json
+from datetime import datetime, timezone
 from http import HTTPStatus
 from typing import Dict, List
 
+from flask import current_app
 from jinja2 import Environment, FileSystemLoader
 from requests import HTTPError
 from sbc_common_components.utils.enums import QueueMessageTypes
@@ -34,6 +36,7 @@
 from auth_api.models import User as UserModel
 from auth_api.models import db
 from auth_api.models.dataclass import Activity
+from auth_api.models.user import UserStatusCode
 from auth_api.schemas import UserSchema
 from auth_api.services.authorization import check_auth
 from auth_api.services.keycloak_user import KeycloakUser

auth-api/src/auth_api/utils/api_gateway.py

@@ -16,17 +16,13 @@
 import uuid
 
 
-def generate_client_representation(account_id: int, client_id_pattern: str, env: str) -> dict:
+def generate_client_representation(account_id: int, client_id: str) -> dict:
     """Return dictionary for api gateway client user."""
     _id = str(uuid.uuid4())
     _secret = secrets.token_urlsafe(36)
-    if env != "prod":
-        client_id_pattern += "-sandbox"
-    _client_id = client_id_pattern.format(account_id=account_id)
-
     client_json: dict = {
         "id": _id,
-        "clientId": _client_id,
+        "clientId": client_id,
         "rootUrl": "",
         "adminUrl": "",
         "baseUrl": "",
@@ -86,7 +82,7 @@ def generate_client_representation(account_id: int, client_id_pattern: str, env:
                 "protocolMapper": "oidc-hardcoded-claim-mapper",
                 "consentRequired": False,
                 "config": {
-                    "claim.value": _client_id,
+                    "claim.value": client_id,
                     "userinfo.token.claim": "true",
                     "id.token.claim": "true",
                     "access.token.claim": "true",
@@ -114,7 +110,7 @@ def generate_client_representation(account_id: int, client_id_pattern: str, env:
                 "protocolMapper": "oidc-hardcoded-claim-mapper",
                 "consentRequired": False,
                 "config": {
-                    "claim.value": _client_id,
+                    "claim.value": client_id,
                     "userinfo.token.claim": "true",
                     "id.token.claim": "true",
                     "access.token.claim": "true",
@@ -128,7 +124,7 @@ def generate_client_representation(account_id: int, client_id_pattern: str, env:
                 "protocolMapper": "oidc-hardcoded-claim-mapper",
                 "consentRequired": False,
                 "config": {
-                    "claim.value": _client_id,
+                    "claim.value": client_id,
                     "userinfo.token.claim": "true",
                     "id.token.claim": "true",
                     "access.token.claim": "true",
@@ -142,7 +138,7 @@ def generate_client_representation(account_id: int, client_id_pattern: str, env:
                 "protocolMapper": "oidc-hardcoded-claim-mapper",
                 "consentRequired": False,
                 "config": {
-                    "claim.value": _client_id,
+                    "claim.value": client_id,
                     "userinfo.token.claim": "true",
                     "id.token.claim": "true",
                     "access.token.claim": "true",
@@ -156,7 +152,7 @@ def generate_client_representation(account_id: int, client_id_pattern: str, env:
                 "protocolMapper": "oidc-hardcoded-claim-mapper",
                 "consentRequired": False,
                 "config": {
-                    "claim.value": _client_id,
+                    "claim.value": client_id,
                     "userinfo.token.claim": "true",
                     "id.token.claim": "true",
                     "access.token.claim": "true",
@@ -215,7 +211,7 @@ def generate_client_representation(account_id: int, client_id_pattern: str, env:
                 "protocolMapper": "oidc-hardcoded-claim-mapper",
                 "consentRequired": False,
                 "config": {
-                    "claim.value": _client_id,
+                    "claim.value": client_id,
                     "userinfo.token.claim": "true",
                     "id.token.claim": "true",
                     "access.token.claim": "true",

auth-api/tests/docker/docker-compose.yml

@@ -2,13 +2,15 @@ version: "3"
 
 services:
   keycloak:
-    image: quay.io/keycloak/keycloak:12.0.2
+    image: quay.io/keycloak/keycloak:26.0
     ports:
       - "8081:8081"
     environment:
-      - KEYCLOAK_USER=admin
-      - KEYCLOAK_PASSWORD=admin
-    command: -b 0.0.0.0 -Djboss.http.port=8081 -Dkeycloak.migration.action=import -Dkeycloak.migration.provider=dir -Dkeycloak.migration.dir=/tmp/keycloak/test -Dkeycloak.migration.strategy=OVERWRITE_EXISTING
+      - KEYCLOAK_ADMIN=admin
+      - KEYCLOAK_ADMIN_PASSWORD=admin
+    volumes:
+      - ./setup:/opt/keycloak/data/import
+    command:  ["start-dev", "--import-realm",  "--http-port=8081", "--http-relative-path=/auth"]
     healthcheck:
       test:
         [
@@ -20,8 +22,6 @@ services:
       interval: 30s
       timeout: 10s
       retries: 10
-    volumes:
-      - ./setup:/tmp/keycloak/test/
   nats:
     image: nats-streaming
     restart: always

auth-api/tests/unit/api/test_org.py

@@ -43,6 +43,7 @@
     AccessType,
     AffidavitStatus,
     CorpType,
+    LoginSource,
     NRActionCodes,
     NRStatus,
     OrgStatus,
@@ -66,6 +67,7 @@
     TestJwtClaims,
     TestOrgInfo,
     TestPaymentMethodInfo,
+    TestUserInfo,
 )
 from tests.utilities.factory_utils import (
     convert_org_to_staff_org,
@@ -1276,6 +1278,11 @@ def test_get_members(client, jwt, session, keycloak_mock):  # pylint:disable=unu
     )
     dictionary = json.loads(rv.data)
     org_id = dictionary["id"]
+    # Create API_GW user, this shouldn't show up in the users list
+    user_dict = dict(TestUserInfo.user1)
+    user_dict["login_source"] = LoginSource.API_GW.value
+    api_user = factory_user_model(user_dict)
+    factory_membership_model(api_user.id, org_id=org_id)
 
     rv = client.get("/api/v1/orgs/{}/members".format(org_id), headers=headers, content_type="application/json")
 
@@ -2936,7 +2943,7 @@ def test_update_org_api_access(client, jwt, session, keycloak_mock):  # pylint:d
 
 def test_search_org_members(client, jwt, session, keycloak_mock):  # pylint:disable=unused-argument
     """Assert that a list of members for an org search can be retrieved."""
-    user_info = TestJwtClaims.public_user_role
+    user_info = dict(TestJwtClaims.public_user_role)
     headers = factory_auth_header(jwt=jwt, claims=user_info)
     client.post("/api/v1/users", headers=headers, content_type="application/json")
     client.post("/api/v1/orgs", data=json.dumps(TestOrgInfo.org1), headers=headers, content_type="application/json")

auth-api/tests/unit/api/test_org_api_keys.py

@@ -18,8 +18,12 @@
 """
 
 import json
+from contextlib import suppress
 from http import HTTPStatus
 
+from auth_api.services.api_gateway import ApiGateway
+from auth_api.services.keycloak import KeycloakService
+from auth_api.utils.api_gateway import generate_client_representation
 from tests.utilities.factory_scenarios import TestJwtClaims, TestOrgInfo
 from tests.utilities.factory_utils import factory_auth_header