From 5dd0187261504fd263f9edca0fb0152fff3f957b Mon Sep 17 00:00:00 2001
From: Ian Mckay <contact@ianmckay.com.au>
Date: Sun, 9 Feb 2025 10:42:43 +1100
Subject: [PATCH] Add extra flag to override map file

---
 README.md              |  2 ++
 iamlivecore/logger.go  | 20 ++++++++++++++++----
 iamlivecore/service.go |  6 ++++++
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 6b9eb808..7049e43f 100644
--- a/README.md
+++ b/README.md
@@ -78,6 +78,8 @@ You can optionally also include the following arguments to the `iamlive` command
 
 **--account-id:** the AWS account ID to use in policy outputs within proxy mode (_default: 123456789012 unless detected_) (_AWS only_)
 
+**--override-aws-map:** overrides the embedded AWS mapping JSON file with the filepath provided (_AWS only_)
+
 **--debug:** dumps associated HTTP requests when set in proxy mode (_default: false_)
 
 _Basic Example (CSM Mode)_
diff --git a/iamlivecore/logger.go b/iamlivecore/logger.go
index e7b776e2..4df4d479 100644
--- a/iamlivecore/logger.go
+++ b/iamlivecore/logger.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"log"
 	"net/url"
+	"os"
 	"reflect"
 	"regexp"
 	"sort"
@@ -88,11 +89,22 @@ type AzureIAMPolicy struct {
 
 func loadMaps() {
 	if *providerFlag == "aws" {
-		err := json.Unmarshal(bIAMMap, &iamMap)
-		if err != nil {
-			log.Fatal(err)
+		if *overrideAwsMapFlag != "" {
+			bIAMMap, err := os.ReadFile(*overrideAwsMapFlag)
+			if err != nil {
+				log.Fatal(err)
+			}
+			err = json.Unmarshal(bIAMMap, &iamMap)
+			if err != nil {
+				log.Fatal(err)
+			}
+		} else {
+			err := json.Unmarshal(bIAMMap, &iamMap)
+			if err != nil {
+				log.Fatal(err)
+			}
 		}
-		err = json.Unmarshal(bIAMSAR, &iamDef)
+		err := json.Unmarshal(bIAMSAR, &iamDef)
 		if err != nil {
 			panic(err)
 		}
diff --git a/iamlivecore/service.go b/iamlivecore/service.go
index ebad0551..c1c0d38d 100644
--- a/iamlivecore/service.go
+++ b/iamlivecore/service.go
@@ -27,6 +27,7 @@ var caBundleFlag *string
 var caKeyFlag *string
 var accountIDFlag *string
 var backgroundFlag *bool
+var overrideAwsMapFlag *string
 var debugFlag *bool
 var forceWildcardResourceFlag *bool
 var cpuProfileFlag = flag.String("cpu-profile", "", "write a CPU profile to this file (for performance testing purposes)")
@@ -48,6 +49,7 @@ func parseConfig() {
 	caKey := "~/.iamlive/ca.key"
 	accountID := ""
 	background := false
+	overrideAwsMap := ""
 	debug := false
 	forceWildcardResource := false
 	csmPort := 31000
@@ -99,6 +101,9 @@ func parseConfig() {
 			if cfg.Section("").HasKey("background") {
 				background, _ = cfg.Section("").Key("background").Bool()
 			}
+			if cfg.Section("").HasKey("override-aws-map") {
+				overrideAwsMap = cfg.Section("").Key("override-aws-map").String()
+			}
 			if cfg.Section("").HasKey("debug") {
 				debug, _ = cfg.Section("").Key("debug").Bool()
 			}
@@ -126,6 +131,7 @@ func parseConfig() {
 	caKeyFlag = flag.String("ca-key", caKey, "the CA certificate key to use for proxy mode")
 	accountIDFlag = flag.String("account-id", accountID, "the AWS account ID to use in policy outputs within proxy mode")
 	backgroundFlag = flag.Bool("background", background, "when set, the process will return the current PID and run in the background without output")
+	overrideAwsMapFlag = flag.String("override-aws-map", overrideAwsMap, "overrides the embedded AWS mapping JSON file with the filepath provided")
 	debugFlag = flag.Bool("debug", debug, "dumps associated HTTP requests when set in proxy mode")
 	forceWildcardResourceFlag = flag.Bool("force-wildcard-resource", forceWildcardResource, "when set, the Resource will always be a wildcard")
 	csmPortFlag = flag.Int("csm-port", csmPort, "port to listen on for CSM")