Add tvm-bot lambda

.github/workflows/terraform_apply.yml

@@ -18,6 +18,7 @@ jobs:
       AWS_ACCESS_KEY_ID: ${{ secrets.TERRAFORM_AWS_ACCESS_KEY_ID }}
       AWS_SECRET_ACCESS_KEY: ${{ secrets.TERRAFORM_AWS_SECRET_ACCESS_KEY }}
       GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      TF_VAR_tvm_bot_webhook_secret: ${{ secrets.TVM_BOT_WEBHOOK_SECRET }}
     defaults:
       run:
         working-directory: ./terraform

.github/workflows/terraform_plan.yml

@@ -52,13 +52,15 @@ jobs:
       AWS_ACCESS_KEY_ID: ${{ secrets.TERRAFORM_AWS_ACCESS_KEY_ID }}
       AWS_SECRET_ACCESS_KEY: ${{ secrets.TERRAFORM_AWS_SECRET_ACCESS_KEY }}
       GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      TF_VAR_tvm_bot_webhook_secret: ${{ secrets.TVM_BOT_WEBHOOK_SECRET }}
     runs-on: ubuntu-latest
     defaults:
       run:
         working-directory: ./terraform
     needs: validate
     if: needs.validate.outputs.valid_workflow == 'True'
-    #These steps run if either the PR is within the same repo or if the PR is on a fork and the committer has deployer access
+    # These steps run if either the PR is within the same repo or if the PR is
+    # on a fork and the committer has deployer access
     steps:
       - uses: actions/checkout@v2
         with:

.github/workflows/test.yml

@@ -0,0 +1,33 @@
+name: test tvm_bot
+on:
+  push:
+  pull_request:
+    branches:
+      - main
+    paths:
+      - terraform/tvm_bot/**
+  pull_request_target:
+    branches:
+      - main
+    paths:
+      - terraform/tvm_bot/**
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ./terraform/tvm_bot
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+      - name: Install dependencies
+        run: |
+          set -eux
+          pip install -r requirements.txt
+      - name: Run unit tests
+        run: |
+          set -eux
+          PYTHONPATH=$(pwd) pytest --tb=native

terraform/README.md

@@ -1,3 +1,24 @@
 # terraform
 
 This folder handles the Terraform configuration for TVM Jenkins Infrastructure.
+
+## Local Usage
+
+```bash
+# if anything is broken, remove all terraform local files
+git clean -xfd .
+
+# set credentials
+export AWS_ACCESS_KEY_ID=...
+export AWS_SECRET_ACCESS_KEY=...
+
+# get terraform state
+terraform init
+
+# the workspace must be selected or else 'plan' will not read the correct state
+terraform workspace new tvm-ci-prod
+terraform workspace select tvm-ci-prod
+
+# run the actual plan against AWS
+terraform plan -var-file=vars/tvm-ci-prod.auto.tfvars
+```

terraform/tvm_bot.tf

@@ -0,0 +1,136 @@
+# Note: The folder must be prepared (i.e. dependencies inlined) by running
+# 'make' in tvm_bot/ first
+data "archive_file" "tvm_bot_archive" {
+  type        = "zip"
+  source_dir  = "tvm_bot"
+  output_path = "tvm_bot.zip"
+}
+
+# See https://registry.terraform.io/providers/hashicorp/aws/2.34.0/docs/guides/serverless-with-aws-lambda-and-api-gateway
+resource "aws_lambda_function" "tvm_bot_lambda" {
+  function_name = "tvm_bot"
+
+  filename         = "tvm_bot.zip"
+  source_code_hash = data.archive_file.tvm_bot_archive.output_base64sha256
+
+  handler = "lambda_function.lambda_handler"
+  runtime = "python3.9"
+
+  role = aws_iam_role.lambda_tvm_bot_exec.arn
+
+  depends_on = [aws_iam_role_policy.logs]
+
+  environment {
+    variables = {
+      WEBHOOK_SECRET = var.tvm_bot_webhook_secret
+      GITHUB_TOKEN = var.tvm_bot_github_token
+    }
+  }
+}
+
+# IAM role which dictates what other AWS services the Lambda function
+# may access.
+resource "aws_iam_role" "lambda_tvm_bot_exec" {
+  name = "tvm_bot_lambda"
+
+  assume_role_policy = <<EOF
+{
+  "Version": "2012-10-17",
+  "Statement": [
+    {
+      "Action": "sts:AssumeRole",
+      "Principal": {
+        "Service": "lambda.amazonaws.com"
+      },
+      "Effect": "Allow",
+      "Sid": ""
+    }
+  ]
+}
+EOF
+}
+
+resource "aws_iam_role_policy" "logs" {
+  name = "lambda-logs"
+  role = aws_iam_role.lambda_tvm_bot_exec.name
+  policy = jsonencode({
+    "Statement" : [
+      {
+        "Action" : [
+          "logs:CreateLogGroup",
+          "logs:CreateLogStream",
+          "logs:PutLogEvents",
+        ],
+        "Effect" : "Allow",
+        "Resource" : "arn:aws:logs:*:*:*",
+      }
+    ]
+  })
+}
+
+resource "aws_api_gateway_rest_api" "tvm_bot" {
+  name        = "tvm_bot"
+  description = "API gateway for tvm_bot lambda"
+}
+
+
+resource "aws_api_gateway_resource" "proxy" {
+  rest_api_id = aws_api_gateway_rest_api.tvm_bot.id
+  parent_id   = aws_api_gateway_rest_api.tvm_bot.root_resource_id
+  path_part   = "{proxy+}"
+}
+
+resource "aws_api_gateway_method" "proxy" {
+  rest_api_id   = aws_api_gateway_rest_api.tvm_bot.id
+  resource_id   = aws_api_gateway_resource.proxy.id
+  http_method   = "POST"
+  authorization = "NONE"
+}
+
+resource "aws_api_gateway_integration" "lambda" {
+  rest_api_id = aws_api_gateway_rest_api.tvm_bot.id
+  resource_id = aws_api_gateway_method.proxy.resource_id
+  http_method = aws_api_gateway_method.proxy.http_method
+
+  integration_http_method = "POST"
+  type                    = "AWS_PROXY"
+  uri                     = aws_lambda_function.tvm_bot_lambda.invoke_arn
+}
+
+resource "aws_api_gateway_method" "proxy_root" {
+  rest_api_id   = aws_api_gateway_rest_api.tvm_bot.id
+  resource_id   = aws_api_gateway_rest_api.tvm_bot.root_resource_id
+  http_method   = "POST"
+  authorization = "NONE"
+}
+
+resource "aws_api_gateway_integration" "lambda_root" {
+  rest_api_id = aws_api_gateway_rest_api.tvm_bot.id
+  resource_id = aws_api_gateway_method.proxy_root.resource_id
+  http_method = aws_api_gateway_method.proxy_root.http_method
+
+  integration_http_method = "POST"
+  type                    = "AWS_PROXY"
+  uri                     = aws_lambda_function.tvm_bot_lambda.invoke_arn
+}
+
+resource "aws_api_gateway_deployment" "tvm_bot" {
+  depends_on = [
+    aws_api_gateway_integration.lambda,
+    aws_api_gateway_integration.lambda_root,
+  ]
+
+  rest_api_id = aws_api_gateway_rest_api.tvm_bot.id
+  stage_name  = "webhook"
+}
+
+resource "aws_lambda_permission" "apigw" {
+  statement_id  = "AllowAPIGatewayInvoke"
+  action        = "lambda:InvokeFunction"
+  function_name = aws_lambda_function.tvm_bot_lambda.function_name
+  principal     = "apigateway.amazonaws.com"
+
+  # The /*/* portion grants access from any method on any resource
+  # within the API Gateway "REST API".
+  source_arn = "${aws_api_gateway_rest_api.tvm_bot.execution_arn}/*/*"
+}

terraform/tvm_bot/Makefile

@@ -0,0 +1,5 @@
+tests:
+	PYTHONPATH=. pytest --tb=native test/test_welcome_message.py
+
+prepare:
+	zip tvm_bot.zip lambda_function.py

terraform/tvm_bot/conftest.py

@@ -0,0 +1,7 @@
+import os
+
+
+def pytest_generate_tests(metafunc):
+    # lambda_function needs a WEBHOOK_SECRET to run, so set that up for local
+    # testing
+    os.environ["WEBHOOK_SECRET"] = "test"

terraform/tvm_bot/lambda_function.py

@@ -0,0 +1,61 @@
+import hashlib
+import hmac
+import json
+import os
+
+from typing import Dict, Any
+from tvm_bot import github_pr_comment
+
+
+def check_hash(payload: bytes, expected: str) -> bool:
+    """
+    GitHub webhooks should be signed with a predetermined secret. This returns
+    True if the signature is valid.
+    """
+    signature = hmac.new(
+        os.environ["WEBHOOK_SECRET"].encode("utf-8"), payload, hashlib.sha256
+    ).hexdigest()
+    return hmac.compare_digest(signature, expected)
+
+
+def should_handle_event(event_type: str) -> bool:
+    return event_type in {"status", "pull_request"}
+
+
+def handle_event(event_type: str, data: Dict[str, Any]) -> None:
+    if event_type == "pull_request":
+        github_pr_comment.github_pr_comment(
+            data, user="driazati", repo="tvm", dry_run=False
+        )
+
+
+def lambda_handler(event, context):
+    expected = event["headers"].get("X-Hub-Signature-256", "").split("=")[1]
+    payload = event["body"].encode("utf-8")
+
+    # Check that the signature matches the secret on GitHub
+    if not check_hash(payload, expected):
+        return {"statusCode": 403, "body": "Forbidden"}
+
+    # Check that the webhook event is one that matters
+    event_type = event["headers"]["X-GitHub-Event"]
+    if not should_handle_event(event_type):
+        return {"statusCode": 400, "body": "Not processing event"}
+
+    data = json.loads(event["body"])
+    handle_event(event_type, data)
+    return {"statusCode": 200, "body": f"ok: {event_type}"}
+
+
+if __name__ == "__main__":
+    # For local runs
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run the lambda handler")
+    parser.add_argument("--payload", help="webhook JSON payload", required=True)
+    parser.add_argument(
+        "--event", help="webhook event type (e.g. pull_request, status)", required=True
+    )
+    args = parser.parse_args()
+
+    handle_event(args.event, json.loads(args.payload))

terraform/tvm_bot/requirements.txt

@@ -0,0 +1 @@
+pytest==6.2.5

terraform/tvm_bot/test/test_tvm_bot.py

@@ -0,0 +1,46 @@
+import hashlib
+import hmac
+import json
+
+from typing import Dict, Any
+
+import lambda_function
+
+
+def sign_payload(payload: Dict[str, Any], webhook_secret: str) -> str:
+    payload_data = json.dumps(payload).encode("utf-8")
+    signature = hmac.new(
+        webhook_secret.encode("utf-8"), payload_data, hashlib.sha256
+    ).hexdigest()
+    return signature
+
+
+def invoke_lambda(
+    event_type: str, data: Dict[str, Any], webhook_secret: str = "test"
+) -> Any:
+    event = {
+        "body": json.dumps(data),
+        "headers": {
+            "X-Hub-Signature-256": f"payload={sign_payload(payload=data, webhook_secret=webhook_secret)}",
+            "X-GitHub-Event": event_type,
+        },
+    }
+    return lambda_function.lambda_handler(event=event, context=None)
+
+
+def test_invalid_secret() -> None:
+    """
+    Ensure that requests without the proper signature will get 403 errors
+    """
+    result = invoke_lambda("status", {}, webhook_secret="bad")
+    assert result["statusCode"] == 403
+    assert result["body"] == "Forbidden"
+
+
+def test_invalid_event() -> None:
+    """
+    Ensure that irrelevant events are ignored
+    """
+    result = invoke_lambda("something", {})
+    assert result["statusCode"] == 400
+    assert result["body"] == "Not processing event"