• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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