Skip to content

Commit

Permalink
Extra tags (#169)
Browse files Browse the repository at this point in the history
* support new --extra-tag argument, tag all resources

We now tag all resources with the arkime_cluster tag.
Can also use --extra-tag <key> <value> to have more tags added to all
resources

* fix eslint

* fix ruff lint

* clarify previous tags are not deleted
  • Loading branch information
awick authored May 16, 2024
1 parent c23499d commit 04a4732
Show file tree
Hide file tree
Showing 6 changed files with 95 additions and 103 deletions.
21 changes: 21 additions & 0 deletions cdk-lib/cloud-demo.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,25 +23,40 @@ const env: Environment = {
region: params.awsRegion
};

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);
}
}
}

switch(params.type) {
case 'ClusterMgmtParams': {
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, {
Expand All @@ -59,6 +74,8 @@ case 'ClusterMgmtParams': {
ssmParamNameCluster: params.nameClusterSsmParam,
userConfig: params.userConfig
});
addTags(captureNodesStack);

captureNodesStack.addDependency(captureBucketStack);
captureNodesStack.addDependency(captureVpcStack);
captureNodesStack.addDependency(osDomainStack);
Expand All @@ -73,6 +90,7 @@ case 'ClusterMgmtParams': {
env: env,
captureVpc: captureVpcStack.vpc
});
addTags(captureTgwStack);
captureTgwStack.addDependency(captureVpcStack);

const viewerVpcStack = new ViewerVpcStack(app, params.stackNames.viewerVpc, {
Expand All @@ -81,6 +99,8 @@ case 'ClusterMgmtParams': {
captureVpc: captureVpcStack.vpc,
viewerVpcPlan: params.planCluster.viewerVpc
});
addTags(viewerVpcStack);

viewerVpcStack.addDependency(captureVpcStack);
viewerVpcStack.addDependency(captureNodesStack);
viewerVpcStack.addDependency(captureTgwStack);
Expand All @@ -102,6 +122,7 @@ case 'ClusterMgmtParams': {
planCluster: params.planCluster,
userConfig: params.userConfig,
});
addTags(viewerNodesStack);
viewerNodesStack.addDependency(captureBucketStack);
viewerNodesStack.addDependency(vpcStackToUse);
viewerNodesStack.addDependency(osDomainStack);
Expand Down
9 changes: 9 additions & 0 deletions cdk-lib/core/context-types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ export interface UserConfig {
replicas: number;
pcapDays: number;
viewerPrefixList: string;
extraTags: ExtraTag[]
}

/**
Expand All @@ -112,3 +113,11 @@ export interface ClusterMgmtStackNames {
viewerNodes: string;
viewerVpc: string;
}

/**
* Structure to hold the Extra Tags
*/
export interface ExtraTag {
key: string;
value: string
}
16 changes: 14 additions & 2 deletions manage_arkime.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 add to all resources, previously added tags are not deleted."),
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")
Expand Down
33 changes: 9 additions & 24 deletions manage_arkime/commands/cluster_create.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import logging
import shutil
import sys
from typing import Callable, List
from typing import Callable, List, Dict

import arkime_interactions.config_wrangling as config_wrangling
from aws_interactions.acm_interactions import upload_default_elb_cert
Expand Down Expand Up @@ -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)
Expand All @@ -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)

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down
43 changes: 27 additions & 16 deletions manage_arkime/core/user_config.py
Original file line number Diff line number Diff line change
@@ -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)

Expand All @@ -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
Expand All @@ -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 {
Expand All @@ -62,5 +72,6 @@ def to_dict(self) -> Dict[str, any]:
'pcapDays': self.pcapDays,
'historyDays': self.historyDays,
'viewerPrefixList': self.viewerPrefixList,
'extraTags': self.extraTags,
}

Loading

0 comments on commit 04a4732

Please sign in to comment.