Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Extra tags #169

Merged
merged 4 commits into from
May 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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);
awick marked this conversation as resolved.
Show resolved Hide resolved
}
}
}

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