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// Get returns the value associated with the key, and whether it exists.
37// It's equivalent to the os.LookupEnv function, but with this copy of the
38// Environment.
39func (e *Environment) Get(key string) (string, bool) {
40	for _, env := range *e {
41		if k, v, ok := decodeKeyValue(env); ok && k == key {
42			return v, true
43		}
44	}
45	return "", false
46}
47
48// Get returns the int value associated with the key, and whether it exists
49// and is a valid int.
50func (e *Environment) GetInt(key string) (int, bool) {
51	if v, ok := e.Get(key); ok {
52		if i, err := strconv.Atoi(v); err == nil {
53			return i, true
54		}
55	}
56	return 0, false
57}
58
59// Set sets the value associated with the key, overwriting the current value
60// if it exists.
61func (e *Environment) Set(key, value string) {
62	e.Unset(key)
63	*e = append(*e, key+"="+value)
64}
65
66// Unset removes the specified keys from the Environment.
67func (e *Environment) Unset(keys ...string) {
68	out := (*e)[:0]
69	for _, env := range *e {
70		if key, _, ok := decodeKeyValue(env); ok && inList(key, keys) {
71			continue
72		}
73		out = append(out, env)
74	}
75	*e = out
76}
77
78// UnsetWithPrefix removes all keys that start with prefix.
79func (e *Environment) UnsetWithPrefix(prefix string) {
80	out := (*e)[:0]
81	for _, env := range *e {
82		if key, _, ok := decodeKeyValue(env); ok && strings.HasPrefix(key, prefix) {
83			continue
84		}
85		out = append(out, env)
86	}
87	*e = out
88}
89
90// Allow removes all keys that are not present in the input list
91func (e *Environment) Allow(keys ...string) {
92	out := (*e)[:0]
93	for _, env := range *e {
94		if key, _, ok := decodeKeyValue(env); ok && inList(key, keys) {
95			out = append(out, env)
96		}
97	}
98	*e = out
99}
100
101// Environ returns the []string required for exec.Cmd.Env
102func (e *Environment) Environ() []string {
103	return []string(*e)
104}
105
106// Copy returns a copy of the Environment so that independent changes may be made.
107func (e *Environment) Copy() *Environment {
108	ret := Environment(make([]string, len(*e)))
109	for i, v := range *e {
110		ret[i] = v
111	}
112	return &ret
113}
114
115// IsTrue returns whether an environment variable is set to a positive value (1,y,yes,on,true)
116func (e *Environment) IsEnvTrue(key string) bool {
117	if value, ok := e.Get(key); ok {
118		return value == "1" || value == "y" || value == "yes" || value == "on" || value == "true"
119	}
120	return false
121}
122
123// IsFalse returns whether an environment variable is set to a negative value (0,n,no,off,false)
124func (e *Environment) IsFalse(key string) bool {
125	if value, ok := e.Get(key); ok {
126		return value == "0" || value == "n" || value == "no" || value == "off" || value == "false"
127	}
128	return false
129}
130
131// AppendFromKati reads a shell script written by Kati that exports or unsets
132// environment variables, and applies those to the local Environment.
133func (e *Environment) AppendFromKati(filename string) error {
134	file, err := os.Open(filename)
135	if err != nil {
136		return err
137	}
138	defer file.Close()
139
140	return e.appendFromKati(file)
141}
142
143func (e *Environment) appendFromKati(reader io.Reader) error {
144	scanner := bufio.NewScanner(reader)
145	for scanner.Scan() {
146		text := strings.TrimSpace(scanner.Text())
147
148		if len(text) == 0 || text[0] == '#' {
149			continue
150		}
151
152		cmd := strings.SplitN(text, " ", 2)
153		if len(cmd) != 2 {
154			return fmt.Errorf("Unknown kati environment line: %q", text)
155		}
156
157		if cmd[0] == "unset" {
158			str, ok := singleUnquote(cmd[1])
159			if !ok {
160				return fmt.Errorf("Failed to unquote kati line: %q", text)
161			}
162			e.Unset(str)
163		} else if cmd[0] == "export" {
164			key, value, ok := decodeKeyValue(cmd[1])
165			if !ok {
166				return fmt.Errorf("Failed to parse export: %v", cmd)
167			}
168
169			key, ok = singleUnquote(key)
170			if !ok {
171				return fmt.Errorf("Failed to unquote kati line: %q", text)
172			}
173			value, ok = singleUnquote(value)
174			if !ok {
175				return fmt.Errorf("Failed to unquote kati line: %q", text)
176			}
177
178			e.Set(key, value)
179		} else {
180			return fmt.Errorf("Unknown kati environment command: %q", text)
181		}
182	}
183	if err := scanner.Err(); err != nil {
184		return err
185	}
186	return nil
187}
188