environment.go 5.5 KB

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