github.com/alibaba/sealer@v0.8.6-0.20220430115802-37a2bdaa8173/pkg/env/env.go (about)

     1  // Copyright © 2021 Alibaba Group Holding Ltd.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package env
    16  
    17  import (
    18  	"fmt"
    19  	"html/template"
    20  	"os"
    21  	"path/filepath"
    22  	"strings"
    23  
    24  	"github.com/Masterminds/sprig/v3"
    25  
    26  	v2 "github.com/alibaba/sealer/types/api/v2"
    27  )
    28  
    29  const templateSuffix = ".tmpl"
    30  
    31  type Interface interface {
    32  	PreProcessor
    33  	// WrapperShell :If host already set env like DATADISK=/data
    34  	// This function add env to the shell, like:
    35  	// Input shell: cat /etc/hosts
    36  	// Output shell: DATADISK=/data cat /etc/hosts
    37  	// So that you can get env values in you shell script
    38  	WrapperShell(host, shell string) string
    39  	// RenderAll :render env to all the files in dir
    40  	RenderAll(host, dir string) error
    41  }
    42  
    43  type processor struct {
    44  	*v2.Cluster
    45  }
    46  
    47  func NewEnvProcessor(cluster *v2.Cluster) Interface {
    48  	return &processor{cluster}
    49  }
    50  
    51  func (p *processor) WrapperShell(host, shell string) string {
    52  	var env string
    53  	for k, v := range p.getHostEnv(host) {
    54  		switch value := v.(type) {
    55  		case []string:
    56  			env = fmt.Sprintf("%s%s=(%s) ", env, k, strings.Join(value, " "))
    57  		case string:
    58  			env = fmt.Sprintf("%s%s=%s ", env, k, value)
    59  		}
    60  	}
    61  	if env == "" {
    62  		return shell
    63  	}
    64  	return fmt.Sprintf("%s && %s", env, shell)
    65  }
    66  
    67  func (p *processor) RenderAll(host, dir string) error {
    68  	return filepath.Walk(dir, func(path string, info os.FileInfo, errIn error) error {
    69  		if errIn != nil {
    70  			return errIn
    71  		}
    72  		if info.IsDir() || !strings.HasSuffix(info.Name(), templateSuffix) {
    73  			return nil
    74  		}
    75  		writer, err := os.OpenFile(strings.TrimSuffix(path, templateSuffix), os.O_CREATE|os.O_RDWR, os.ModePerm)
    76  		if err != nil {
    77  			return fmt.Errorf("failed to open file [%s] when render env: %v", path, err)
    78  		}
    79  		defer func() {
    80  			_ = writer.Close()
    81  		}()
    82  		t, err := template.New(info.Name()).Funcs(sprig.FuncMap()).ParseFiles(path)
    83  		if err != nil {
    84  			return fmt.Errorf("failed to create template: %s %v", path, err)
    85  		}
    86  		if err := t.Execute(writer, p.getHostEnv(host)); err != nil {
    87  			return fmt.Errorf("failed to render env template: %s %v", path, err)
    88  		}
    89  		return nil
    90  	})
    91  }
    92  
    93  func mergeList(hostEnv, globalEnv map[string]interface{}) map[string]interface{} {
    94  	if len(hostEnv) == 0 {
    95  		return globalEnv
    96  	}
    97  	for globalEnvKey, globalEnvValue := range globalEnv {
    98  		if _, ok := hostEnv[globalEnvKey]; ok {
    99  			continue
   100  		}
   101  		hostEnv[globalEnvKey] = globalEnvValue
   102  	}
   103  	return hostEnv
   104  }
   105  
   106  // Merge the host ENV and global env, the host env will overwrite cluster.Spec.Env
   107  func (p *processor) getHostEnv(hostIP string) (env map[string]interface{}) {
   108  	hostEnv, globalEnv := map[string]interface{}{}, ConvertEnv(p.Spec.Env)
   109  
   110  	for _, host := range p.Spec.Hosts {
   111  		for _, ip := range host.IPS {
   112  			if ip == hostIP {
   113  				hostEnv = ConvertEnv(host.Env)
   114  			}
   115  		}
   116  	}
   117  	return mergeList(hostEnv, globalEnv)
   118  }
   119  
   120  // ConvertEnv []string to map[string]interface{}, example [IP=127.0.0.1,IP=192.160.0.2,Key=value] will convert to {IP:[127.0.0.1,192.168.0.2],key:value}
   121  func ConvertEnv(envList []string) (env map[string]interface{}) {
   122  	temp := make(map[string][]string)
   123  	env = make(map[string]interface{})
   124  
   125  	for _, e := range envList {
   126  		var kv []string
   127  		if kv = strings.SplitN(e, "=", 2); len(kv) != 2 {
   128  			continue
   129  		}
   130  		temp[kv[0]] = append(temp[kv[0]], strings.Split(kv[1], ";")...)
   131  	}
   132  
   133  	for k, v := range temp {
   134  		if len(v) > 1 {
   135  			env[k] = v
   136  			continue
   137  		}
   138  		if len(v) == 1 {
   139  			env[k] = v[0]
   140  		}
   141  	}
   142  
   143  	return
   144  }