Improve functions

.github/workflows/provisioning.yml

@@ -26,6 +26,11 @@ jobs:
       pull-requests: write
 
     steps:
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'
+
       - name: Checkout Repository
         uses: actions/checkout@v4
 
@@ -62,5 +67,5 @@ jobs:
         run: exit 1
 
       - name: Terraform Apply
-        if: github.ref == 'refs/heads/main' && github.event_name == 'push'
+        #if: github.ref == 'refs/heads/main' && github.event_name == 'push'
         run: terraform apply -auto-approve -input=false

README.md

@@ -1,12 +1,12 @@
-# IaC for Lambda functions
+# IaC for Auth
 
-IaC para funções Lambda provisionada com Terraform.
+IaC para usar pool do Cognito e funções Lambda de autenticação provisionada com Terraform.
 
 Repositório principal: [tech-challenge](https://github.com/FIAP-3SOAT-G15/tech-challenge)
 
 ## Recursos criados
 
-Funções Lambda para sign up e sign in de clientes, e auth challenge, para serem usados com o Cognito.
+User pool do Cognito com grupos para clientes e administradores, funções Lambda para sign up e sign in de clientes, e função que define auth challenge como trigger de autenticação customizada no Cognito (com CPF ou e-mail e senha).
 
 ## Estrutura
 

src/sign-in/lambda_function.py

@@ -1,15 +1,17 @@
 import json
+import logging
 import os
+
 import boto3
+import botocore.exceptions
 
 cognito = boto3.client('cognito-idp')
 
-USER_POOL_ID = os.getenv('USER_POOL_ID')
-CLIENT_ID = os.getenv('CLIENT_ID')
+USER_POOL_ID = os.environ.get('USER_POOL_ID')
+CLIENT_ID = os.environ.get('CLIENT_ID')
 
 
 def lambda_handler(event, context):
-    print(event)
     body = json.loads(event['body'])
 
     identifier = body.get('cpf') or body.get('email')
@@ -33,10 +35,27 @@ def lambda_handler(event, context):
             'headers': {'Content-Type': 'application/json'},
             'body': json.dumps(response),
         }
-    except cognito.exceptions.ClientError as e:
-        print(e)
-        return {
-            'statusCode': 500,
-            'headers': {'Content-Type': 'application/json'},
-            'body': "{ 'message': 'Error initiating authentication' }",
-        }
+
+    except botocore.exceptions.ClientError as error:
+        logging.error(error)
+
+        if error.response['Error']['Code'] == 'UserNotFoundException':
+            return {
+                'statusCode': 401,
+                'headers': {'Content-Type': 'application/json'},
+                'body': "{ 'message': 'Unauthorized' }",
+            }
+
+        return internal_error(error)
+
+    except Exception as error:
+        return internal_error(error)
+
+
+def internal_error(error):
+    logging.error(error)
+    return {
+        'statusCode': 500,
+        'headers': {'Content-Type': 'application/json'},
+        'body': "{ 'message': 'Internal server error' }",
+    }

src/sign-up/lambda_function.py

@@ -1,19 +1,44 @@
+import json
+import logging
 import os
+
 import boto3
-import json
+import botocore.exceptions
+import psycopg2
+import requests
 
 cognito_client = boto3.client('cognito-idp')
 
-user_pool_id = os.getenv('USER_POOL_ID')
+USER_POOL_ID = os.environ.get('USER_POOL_ID')
+
+headers = {'X-Aws-Parameters-Secrets-Token': os.environ.get('AWS_SESSION_TOKEN')}
+params_extension_endpoint = 'http://localhost:2773/systemsmanager/parameters/get?name='
+secrets_extension_endpoint = 'http://localhost:2773/secretsmanager/get?secretId='
+
+
+def get_param(param_name):
+    return json.loads(requests.get(params_extension_endpoint + param_name, headers=headers).text)
+
+
+def get_secret(secret_id):
+    return json.loads(requests.get(secrets_extension_endpoint + secret_id, headers=headers).text)
+
+
+def internal_error(error):
+    logging.error(error)
+    return {
+        'statusCode': 500,
+        'headers': {'Content-Type': 'application/json'},
+        'body': "{ 'message': 'Internal server error' }",
+    }
 
 
 def lambda_handler(event, context):
-    print(event)
-    body = json.loads(event.get("body", "{}"))
+    body = json.loads(event.get('body', '{}'))
 
-    email = body.get('email')
-    name = body.get('name')
-    cpf = body.get('cpf')
+    email = body.get('email')  # TODO: validate email address
+    name = body.get('name')  # TODO: validate name
+    cpf = body.get('cpf')  # TODO: validate CPF
 
     user_attributes = []
     if cpf:
@@ -34,22 +59,51 @@ def lambda_handler(event, context):
         }
 
     try:
+        RDS_PARAMS = get_param(os.environ.get('RDS_PARAM_ID'))
+        RDS_SECRETS = get_secret(os.environ.get('RDS_SECRET_ID'))
+
         response = cognito_client.admin_create_user(
-            UserPoolId=user_pool_id,
+            UserPoolId=USER_POOL_ID,
             Username=username,
             UserAttributes=user_attributes,
             MessageAction='SUPPRESS'
         )
-        print(response)
+
+        logging.info(response)
+
+        response = cognito_client.admin_add_user_to_group(
+            UserPoolId=USER_POOL_ID,
+            Username=username,
+            GroupName='customer'
+        )
+
+        logging.info(response)
+
+        # try:
+        #     conn = psycopg2.connect(
+        #         user=RDS_SECRETS['username'], password=RDS_SECRETS['password'],
+        #         host=RDS_PARAMS['endpoint'].split(':')[0], port=RDS_PARAMS['port'], dbname=RDS_PARAMS['name']
+        #     )
+        # except psycopg2.Error as error:
+        #     internal_error(error)
+
         return {
             'statusCode': 200,
             'headers': {'Content-Type': 'application/json'},
             'body': "{ 'message': 'User created successfully' }"
         }
-    except Exception as e:
-        print(e)
-        return {
-            'statusCode': 500,
-            'headers': {'Content-Type': 'application/json'},
-            'body': "{ 'message': 'Error creating user' }"
-        }
+
+    except botocore.exceptions.ClientError as error:
+        logging.error(error)
+
+        if error.response['Error']['Code'] == 'UsernameExistsException':
+            return {
+                'statusCode': 403,
+                'headers': {'Content-Type': 'application/json'},
+                'body': "{ 'message': 'Unauthorized' }"
+            }
+
+        return internal_error(error)
+
+    except Exception as error:
+        return internal_error(error)

src/sign-up/requirements.txt

@@ -0,0 +1,2 @@
+requests
+psycopg2-binary

terraform/main.tf

@@ -2,6 +2,26 @@ locals {
   runtime = "python3.12"
 }
 
+data "terraform_remote_state" "tech-challenge" {
+  backend = "s3"
+
+  config = {
+    bucket = "fiap-3soat-g15-infra-tech-challenge-state"
+    key    = "live/terraform.tfstate"
+    region = var.region
+  }
+}
+
+data "terraform_remote_state" "rds" {
+  backend = "s3"
+
+  config = {
+    bucket = "fiap-3soat-g15-infra-db-state"
+    key    = "live/terraform.tfstate"
+    region = var.region
+  }
+}
+
 resource "aws_cognito_user_pool" "user_pool" {
   name = "self-order-management"
 
@@ -55,6 +75,14 @@ resource "aws_cognito_user_pool_client" "client" {
   user_pool_id = aws_cognito_user_pool.user_pool.id
 }
 
+data "aws_ssm_parameter" "rds_param" {
+  name = data.terraform_remote_state.rds.outputs.rds_ssm_parameter_name
+}
+
+data "aws_secretsmanager_secret" "rds_secret" {
+  arn = data.terraform_remote_state.rds.outputs.db_instance_master_user_secret_arn
+}
+
 module "lambda_auth_sign_up" {
   source  = "terraform-aws-modules/lambda/aws"
   version = "7.2.2"
@@ -63,21 +91,48 @@ module "lambda_auth_sign_up" {
   handler       = "lambda_function.lambda_handler"
   runtime       = local.runtime
 
-  source_path = "../src/sign-up"
+  source_path = {
+    path             = "../src/sign-up"
+    pip_requirements = true
+  }
 
   environment_variables = {
-    USER_POOL_ID = aws_cognito_user_pool.user_pool.id
+    USER_POOL_ID  = aws_cognito_user_pool.user_pool.id
+    RDS_PARAM_ID  = data.aws_ssm_parameter.rds_param.name
+    RDS_SECRET_ID = data.aws_secretsmanager_secret.rds_secret.name
   }
 
   attach_policy_statements = true
   policy_statements = {
     cognito = {
-      effect    = "Allow"
-      actions   = ["cognito-idp:AdminCreateUser"]
-      resources = [aws_cognito_user_pool.user_pool.arn]
+      effect = "Allow"
+      actions = [
+        "cognito-idp:AdminCreateUser",
+        "cognito-idp:AdminAddUserToGroup"
+      ]
+      resources = [
+        aws_cognito_user_pool.user_pool.arn
+      ]
     }
   }
 
+
+  vpc_subnet_ids = data.terraform_remote_state.tech-challenge.outputs.private_subnets
+
+  attach_policies = true
+  policies = [
+    data.terraform_remote_state.rds.outputs.rds_secrets_read_only_policy_arn,
+    data.terraform_remote_state.rds.outputs.rds_params_read_only_policy_arn
+  ]
+  number_of_policies = 2
+
+  layers = [
+    # AWS Parameters and Secrets Lambda Extension for us-east-1
+    # https://docs.aws.amazon.com/secretsmanager/latest/userguide/retrieving-secrets_lambda.html
+    # https://docs.aws.amazon.com/systems-manager/latest/userguide/ps-integration-lambda-extensions.html#ps-integration-lambda-extensions-add
+    "arn:aws:lambda:us-east-1:177933569100:layer:AWS-Parameters-and-Secrets-Lambda-Extension:11"
+  ]
+
   tags = var.tags
 
   depends_on = [
@@ -103,9 +158,13 @@ module "lambda_auth_sign_in" {
   attach_policy_statements = true
   policy_statements = {
     cognito = {
-      effect    = "Allow"
-      actions   = ["cognito-idp:AdminInitiateAuth"]
-      resources = [aws_cognito_user_pool.user_pool.arn]
+      effect = "Allow"
+      actions = [
+        "cognito-idp:AdminInitiateAuth"
+      ]
+      resources = [
+        aws_cognito_user_pool.user_pool.arn
+      ]
     }
   }