From 7b69dab6220171ed98af951653f6d0df52f0b4ff Mon Sep 17 00:00:00 2001
From: Seth Grover <seth.d.grover@gmail.com>
Date: Thu, 9 May 2024 08:58:43 -0600
Subject: [PATCH] for idaholab/Malcolm#389, used tagged images

---
 scripts/install.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/scripts/install.py b/scripts/install.py
index a1ecff792..a884b0f11 100755
--- a/scripts/install.py
+++ b/scripts/install.py
@@ -78,6 +78,7 @@
     eprint,
     flatten,
     LoadFileIfJson,
+    remove_prefix,
     remove_suffix,
     run_process,
     same_file_or_dir,
@@ -2201,6 +2202,12 @@ def tweak_malcolm_runtime(self, malcolm_install_path):
                                 # whether or not to restart services automatically (on boot, etc.)
                                 line = f"{sectionIndents[currentSection] * 2}restart: {restartMode}"
 
+                            elif re.match(r'^\s*image\s*:.*$', line):
+                                # use architecture-specific images
+                                imageLineSpit = line.rstrip().split(":")
+                                imageLineSpit[-1] = imageLineSpit[-1].split("-", 1)[0] + args.imageArch
+                                line = ":".join(imageLineSpit)
+
                             elif (currentService == 'arkime') or (currentService == 'arkime-live'):
                                 # stuff specifically in the arkime section
                                 if re.match(r'^\s*-.+:/data/pcap(:.+)?\s*$', line):
@@ -3485,6 +3492,15 @@ def main():
         default=True,
         help="Enable dark mode for OpenSearch Dashboards",
     )
+    runtimeOptionsArgGroup.add_argument(
+        '--image-arch',
+        dest='imageArch',
+        required=False,
+        metavar='<amd64|arm64>',
+        type=str,
+        default=None,
+        help='Architecture for container image',
+    )
 
     authencOptionsArgGroup = parser.add_argument_group('Entryption and authentication options')
     authencOptionsArgGroup.add_argument(
@@ -4266,6 +4282,16 @@ def main():
     else:
         sys.tracebacklimit = 0
 
+    if args.imageArch is not None:
+        args.imageArch = (
+            '' if (args.imageArch.lower() == 'amd64') else ('-' + remove_suffix(args.imageArch.lower(), '-'))
+        )
+    else:
+        args.imageArch = ''
+        if rawPlatform := platform.machine().lower():
+            if rawPlatform == 'aarch64':
+                args.imageArch = '-arm64'
+
     orchMode = OrchestrationFramework.UNKNOWN
     if args.configFile and os.path.isfile(args.configFile):
         if not (