Add PYTHONPATH to subprocess

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,

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"

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
+}

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" {}