github.com/w3security/driftctl@v0.38.0/pkg/envproxy/env_proxy.go (about) 1 package envproxy 2 3 import ( 4 "os" 5 "strings" 6 ) 7 8 type EnvProxy struct { 9 fromPrefix string 10 toPrefix string 11 defaultEnv map[string]string 12 } 13 14 func NewEnvProxy(fromPrefix, toPrefix string) *EnvProxy { 15 envMap := map[string]string{} 16 for _, variable := range os.Environ() { 17 tmp := strings.SplitN(variable, "=", 2) 18 envMap[tmp[0]] = tmp[1] 19 } 20 return &EnvProxy{ 21 fromPrefix: fromPrefix, 22 toPrefix: toPrefix, 23 defaultEnv: envMap, 24 } 25 } 26 27 func (s *EnvProxy) Apply() { 28 if s.fromPrefix == "" || s.toPrefix == "" { 29 return 30 } 31 for key, value := range s.defaultEnv { 32 if strings.HasPrefix(key, s.fromPrefix) { 33 key = strings.Replace(key, s.fromPrefix, s.toPrefix, 1) 34 os.Setenv(key, value) 35 } 36 } 37 } 38 39 func (s *EnvProxy) Restore() { 40 if s.fromPrefix == "" || s.toPrefix == "" { 41 return 42 } 43 for key, value := range s.defaultEnv { 44 if strings.HasPrefix(key, s.fromPrefix) { 45 key = strings.Replace(key, s.fromPrefix, s.toPrefix, 1) 46 value = s.defaultEnv[key] 47 } 48 os.Setenv(key, value) 49 } 50 }