1// Copyright 2017 Google Inc. All rights reserved. 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 15package build 16 17import ( 18 "bufio" 19 "fmt" 20 "io" 21 "os" 22 "strconv" 23 "strings" 24) 25 26// Environment adds a number of useful manipulation functions to the list of 27// strings returned by os.Environ() and used in exec.Cmd.Env. 28type Environment []string 29 30// OsEnvironment wraps the current environment returned by os.Environ() 31func OsEnvironment() *Environment { 32 env := Environment(os.Environ()) 33 return &env 34} 35 36// Returns a copy of the environment as a map[string]string. 37func (e *Environment) AsMap() map[string]string { 38 result := make(map[string]string) 39 40 for _, envVar := range *e { 41 if k, v, ok := decodeKeyValue(envVar); ok { 42 result[k] = v 43 } 44 } 45 46 return result 47} 48 49// Get returns the value associated with the key, and whether it exists. 50// It's equivalent to the os.LookupEnv function, but with this copy of the 51// Environment. 52func (e *Environment) Get(key string) (string, bool) { 53 for _, envVar := range *e { 54 if k, v, ok := decodeKeyValue(envVar); ok && k == key { 55 return v, true 56 } 57 } 58 return "", false 59} 60 61// Get returns the int value associated with the key, and whether it exists 62// and is a valid int. 63func (e *Environment) GetInt(key string) (int, bool) { 64 if v, ok := e.Get(key); ok { 65 if i, err := strconv.Atoi(v); err == nil { 66 return i, true 67 } 68 } 69 return 0, false 70} 71 72// Set sets the value associated with the key, overwriting the current value 73// if it exists. 74func (e *Environment) Set(key, value string) { 75 e.Unset(key) 76 *e = append(*e, key+"="+value) 77} 78 79// Unset removes the specified keys from the Environment. 80func (e *Environment) Unset(keys ...string) { 81 newEnv := (*e)[:0] 82 for _, envVar := range *e { 83 if key, _, ok := decodeKeyValue(envVar); ok && inList(key, keys) { 84 // Delete this key. 85 continue 86 } 87 newEnv = append(newEnv, envVar) 88 } 89 *e = newEnv 90} 91 92// UnsetWithPrefix removes all keys that start with prefix. 93func (e *Environment) UnsetWithPrefix(prefix string) { 94 newEnv := (*e)[:0] 95 for _, envVar := range *e { 96 if key, _, ok := decodeKeyValue(envVar); ok && strings.HasPrefix(key, prefix) { 97 // Delete this key. 98 continue 99 } 100 newEnv = append(newEnv, envVar) 101 } 102 *e = newEnv 103} 104 105// Allow removes all keys that are not present in the input list 106func (e *Environment) Allow(keys ...string) { 107 newEnv := (*e)[:0] 108 for _, envVar := range *e { 109 if key, _, ok := decodeKeyValue(envVar); ok && inList(key, keys) { 110 // Keep this key. 111 newEnv = append(newEnv, envVar) 112 } 113 } 114 *e = newEnv 115} 116 117// Environ returns the []string required for exec.Cmd.Env 118func (e *Environment) Environ() []string { 119 return []string(*e) 120} 121 122// Copy returns a copy of the Environment so that independent changes may be made. 123func (e *Environment) Copy() *Environment { 124 envCopy := Environment(make([]string, len(*e))) 125 for i, envVar := range *e { 126 envCopy[i] = envVar 127 } 128 return &envCopy 129} 130 131// IsTrue returns whether an environment variable is set to a positive value (1,y,yes,on,true) 132func (e *Environment) IsEnvTrue(key string) bool { 133 if value, ok := e.Get(key); ok { 134 return value == "1" || value == "y" || value == "yes" || value == "on" || value == "true" 135 } 136 return false 137} 138 139// IsFalse returns whether an environment variable is set to a negative value (0,n,no,off,false) 140func (e *Environment) IsFalse(key string) bool { 141 if value, ok := e.Get(key); ok { 142 return value == "0" || value == "n" || value == "no" || value == "off" || value == "false" 143 } 144 return false 145} 146 147// AppendFromKati reads a shell script written by Kati that exports or unsets 148// environment variables, and applies those to the local Environment. 149func (e *Environment) AppendFromKati(filename string) error { 150 file, err := os.Open(filename) 151 if err != nil { 152 return err 153 } 154 defer file.Close() 155 156 return e.appendFromKati(file) 157} 158 159// Helper function for AppendFromKati. Accepts an io.Reader to make testing easier. 160func (e *Environment) appendFromKati(reader io.Reader) error { 161 scanner := bufio.NewScanner(reader) 162 for scanner.Scan() { 163 text := strings.TrimSpace(scanner.Text()) 164 165 if len(text) == 0 || text[0] == '#' { 166 // Skip blank lines and comments. 167 continue 168 } 169 170 // We expect two space-delimited strings, like: 171 // unset 'HOME' 172 // export 'BEST_PIZZA_CITY'='NYC' 173 cmd := strings.SplitN(text, " ", 2) 174 if len(cmd) != 2 { 175 return fmt.Errorf("Unknown kati environment line: %q", text) 176 } 177 178 if cmd[0] == "unset" { 179 str, ok := singleUnquote(cmd[1]) 180 if !ok { 181 return fmt.Errorf("Failed to unquote kati line: %q", text) 182 } 183 184 // Actually unset it. 185 e.Unset(str) 186 } else if cmd[0] == "export" { 187 key, value, ok := decodeKeyValue(cmd[1]) 188 if !ok { 189 return fmt.Errorf("Failed to parse export: %v", cmd) 190 } 191 192 key, ok = singleUnquote(key) 193 if !ok { 194 return fmt.Errorf("Failed to unquote kati line: %q", text) 195 } 196 value, ok = singleUnquote(value) 197 if !ok { 198 return fmt.Errorf("Failed to unquote kati line: %q", text) 199 } 200 201 // Actually set it. 202 e.Set(key, value) 203 } else { 204 return fmt.Errorf("Unknown kati environment command: %q", text) 205 } 206 } 207 if err := scanner.Err(); err != nil { 208 return err 209 } 210 return nil 211} 212