feat: code for the lambda for exporting account tags (#235)

Add the Terraform and GitHub workflow updates to manage
the Lambda function.

Co-authored-by: Pat Heard <patrick.heard@cds-snc.ca>

.github/workflows/tf-apply.yml

@@ -59,6 +59,11 @@ jobs:
             account: 659087519042
             role: cds-aws-lz-apply
 
+          - account_folder: org_account
+            module: billing_extract_tags
+            account: 659087519042
+            role: cds-aws-lz-apply
+
           - account_folder: log_archive
             module: main
             account: 274536870005

.github/workflows/tf-plan.yml

@@ -50,7 +50,11 @@ jobs:
             module: sentinel_oidc
             account: 659087519042
             role: cds-aws-lz-plan
-            admin_sso_role_arn: ADMIN_SSO_ROLE_ARN
+
+          - account_folder: org_account
+            module: billing_extract_tags
+            account: 659087519042
+            role: cds-aws-lz-plan
 
           - account_folder: log_archive
             module: main

.gitignore

@@ -12,4 +12,6 @@ debug.log
 creds.sh
 
 .tool-versions
-.DS_Store
+.DS_Store
+
+__pycache__

terragrunt/org_account/billing_extract_tags/eventbridge.tf

@@ -0,0 +1,11 @@
+resource "aws_cloudwatch_event_rule" "billing_extract_tags" {
+  name                = "billing_extract_tags_daily"
+  schedule_expression = "cron(0 5 * * ? *)"
+
+  tags = local.common_tags
+}
+
+resource "aws_cloudwatch_event_target" "billing_extract_tags" {
+  rule = aws_cloudwatch_event_rule.billing_extract_tags.name
+  arn  = aws_lambda_function.billing_extract_tags.arn
+}

terragrunt/org_account/billing_extract_tags/iam.tf

@@ -0,0 +1,80 @@
+resource "aws_iam_role" "billing_extract_tags" {
+  name               = "BillingExtractTags"
+  assume_role_policy = data.aws_iam_policy_document.billing_extract_tags_assume.json
+  tags               = local.common_tags
+}
+
+data "aws_iam_policy_document" "billing_extract_tags_assume" {
+  statement {
+    effect  = "Allow"
+    actions = ["sts:AssumeRole"]
+    principals {
+      type = "Service"
+      identifiers = [
+      "lambda.amazonaws.com"]
+    }
+  }
+}
+
+data "aws_iam_policy_document" "billing_extract_tags" {
+  statement {
+    effect    = "Allow"
+    actions   = ["logs:CreateLogGroup"]
+    resources = ["arn:aws:logs:${var.region}:${var.account_id}:*"]
+  }
+
+  statement {
+    effect = "Allow"
+    actions = [
+      "logs:CreateLogStream",
+      "logs:PutLogEvents"
+    ]
+    resources = [
+      "arn:aws:logs:${var.region}:${var.account_id}:log-group:/aws/lambda/billing_extract_tags:*"
+    ]
+  }
+
+  statement {
+    effect = "Allow"
+    actions = [
+      "s3:PutObject*",
+      "s3:ListBucket",
+      "s3:GetObject*",
+      "s3:DeleteObject*",
+      "s3:GetBucketLocation"
+    ]
+    resources = [
+      module.billing_extract_tags.s3_bucket_arn,
+      "${module.billing_extract_tags.s3_bucket_arn}/*",
+    ]
+  }
+}
+
+resource "aws_iam_policy" "billing_extract_tags" {
+  name   = "BillingExtractTags"
+  policy = data.aws_iam_policy_document.billing_extract_tags.json
+  tags   = local.common_tags
+}
+
+resource "aws_iam_role_policy_attachment" "billing_extract_tags" {
+  role       = aws_iam_role.billing_extract_tags.name
+  policy_arn = aws_iam_policy.billing_extract_tags.arn
+}
+
+data "aws_iam_policy" "org_read_only" {
+  arn = "arn:aws:iam::aws:policy/AWSOrganizationsReadOnlyAccess"
+}
+
+resource "aws_iam_role_policy_attachment" "org_read_only" {
+  role       = aws_iam_role.billing_extract_tags.name
+  policy_arn = data.aws_iam_policy.org_read_only.arn
+}
+
+data "aws_iam_policy" "lambda_insights" {
+  name = "CloudWatchLambdaInsightsExecutionRolePolicy"
+}
+
+resource "aws_iam_role_policy_attachment" "lambda_insights" {
+  role       = aws_iam_role.billing_extract_tags.name
+  policy_arn = data.aws_iam_policy.lambda_insights.arn
+}

terragrunt/org_account/billing_extract_tags/lambda.tf

@@ -0,0 +1,44 @@
+data "archive_file" "billing_extract_tags" {
+  type        = "zip"
+  source_file = "${path.module}/lambdas/billing_extract_tags/main.py"
+  output_path = "/tmp/main.py.zip"
+}
+
+resource "aws_lambda_function" "billing_extract_tags" {
+  function_name = "billing_extract_tags"
+  role          = aws_iam_role.billing_extract_tags.arn
+  runtime       = "python3.11"
+  handler       = "main.handler"
+  memory_size   = 1024
+  timeout       = 30
+
+  filename         = data.archive_file.billing_extract_tags.output_path
+  source_code_hash = filebase64sha256(data.archive_file.billing_extract_tags.output_path)
+
+  environment {
+    variables = {
+      TARGET_BUCKET = module.billing_extract_tags.s3_bucket_id
+    }
+  }
+
+  tracing_config {
+    mode = "PassThrough"
+  }
+
+  tags = local.common_tags
+}
+
+resource "aws_lambda_permission" "billing_extract_tags" {
+  statement_id  = "AllowBillingExtractTagsDaily"
+  action        = "lambda:InvokeFunction"
+  function_name = aws_lambda_function.billing_extract_tags.function_name
+  principal     = "events.amazonaws.com"
+  source_arn    = aws_cloudwatch_event_rule.billing_extract_tags.arn
+}
+
+resource "aws_cloudwatch_log_group" "billing_extract_tags" {
+  #checkov:skip=CKV_AWS_158:We trust the AWS provided keys
+  name              = "/aws/lambda/${aws_lambda_function.billing_extract_tags.function_name}"
+  retention_in_days = "14"
+  tags              = local.common_tags
+}

terragrunt/org_account/billing_extract_tags/lambdas/billing_extract_tags/Makefile

@@ -0,0 +1,20 @@
+default: 
+	python3 main.py
+
+fmt:
+	black . $(ARGS)
+
+install:
+	pip3 install --user -r requirements_dev.txt
+
+lint:
+	flake8 main.py
+
+test:
+	python -m pytest -s -vv .
+
+.PHONY: \
+	fmt \
+	install	\
+	lint \
+	test

terragrunt/org_account/billing_extract_tags/lambdas/billing_extract_tags/main.py

@@ -0,0 +1,73 @@
+"""
+Get the tags for all accounts in the organization and save them to an s3 bucket.
+This is then used to enrich the billing data with the account tags to allow
+for business unit filtering.
+"""
+import json
+import logging
+import os
+
+import boto3
+
+orgs = boto3.client("organizations")
+s3 = boto3.client("s3")
+
+TARGET_BUCKET = os.getenv("TARGET_BUCKET")
+
+
+def lambda_handler(event, context):
+    """
+    Get the tags for all accounts in the organization and save them to an s3 bucket
+    """
+    logging.info("Getting account tags")
+    accounts = []
+    accounts_result = orgs.list_accounts()
+    accounts += accounts_result["Accounts"]
+    while "NextToken" in accounts_result:
+        logging.info("Paginating accounts...")
+        accounts_result = orgs.list_accounts(NextToken=accounts_result["NextToken"])
+        accounts += accounts_result["Accounts"]
+
+    # Iterate over the accounts and get the tags and then add them to the account in the list
+    logging.info("Getting account tags")
+    for account in accounts:
+        account_tags = orgs.list_tags_for_resource(ResourceId=account["Id"])
+
+        # Convert the tags from {'Key': 'Name', 'Value': 'Dev'} to {'Name': 'Dev'}
+        account_tags["Tags"] = {
+            tag["Key"]: tag["Value"] for tag in account_tags["Tags"]
+        }
+        account["Tags"] = account_tags["Tags"]
+
+    # Get a set of all possible tag keys
+    tag_keys = set()
+    for account in accounts:
+        tag_keys.update(account["Tags"].keys())
+    logging.info(f"Found tag keys: {tag_keys}")
+
+    # Add empty strings for all the tags that are not present in the account
+    logging.info("Adding empty strings for missing tags")
+    for account in accounts:
+        for tag_key in tag_keys:
+            if tag_key not in account["Tags"]:
+                account["Tags"][tag_key] = ""
+
+    # Convert the tags into the format tag_key_name: tag_value and add them to the base object
+    logging.info("Converting tags to tag_key_name: tag_value")
+    for account in accounts:
+        for tag_key, tag_value in account["Tags"].items():
+            account[f"tag_{tag_key}"] = tag_value
+        del account["Tags"]
+
+    # .write json to string and add a newline between each record
+    logging.info("Writing account tags to json")
+    accounts = json.dumps(accounts, default=str)
+    # accounts = accounts.replace('},', '},\n')
+    # accounts = accounts.replace('[{', '[\n{')
+    logging.info(f"Accounts: {accounts}")
+
+    # save accounts to an s3 bucket
+    logging.info("Saving account tags to s3")
+    s3.put_object(Bucket=TARGET_BUCKET, Key="account_tags.json", Body=accounts)
+
+    return {"statusCode": 200}

terragrunt/org_account/billing_extract_tags/lambdas/billing_extract_tags/requirements_dev.txt

@@ -0,0 +1,4 @@
+boto3==1.34.54
+black==23.12.1 
+flake8==7.0.0
+pytest==7.4.4

terragrunt/org_account/billing_extract_tags/lambdas/billing_extract_tags/test_main.py

@@ -0,0 +1,64 @@
+import unittest
+import os
+from unittest.mock import call, patch, MagicMock
+from main import lambda_handler
+
+TARGET_BUCKET = "TARGET_BUCKET"
+ACCOUNT_TAGS_KEY = "account_tags.json"
+
+
+class TestLambdaHandler(unittest.TestCase):
+    def setUp(self):
+        self.event = {}
+        self.context = MagicMock()
+
+    @patch("main.orgs.list_accounts")
+    @patch("main.orgs.list_tags_for_resource")
+    @patch("main.s3.put_object")
+    @patch("main.TARGET_BUCKET", TARGET_BUCKET)
+    def test_lambda_handler(
+        self, mock_s3_put, mock_orgs_list_tags, mock_orgs_list_accounts
+    ):
+        mock_orgs_list_accounts.return_value = {
+            "Accounts": [{"Id": "123"}],
+        }
+        mock_orgs_list_tags.return_value = {"Tags": [{"Key": "Name", "Value": "Dev"}]}
+
+        response = lambda_handler(self.event, self.context)
+
+        mock_orgs_list_accounts.assert_called()
+        mock_orgs_list_tags.assert_called_with(ResourceId="123")
+        mock_s3_put.assert_called_with(
+            Bucket=TARGET_BUCKET,
+            Key=ACCOUNT_TAGS_KEY,
+            Body='[{"Id": "123", "tag_Name": "Dev"}]',
+        )
+
+        self.assertEqual(response, {"statusCode": 200})
+
+    @patch("main.orgs.list_accounts")
+    @patch("main.orgs.list_tags_for_resource")
+    @patch("main.s3.put_object")
+    @patch("main.TARGET_BUCKET", TARGET_BUCKET)
+    def test_lambda_handler_pagination(
+        self, mock_s3_put, mock_orgs_list_tags, mock_orgs_list_accounts
+    ):
+        mock_orgs_list_accounts.side_effect = [
+            {"Accounts": [{"Id": "123"}], "NextToken": "token"},
+            {"Accounts": [{"Id": "456"}]},
+        ]
+        mock_orgs_list_tags.side_effect = [
+            {"Tags": [{"Key": "Name", "Value": "Dev"}]},
+            {"Tags": [{"Key": "Name", "Value": "Prod"}]},
+        ]
+
+        lambda_handler(self.event, self.context)
+
+        mock_orgs_list_accounts.assert_any_call(NextToken="token")
+        mock_orgs_list_tags.assert_any_call(ResourceId="123")
+        mock_orgs_list_tags.assert_any_call(ResourceId="456")
+        mock_s3_put.assert_called_with(
+            Bucket=TARGET_BUCKET,
+            Key=ACCOUNT_TAGS_KEY,
+            Body='[{"Id": "123", "tag_Name": "Dev"}, {"Id": "456", "tag_Name": "Prod"}]',
+        )

terragrunt/org_account/billing_extract_tags/locals.tf

@@ -0,0 +1,6 @@
+locals {
+  common_tags = {
+    CostCentre = var.billing_code
+    Terraform  = "true"
+  }
+}