Finalised Token Service fixes and improvement

lib/workload/stateful/token_service/README.md

@@ -2,7 +2,9 @@
 
 This service provides the JWT token for API authentication and authorization (AAI) purpose. We use the Cognito as AAI service broker and, it is set up at our infrastructure repo. This service maintains 2 secrets with rotation enabled.
 
-## JWT
+## User Guide
+
+### JWT
 
 For most cases, you would want to lookup JWT token from the secret manager at the following coordinate.
 ```
@@ -17,7 +19,7 @@ import boto3
 
 sm_client = boto3.client('secretsmanager')
 
-resp = sm_client().get_secret_value(SecretId='orcabus/token-service-jwt')
+resp = sm_client.get_secret_value(SecretId='orcabus/token-service-jwt')
 
 jwt_json = resp['SecretString']
 jwt_dict = json.loads(jwt_json)
@@ -34,7 +36,7 @@ from libumccr.aws import libsm
 tok = json.loads(libsm.get_secret('orcabus/token-service-jwt'))['id_token']
 ```
 
-## Service User
+### Service User
 
 As an admin, you must register the service user. This has to be done at Cognito AAI terraform stack. Please follow `AdminCreateUser` [flow noted](https://github.com/umccr/infrastructure/pull/412/files) in `users.tf` at upstream infrastructure repo.
 
@@ -55,9 +57,7 @@ aws secretsmanager put-secret-value \
 
 After then, the scheduled secret rotation should carry on rotating password, every set days.
 
----
-
-## Stack
+## Development
 
 ### TL;DR
 
@@ -77,7 +77,7 @@ The stack contains 2 Lambda Python code that do secret rotation. This code is de
 ### Cognitor
 And, there is the thin service layer package called `cognitor` for interfacing with AWS Cognito through boto3 - in fact it just [a façade](https://www.google.com/search?q=fa%C3%A7ade+pattern) of boto3 for Cognito. See its test cases for how to use and operate it.
 
-### Local Dev
+### Local DX
 
 #### App
 
@@ -136,7 +136,7 @@ Then, do CloudFormation lint check:
 cfn-lint .local/template.yml
 ```
 
-If that all good, then you may diff e & push straight to dev for giving it the WIP a try...
+If that all good, then you may `diff -e` & `deploy -e` straight to dev for giving it the WIP a try...
 
 ```
 export AWS_PROFILE=umccr-dev-admin

lib/workload/stateful/token_service/deploy/construct/policy/index.ts

@@ -12,7 +12,6 @@ export function getLambdaBasicExecPolicy(resources: string[]) {
    * https://docs.aws.amazon.com/aws-managed-policy/latest/reference/AWSLambdaBasicExecutionRole.html
    */
   return new PolicyStatement({
-    sid: 'LambdaBasicExecStmt1711498867457',
     effect: Effect.ALLOW,
     actions: ['logs:CreateLogGroup', 'logs:CreateLogStream', 'logs:PutLogEvents'],
     resources: resources,
@@ -31,7 +30,6 @@ export function getLambdaVPCPolicy() {
    * https://docs.aws.amazon.com/aws-managed-policy/latest/reference/AWSLambdaVPCAccessExecutionRole.html
    */
   return new PolicyStatement({
-    sid: 'LambdaVPCStmt1711498867457',
     effect: Effect.ALLOW,
     actions: [
       'ec2:CreateNetworkInterface',
@@ -51,7 +49,6 @@ export function getSSMPolicy(resources: string[]) {
    * https://docs.aws.amazon.com/aws-managed-policy/latest/reference/AmazonSSMReadOnlyAccess.html
    */
   return new PolicyStatement({
-    sid: 'SSMStmt1711498867457',
     effect: Effect.ALLOW,
     actions: ['ssm:Describe*', 'ssm:Get*', 'ssm:List*'],
     resources: resources,
@@ -70,8 +67,6 @@ export const getCognitoAdminActions = () => {
     'cognito-idp:DescribeUserPool',
     'cognito-idp:AdminGetUser',
     'cognito-idp:AdminSetUserPassword',
-    'cognito-idp:InitiateAuth',
-    'cognito-idp:ListUserPools',
     'cognito-idp:ListUsers',
   ];
 };
@@ -81,29 +76,28 @@ export function getCognitoAdminPolicy(resources: string[]) {
    * The Cognito policy that specifically required for Token Service `cognitor.py` application
    */
   return new PolicyStatement({
-    sid: 'CognitoAdminStmt1711498867457',
     effect: Effect.ALLOW,
     actions: getCognitoAdminActions(),
     resources: resources,
   });
 }
 
-export const getCognitoJWTActions = () => {
-  /**
-   * Always return new string array of permission flags that is allowed.
-   */
-  return ['cognito-idp:InitiateAuth'];
-};
-
-export function getCognitoJWTPolicy(resources: string[]) {
+export function getCognitoJWTPolicy() {
   /**
+   * NOTE: This function only tailors to policy statement for `InitiateAuth` API call. This API
+   * endpoint is `public` call by design. See `Unauthenticated user operations` section, below.
+   * https://docs.aws.amazon.com/cognito/latest/developerguide/user-pools-API-operations.html
+   *
+   * Therefore, the resource-level permission filtering is not needed/supported. This is try with
+   * Policy Simulator. See thread in `#orcabus` channel for howto.
+   * https://umccr.slack.com/archives/C03ABJTSN7J/p1711576104287009
+   *
    * The Cognito policy that specifically required for Token Service `cognitor.py` application
    */
   return new PolicyStatement({
-    sid: 'CognitoJWTStmt1711498867457',
     effect: Effect.ALLOW,
-    actions: getCognitoJWTActions(),
-    resources: resources,
+    actions: ['cognito-idp:InitiateAuth'],
+    resources: ['*'], // see docstring ^^
   });
 }
 

lib/workload/stateful/token_service/deploy/stack.ts

@@ -11,7 +11,7 @@ import { LogGroup } from 'aws-cdk-lib/aws-logs';
 import { IVpc, Vpc } from 'aws-cdk-lib/aws-ec2';
 import {
   getCognitoAdminActions,
-  getCognitoJWTActions,
+  getCognitoJWTPolicy,
   getLambdaVPCPolicy,
   getServiceUserSecretResourcePolicy,
 } from './construct/policy';
@@ -127,8 +127,7 @@ export class TokenServiceStack extends Stack {
     });
     lambdaLogGroup.grantWrite(jwtRotationFn);
 
-    this.userPool.grant(jwtRotationFn, ...getCognitoJWTActions());
-
+    jwtRotationFn.addToRolePolicy(getCognitoJWTPolicy());
     jwtRotationFn.addToRolePolicy(getLambdaVPCPolicy());
 
     return jwtRotationFn;

lib/workload/stateful/token_service/token_service/cognitor/__init__.py

@@ -22,8 +22,14 @@ def __init__(self, user_pool_id: str, user_pool_app_client_id: str):
 
     @staticmethod
     def generate_password(length: int = 32) -> str:
+        """
+        Must meet the Cognito password requirements policy
+        https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-settings-policies.html
+        """
+        if length < 8:
+            raise ValueError('Length must be at least 8 characters or more better')
         return ''.join(
-            random.SystemRandom().choice(string.ascii_letters + string.digits + '!@#$%&.') for _ in range(length)
+            random.SystemRandom().choice(string.ascii_letters + string.digits + '_-!@#%&.') for _ in range(length)
         )
 
     def list_users(self, **kwargs) -> dict:
@@ -135,10 +141,7 @@ def rotate_service_user_password(self, user_dto: ServiceUserDto):
         """
         Forced set Cognito Service User login credential info from pass-in dto
         """
-        if not self.username_exists(user_dto.username):
-            return
-
-        _ = self.client.admin_set_user_password(
+        self.client.admin_set_user_password(
             UserPoolId=self.user_pool_id,
             Username=user_dto.username,
             Password=user_dto.password,

lib/workload/stateful/token_service/token_service/cognitor/tests.py

@@ -194,7 +194,6 @@ def test_rotate_service_user_password(self):
         python -m unittest token_service.cognitor.tests.CognitorUnitTest.test_rotate_service_user_password
         """
         with Stubber(self.mock_client) as stubber:
-            stubber.add_response('admin_get_user', {'Username': self.jjb_dto.username, })
             stubber.add_response('admin_set_user_password', {})
 
             self.srv.rotate_service_user_password(user_dto=self.jjb_dto)

lib/workload/stateful/token_service/token_service/rotate_service_jwt.py

@@ -21,7 +21,6 @@
 import boto3
 
 from cognitor import CognitoTokenService, ServiceUserDto
-from helper import get_secret_dict
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -123,15 +122,16 @@ def create_secret(service_client, arn, token):
 
     """
     # Make sure the current secret exists
-    current_dict = get_secret_dict(service_client, arn, "AWSCURRENT")
+    current_dict = _get_secret_dict(service_client, arn, "AWSCURRENT")
 
     # Now try to get the secret version, if that fails, put a new secret
     try:
-        get_secret_dict(service_client, arn, "AWSPENDING", token)
+        _get_secret_dict(service_client, arn, "AWSPENDING", token)
         logger.info("createSecret: Successfully retrieved secret for %s." % arn)
     except service_client.exceptions.ResourceNotFoundException:
         # Get Cognito Service User login credential info from another (peer) rotating secret from Secret Manager
-        service_user_info = service_client.get_secret_value(SecretId=service_user_secret_id)
+        resp = service_client.get_secret_value(SecretId=service_user_secret_id)
+        service_user_info = json.loads(resp['SecretString'])
         # Generate new token object
         current_dict = token_srv.generate_service_user_tokens(
             user_dto=ServiceUserDto(
@@ -168,11 +168,11 @@ def set_secret(service_client, arn, token):
 
     """
     try:
-        previous_dict = get_secret_dict(service_client, arn, "AWSPREVIOUS")
+        previous_dict = _get_secret_dict(service_client, arn, "AWSPREVIOUS")
     except (service_client.exceptions.ResourceNotFoundException, KeyError):
         previous_dict = None
-    current_dict = get_secret_dict(service_client, arn, "AWSCURRENT")
-    pending_dict = get_secret_dict(service_client, arn, "AWSPENDING", token)
+    current_dict = _get_secret_dict(service_client, arn, "AWSCURRENT")
+    pending_dict = _get_secret_dict(service_client, arn, "AWSPENDING", token)
 
     # First try to login with the pending secret, if it succeeds, return
     id_token = pending_dict['id_token']
@@ -219,7 +219,7 @@ def test_secret(service_client, arn, token):
 
     """
     # Try to login with the pending secret, if it succeeds, return
-    pending_dict = get_secret_dict(service_client, arn, "AWSPENDING", token)
+    pending_dict = _get_secret_dict(service_client, arn, "AWSPENDING", token)
     id_token = pending_dict['id_token']
     if _is_valid_jwt(id_token):
         # Being able to generate tokens using pending username/password consider success.
@@ -263,6 +263,49 @@ def finish_secret(service_client, arn, token):
 # --- module internal functions
 
 
+def _get_secret_dict(service_client, arn, stage, token=None):
+    """Gets the secret dictionary corresponding for the secret arn, stage, and token
+
+    This helper function gets credentials for the arn and stage passed in and returns the dictionary by parsing the
+    JSON string
+
+    Args:
+        service_client (client): The secrets manager service client
+
+        arn (string): The secret ARN or other identifier
+
+        token (string): The ClientRequestToken associated with the secret version, or None if no validation is desired
+
+        stage (string): The stage identifying the secret version
+
+    Returns:
+        SecretDictionary: Secret dictionary
+
+    Raises:
+        ResourceNotFoundException: If the secret with the specified arn and stage does not exist
+
+        ValueError: If the secret is not valid JSON
+
+    """
+    required_fields = ['id_token']
+
+    # Only do VersionId validation against the stage if a token is passed in
+    if token:
+        secret = service_client.get_secret_value(SecretId=arn, VersionId=token, VersionStage=stage)
+    else:
+        secret = service_client.get_secret_value(SecretId=arn, VersionStage=stage)
+    plaintext = secret['SecretString']
+    secret_dict = json.loads(plaintext)
+
+    # Run validations against the secret
+    for field in required_fields:
+        if field not in secret_dict:
+            raise KeyError("%s key is missing from secret JSON" % field)
+
+    # Parse and return the secret JSON string
+    return secret_dict
+
+
 def _is_valid_jwt(this_id_token: str) -> bool:
     """
     The following should be a good self-contained JWT check for now.