diff --git a/lambda/main.py b/lambda/main.py index a1e7262..a23400b 100644 --- a/lambda/main.py +++ b/lambda/main.py @@ -1,15 +1,17 @@ + import subprocess import uuid import os - import logging +import sys +import json logger = logging.getLogger() logger.setLevel(logging.DEBUG) def lambda_handler(event, context): - logger.debug("Input event: {}".format(event)) + logger.debug("Input event: {}".format(json.dumps(event))) cmd = event['interpreter'] if event['command'] is not None: @@ -25,7 +27,10 @@ def lambda_handler(event, context): # For the subprocess environment, use all of the existing env vars, plus # any new ones. New ones with the same name will overwrite. - new_env = os.environ.copy() | event['environment'] + new_env = os.environ.copy() + # Set the python path to include everything that is given by default to Python functions + new_env['PYTHONPATH'] = ':'.join(sys.path) + new_env.update(event['environment']) # Start the process p = subprocess.Popen( @@ -84,6 +89,9 @@ def lambda_handler(event, context): {} '''.format(exit_code, stdout, stderr)) + logger.debug("Exit code: {}".format(exit_code)) + logger.debug("Stdout {}".format(stdout)) + logger.debug("Stderr {}".format(stderr)) return { 'exit_code': exit_code, 'stdout': stdout, diff --git a/main.tf b/main.tf index 516d69b..55c1c7b 100644 --- a/main.tf +++ b/main.tf @@ -21,12 +21,14 @@ module "shell_lambda" { archive_output_directory = "${path.module}/archives/" lambda_config = { function_name = "invicton-labs-aws-lambda-shell-${random_id.lambda.hex}" + description = var.lambda_description handler = "main.lambda_handler" - runtime = "python3.9" + runtime = var.lambda_runtime timeout = var.lambda_timeout memory_size = var.lambda_memory_size role = local.lambda_role layers = var.lambda_layer_arns + architectures = [var.lambda_architecture] tags = { "ModuleAuthor" = "InvictonLabs" "ModuleUrl" = "https://registry.terraform.io/modules/Invicton-Labs/lambda-shell/aws" diff --git a/outputs.tf b/outputs.tf index 57051f0..52c413a 100644 --- a/outputs.tf +++ b/outputs.tf @@ -3,3 +3,8 @@ output "invicton_labs_lambda_shell_arn" { // Use the "complete" output of the inner Lambda module so that the function ARN can't be used until everything has been completed (permissions have been applied) value = module.shell_lambda.complete ? module.shell_lambda.lambda.arn : module.shell_lambda.lambda.arn } + +output "runtime" { + description = "The Lambda runtime that is used for the Lambda shell." + value = var.lambda_runtime +} diff --git a/variables.tf b/variables.tf index b31a9a4..d3ddc1b 100644 --- a/variables.tf +++ b/variables.tf @@ -1,3 +1,9 @@ +variable "lambda_description" { + description = "The description string to apply to the Lambda function." + type = string + default = "Invicton-Labs/lambda-shell/aws (https://registry.terraform.io/modules/Invicton-Labs/lambda-shell/aws)" +} + variable "lambda_timeout" { description = "The timeout (in seconds) for the Lambda function that is running the shell command." type = number @@ -16,6 +22,16 @@ variable "lambda_role_arn" { default = null } +variable "lambda_architecture" { + description = "The architecture to use for the Lambda function." + type = string + default = "x86_64" + validation { + condition = contains(["x86_64", "arm64"], var.lambda_architecture) + error_message = "The `lambda_architecture` variable must be `x86_64` or `arm64`." + } +} + variable "lambda_role_policies_json" { description = "A list of JSON-encoded policies to apply to a new role that will be created for the Lambda that runs shell commands. Conflicts with `lambda_role_arn`. If neither is provided, the module will attempt to use the role that the Terraform caller has assumed (if a role has been assumed)." type = list(string) @@ -62,8 +78,18 @@ variable "lambda_vpc_config" { variable "lambda_layer_arns" { description = "A list of Lambda Layer ARNs to attach to the Lambda." - type = list(string) - default = [] + type = list(string) + default = [] +} + +variable "lambda_runtime" { + description = "The runtime to use for the lambda shell." + type = string + default = "python3.9" + validation { + condition = contains(["python3.7", "python3.8", "python3.9"], var.lambda_runtime) + error_message = "The `lambda_runtime` variable must be `python3.7`, `python3.8`, or `python3.9`." + } } data "aws_caller_identity" "current" {}