support new --extra-tag argument, tag all resources

awick · awick · commit 865fc43aa1fa · 2024-05-15T15:59:38.000-04:00
We now tag all resources with the arkime_cluster tag.
Can also use --extra-tag &lt;key&gt; &lt;value&gt; to have more tags added to all
resources
diff --git a/cdk-lib/cloud-demo.ts b/cdk-lib/cloud-demo.ts
@@ -25,23 +25,37 @@ const env: Environment = {
 
 switch(params.type) {
 case 'ClusterMgmtParams': {
+    function addTags (stack: cdk.Stack) {
+        const clusterParams = params as prms.ClusterMgmtParams;
+        cdk.Tags.of(stack).add('arkime_cluster', clusterParams.nameCluster);
+
+        if (clusterParams.userConfig.extraTags) {
+            for (const tag of clusterParams.userConfig.extraTags) {
+                cdk.Tags.of(stack).add(tag.key, tag.value);
+            }
+        }
+    }
     const captureBucketStack = new CaptureBucketStack(app, params.stackNames.captureBucket, {
         env: env,
         planCluster: params.planCluster,
         ssmParamName: params.nameCaptureBucketSsmParam,
     });
+    addTags(captureBucketStack);
 
     const captureVpcStack = new CaptureVpcStack(app, params.stackNames.captureVpc, {
         env: env,
         planCluster: params.planCluster,
     });
+    addTags(captureVpcStack);
 
     const osDomainStack = new OpenSearchDomainStack(app, params.stackNames.osDomain, {
         env: env,
         captureVpc: captureVpcStack.vpc,
         planCluster: params.planCluster,
         ssmParamName: params.nameOSDomainSsmParam,
     });
+    addTags(osDomainStack);
+
     osDomainStack.addDependency(captureVpcStack);
 
     const captureNodesStack = new CaptureNodesStack(app, params.stackNames.captureNodes, {
@@ -59,6 +73,8 @@ case 'ClusterMgmtParams': {
         ssmParamNameCluster: params.nameClusterSsmParam,
         userConfig: params.userConfig
     });
+    addTags(captureNodesStack);
+
     captureNodesStack.addDependency(captureBucketStack);
     captureNodesStack.addDependency(captureVpcStack);
     captureNodesStack.addDependency(osDomainStack);
@@ -73,6 +89,7 @@ case 'ClusterMgmtParams': {
             env: env,
             captureVpc: captureVpcStack.vpc
         });
+        addTags(captureTgwStack);
         captureTgwStack.addDependency(captureVpcStack);
 
         const viewerVpcStack = new ViewerVpcStack(app, params.stackNames.viewerVpc, {
@@ -81,6 +98,8 @@ case 'ClusterMgmtParams': {
             captureVpc: captureVpcStack.vpc,
             viewerVpcPlan: params.planCluster.viewerVpc
         });
+        addTags(viewerVpcStack);
+
         viewerVpcStack.addDependency(captureVpcStack);
         viewerVpcStack.addDependency(captureNodesStack);
         viewerVpcStack.addDependency(captureTgwStack);
@@ -102,6 +121,7 @@ case 'ClusterMgmtParams': {
         planCluster: params.planCluster,
         userConfig: params.userConfig,
     });
+    addTags(viewerNodesStack);
     viewerNodesStack.addDependency(captureBucketStack);
     viewerNodesStack.addDependency(vpcStackToUse);
     viewerNodesStack.addDependency(osDomainStack);
diff --git a/cdk-lib/core/context-types.ts b/cdk-lib/core/context-types.ts
@@ -98,6 +98,7 @@ export interface UserConfig {
     replicas: number;
     pcapDays: number;
     viewerPrefixList: string;
+    extraTags: ExtraTag[]
 }
 
 /**
@@ -112,3 +113,11 @@ export interface ClusterMgmtStackNames {
     viewerNodes: string;
     viewerVpc: string;
 }
+
+/**
+ * Structure to hold the Extra Tags
+ */
+export interface ExtraTag {
+    key: string;
+    value: string
+}
diff --git a/manage_arkime.py b/manage_arkime.py
@@ -130,13 +130,25 @@ def demo_traffic_destroy(ctx):
     default=None,
     type=click.STRING,
     required=False)
+@click.option(
+    "--extra-tag",
+    help=("Extra AWS Tag to apply to all resources."),
+    default=None,
+    nargs=2,
+    type=(click.STRING, click.STRING),
+    multiple=True,
+    required=False)
 @click.pass_context
 def cluster_create(ctx, name, expected_traffic, spi_days, history_days, replicas, pcap_days, preconfirm_usage,
-                   just_print_cfn, capture_cidr, viewer_cidr, viewer_prefix_list):
+                   just_print_cfn, capture_cidr, viewer_cidr, viewer_prefix_list, extra_tag):
     profile = ctx.obj.get("profile")
     region = ctx.obj.get("region")
+    extra_tags = []
+    if extra_tag:
+        for key, value in extra_tag:
+            extra_tags.append({"key": key, "value": value})
     cmd_cluster_create(profile, region, name, expected_traffic, spi_days, history_days, replicas, pcap_days,
-                       preconfirm_usage, just_print_cfn, capture_cidr, viewer_cidr, viewer_prefix_list)
+                       preconfirm_usage, just_print_cfn, capture_cidr, viewer_cidr, viewer_prefix_list, extra_tags)
 cli.add_command(cluster_create)
 
 @click.command(help="Tears down the Arkime Cluster in your account; by default, leaves your data intact")
diff --git a/manage_arkime/commands/cluster_create.py b/manage_arkime/commands/cluster_create.py
@@ -3,7 +3,7 @@
 import logging
 import shutil
 import sys
-from typing import Callable, List
+from typing import Callable, List, Tuple, Dict
 
 import arkime_interactions.config_wrangling as config_wrangling
 from aws_interactions.acm_interactions import upload_default_elb_cert
@@ -31,7 +31,8 @@
 logger = logging.getLogger(__name__)
 
 def cmd_cluster_create(profile: str, region: str, name: str, expected_traffic: float, spi_days: int, history_days: int, replicas: int,
-                       pcap_days: int, preconfirm_usage: bool, just_print_cfn: bool, capture_cidr: str, viewer_cidr: str, viewer_prefix_list: str):
+                       pcap_days: int, preconfirm_usage: bool, just_print_cfn: bool, capture_cidr: str, viewer_cidr: str, viewer_prefix_list: str,
+                       extra_tags: List[Dict[str, str]]):
     logger.debug(f"Invoking cluster-create with profile '{profile}' and region '{region}'")
 
     aws_provider = AwsClientProvider(aws_profile=profile, aws_region=region)
@@ -50,7 +51,7 @@ def cmd_cluster_create(profile: str, region: str, name: str, expected_traffic: f
 
     # Generate our capacity plan, then confirm it's what the user expected and it's safe to proceed with the operation
     previous_user_config = _get_previous_user_config(name, aws_provider)
-    next_user_config = _get_next_user_config(name, expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list,  aws_provider)
+    next_user_config = _get_next_user_config(name, expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list, extra_tags, aws_provider)
     previous_capacity_plan = _get_previous_capacity_plan(name, aws_provider)
     next_capacity_plan = _get_next_capacity_plan(next_user_config, previous_capacity_plan, capture_cidr, viewer_cidr, aws_provider)
 
@@ -82,9 +83,6 @@ def cmd_cluster_create(profile: str, region: str, name: str, expected_traffic: f
         # Deploy the CFN resources
         cdk_client.deploy(stacks_to_deploy, context=create_context)
 
-        # Tag the OpenSearch Domain
-        _tag_domain(name, aws_provider)
-
         # Kick off Events to ensure that ISM is set up on the CFN-created OpenSearch Domain
         _configure_ism(name, next_user_config.historyDays, next_user_config.spiDays, next_user_config.replicas, aws_provider)
 
@@ -156,7 +154,7 @@ def _get_previous_user_config(cluster_name: str, aws_provider: AwsClientProvider
         return UserConfig(None, None, None, None, None)
 
 def _get_next_user_config(cluster_name: str, expected_traffic: float, spi_days: int, history_days: int, replicas: int,
-                          pcap_days: int, viewer_prefix_list: str, aws_provider: AwsClientProvider) -> UserConfig:
+                          pcap_days: int, viewer_prefix_list: str, extra_tags: str, aws_provider: AwsClientProvider) -> UserConfig:
     # At least one parameter isn't defined
     if None in [expected_traffic, spi_days, replicas, pcap_days, history_days, viewer_prefix_list]:
         # Re-use the existing configuration if it exists
@@ -183,14 +181,16 @@ def _get_next_user_config(cluster_name: str, expected_traffic: float, spi_days:
                 user_config.pcapDays = pcap_days
             if viewer_prefix_list is not None:
                 user_config.viewerPrefixList = viewer_prefix_list
+            if extra_tags is not None:
+                user_config.extraTags = extra_tags
             return user_config
 
         # Existing configuration doesn't exist, use defaults
         except ssm_ops.ParamDoesNotExist:
-            return UserConfig(expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list)
+            return UserConfig(expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list, extra_tags)
     # All of the parameters defined
     else:
-        return UserConfig(expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list)
+        return UserConfig(expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list, extra_tags)
 
 def _get_previous_capacity_plan(cluster_name: str, aws_provider: AwsClientProvider) -> ClusterPlan:
     # Pull the existing plan, if possible
@@ -358,21 +358,6 @@ def _configure_ism(cluster_name: str, history_days: int, spi_days: int, replicas
         aws_provider
     )
 
-def _tag_domain(cluster_name: str, aws_provider: AwsClientProvider):
-    os_domain_Arn = ssm_ops.get_ssm_param_json_value(
-        constants.get_opensearch_domain_ssm_param_name(cluster_name),
-        "domainArn",
-        aws_provider
-    )
-
-    opensearch_client = aws_provider.get_opensearch()
-    opensearch_client.add_tags(
-        ARN=os_domain_Arn,
-        TagList=[
-            {"Key": "arkime_cluster", "Value": cluster_name},
-        ]
-    )
-
 def _get_stacks_to_deploy(cluster_name: str, next_user_config: UserConfig, next_capacity_plan: ClusterPlan) -> List[str]:
     # This list defines what actually gets deployed, as opposed to what the CDK has in its blueprint as being available
     # to deploy.
diff --git a/manage_arkime/core/user_config.py b/manage_arkime/core/user_config.py
@@ -1,6 +1,6 @@
 from dataclasses import dataclass, fields
 import logging
-from typing import Dict
+from typing import Dict, List
 
 from core.capacity_planning import (MINIMUM_TRAFFIC, DEFAULT_SPI_DAYS, DEFAULT_REPLICAS, DEFAULT_S3_STORAGE_DAYS, DEFAULT_HISTORY_DAYS)
 
@@ -14,30 +14,31 @@ class UserConfig:
     replicas: int
     pcapDays: int
     viewerPrefixList: str = None
+    extraTags: List[Dict[str, str]] = None
+
+    def __init__(self, expectedTraffic: float, spiDays: int, historyDays: int, replicas: int, pcapDays: int, viewerPrefixList: str = None, extraTags: List[Dict[str, str]] = []):
+        self.expectedTraffic = expectedTraffic
+        self.spiDays = spiDays
+        self.historyDays = historyDays
+        self.replicas = replicas
+        self.pcapDays = pcapDays
+        self.viewerPrefixList = viewerPrefixList
+        self.extraTags = extraTags
 
-    def __init__(self, expectedTraffic: float, spiDays: int, historyDays: int, replicas: int, pcapDays: int, viewerPrefixList: str = None):
         if (expectedTraffic is None):
-            expectedTraffic = MINIMUM_TRAFFIC
+            self.expectedTraffic = MINIMUM_TRAFFIC
 
         if (spiDays is None):
-            spiDays = DEFAULT_SPI_DAYS
+            self.spiDays = DEFAULT_SPI_DAYS
 
         if (historyDays is None):
-            historyDays = DEFAULT_HISTORY_DAYS
+            self.historyDays = DEFAULT_HISTORY_DAYS
 
         if (replicas is None):
-            replicas = DEFAULT_REPLICAS
+            self.replicas = DEFAULT_REPLICAS
 
         if (pcapDays is None):
-            pcapDays = DEFAULT_S3_STORAGE_DAYS
-
-
-        self.expectedTraffic = expectedTraffic
-        self.spiDays = spiDays
-        self.historyDays = historyDays
-        self.replicas = replicas
-        self.pcapDays = pcapDays
-        self.viewerPrefixList = viewerPrefixList
+            self.pcapDays = DEFAULT_S3_STORAGE_DAYS
 
     """ Only process fields we still need, this allows us to ignore config no longer used """
     @classmethod
@@ -47,12 +48,21 @@ def from_dict(cls, d):
         return cls(**valid_kwargs)
 
     def __eq__(self, other):
+        set1 = None
+        if self.extraTags is not None:
+            set1 = {frozenset(d.items()) for d in self.extraTags}
+
+        set2 = None
+        if other.extraTags is not None:
+            set2 = {frozenset(d.items()) for d in other.extraTags}
+
         return (self.expectedTraffic == other.expectedTraffic and
                 self.spiDays == other.spiDays and
                 self.historyDays == other.historyDays and
                 self.replicas == other.replicas and
                 self.pcapDays == other.pcapDays and
-                self.viewerPrefixList == other.viewerPrefixList)
+                self.viewerPrefixList == other.viewerPrefixList and
+                set1 == set2)
 
     def to_dict(self) -> Dict[str, any]:
         return {
@@ -62,5 +72,6 @@ def to_dict(self) -> Dict[str, any]:
             'pcapDays': self.pcapDays,
             'historyDays': self.historyDays,
             'viewerPrefixList': self.viewerPrefixList,
+            'extraTags': self.extraTags,
         }
 
diff --git a/test_manage_arkime/commands/test_cluster_create.py b/test_manage_arkime/commands/test_cluster_create.py
