From c23499d75e032a6e2a10d64838a1d1327735e9d9 Mon Sep 17 00:00:00 2001
From: Andy Wick <andywick@amazon.com>
Date: Fri, 12 Apr 2024 09:53:11 -0400
Subject: [PATCH] support only setting some parameters on first create (#168)

* support only setting some parameters on first create

* refactor so defaults are done in UserConfig
---
 manage_arkime/commands/cluster_create.py |  9 ++++----
 manage_arkime/core/user_config.py        | 26 ++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/manage_arkime/commands/cluster_create.py b/manage_arkime/commands/cluster_create.py
index a85ff4e..5dc56d7 100644
--- a/manage_arkime/commands/cluster_create.py
+++ b/manage_arkime/commands/cluster_create.py
@@ -22,8 +22,7 @@
 from core.price_report import PriceReport
 import core.compatibility as compat
 from core.capacity_planning import (get_capture_node_capacity_plan, get_viewer_node_capacity_plan, get_ecs_sys_resource_plan, get_os_domain_plan,
-                                    ClusterPlan, VpcPlan, MINIMUM_TRAFFIC, DEFAULT_SPI_DAYS, DEFAULT_REPLICAS, get_capture_vpc_plan,
-                                    S3Plan, DEFAULT_S3_STORAGE_CLASS, DEFAULT_S3_STORAGE_DAYS, DEFAULT_HISTORY_DAYS,
+                                    ClusterPlan, VpcPlan, get_capture_vpc_plan, S3Plan, DEFAULT_S3_STORAGE_CLASS,
                                     CaptureNodesPlan, ViewerNodesPlan, DataNodesPlan, EcsSysResourcePlan, MasterNodesPlan, OSDomainPlan,
                                     get_viewer_vpc_plan)
 import core.versioning as ver
@@ -167,8 +166,11 @@ def _get_next_user_config(cluster_name: str, expected_traffic: float, spi_days:
                 "userConfig",
                 aws_provider
             )
+
+            # Load UserConfig from what was in SSM
             user_config = UserConfig(**stored_config_json)
 
+            # Now replace what was in SSM with any provided new values
             if expected_traffic is not None:
                 user_config.expectedTraffic = expected_traffic
             if spi_days is not None:
@@ -181,12 +183,11 @@ 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
-
             return user_config
 
         # Existing configuration doesn't exist, use defaults
         except ssm_ops.ParamDoesNotExist:
-            return UserConfig(MINIMUM_TRAFFIC, DEFAULT_SPI_DAYS, DEFAULT_HISTORY_DAYS, DEFAULT_REPLICAS, DEFAULT_S3_STORAGE_DAYS, None)
+            return UserConfig(expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list)
     # All of the parameters defined
     else:
         return UserConfig(expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list)
diff --git a/manage_arkime/core/user_config.py b/manage_arkime/core/user_config.py
index c4737ff..4ba336a 100644
--- a/manage_arkime/core/user_config.py
+++ b/manage_arkime/core/user_config.py
@@ -2,6 +2,8 @@
 import logging
 from typing import Dict
 
+from core.capacity_planning import (MINIMUM_TRAFFIC, DEFAULT_SPI_DAYS, DEFAULT_REPLICAS, DEFAULT_S3_STORAGE_DAYS, DEFAULT_HISTORY_DAYS)
+
 logger = logging.getLogger(__name__)
 
 @dataclass
@@ -13,6 +15,30 @@ class UserConfig:
     pcapDays: int
     viewerPrefixList: str = None
 
+    def __init__(self, expectedTraffic: float, spiDays: int, historyDays: int, replicas: int, pcapDays: int, viewerPrefixList: str = None):
+        if (expectedTraffic is None):
+            expectedTraffic = MINIMUM_TRAFFIC
+
+        if (spiDays is None):
+            spiDays = DEFAULT_SPI_DAYS
+
+        if (historyDays is None):
+            historyDays = DEFAULT_HISTORY_DAYS
+
+        if (replicas is None):
+            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
+
     """ Only process fields we still need, this allows us to ignore config no longer used """
     @classmethod
     def from_dict(cls, d):