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  }