You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

253 lines
7.2 KiB

  1. /*
  2. Copyright 2015 Google LLC
  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. http://www.apache.org/licenses/LICENSE-2.0
  7. Unless required by applicable law or agreed to in writing, software
  8. distributed under the License is distributed on an "AS IS" BASIS,
  9. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  10. See the License for the specific language governing permissions and
  11. limitations under the License.
  12. */
  13. // Package cbtconfig encapsulates common code for reading configuration from .cbtrc and gcloud.
  14. package cbtconfig
  15. import (
  16. "bufio"
  17. "bytes"
  18. "crypto/tls"
  19. "crypto/x509"
  20. "encoding/json"
  21. "flag"
  22. "fmt"
  23. "io/ioutil"
  24. "log"
  25. "os"
  26. "os/exec"
  27. "path/filepath"
  28. "runtime"
  29. "strings"
  30. "time"
  31. "golang.org/x/oauth2"
  32. "google.golang.org/grpc/credentials"
  33. )
  34. // Config represents a configuration.
  35. type Config struct {
  36. Project, Instance string // required
  37. Creds string // optional
  38. AdminEndpoint string // optional
  39. DataEndpoint string // optional
  40. CertFile string // optional
  41. UserAgent string // optional
  42. TokenSource oauth2.TokenSource // derived
  43. TLSCreds credentials.TransportCredentials // derived
  44. }
  45. type RequiredFlags uint
  46. const NoneRequired RequiredFlags = 0
  47. const (
  48. ProjectRequired RequiredFlags = 1 << iota
  49. InstanceRequired
  50. )
  51. const ProjectAndInstanceRequired RequiredFlags = ProjectRequired | InstanceRequired
  52. // RegisterFlags registers a set of standard flags for this config.
  53. // It should be called before flag.Parse.
  54. func (c *Config) RegisterFlags() {
  55. flag.StringVar(&c.Project, "project", c.Project, "project ID, if unset uses gcloud configured project")
  56. flag.StringVar(&c.Instance, "instance", c.Instance, "Cloud Bigtable instance")
  57. flag.StringVar(&c.Creds, "creds", c.Creds, "if set, use application credentials in this file")
  58. flag.StringVar(&c.AdminEndpoint, "admin-endpoint", c.AdminEndpoint, "Override the admin api endpoint")
  59. flag.StringVar(&c.DataEndpoint, "data-endpoint", c.DataEndpoint, "Override the data api endpoint")
  60. flag.StringVar(&c.CertFile, "cert-file", c.CertFile, "Override the TLS certificates file")
  61. flag.StringVar(&c.UserAgent, "user-agent", c.UserAgent, "Override the user agent string")
  62. }
  63. // CheckFlags checks that the required config values are set.
  64. func (c *Config) CheckFlags(required RequiredFlags) error {
  65. var missing []string
  66. if c.CertFile != "" {
  67. b, err := ioutil.ReadFile(c.CertFile)
  68. if err != nil {
  69. return fmt.Errorf("Failed to load certificates from %s: %v", c.CertFile, err)
  70. }
  71. cp := x509.NewCertPool()
  72. if !cp.AppendCertsFromPEM(b) {
  73. return fmt.Errorf("Failed to append certificates from %s", c.CertFile)
  74. }
  75. c.TLSCreds = credentials.NewTLS(&tls.Config{RootCAs: cp})
  76. }
  77. if required != NoneRequired {
  78. c.SetFromGcloud()
  79. }
  80. if required&ProjectRequired != 0 && c.Project == "" {
  81. missing = append(missing, "-project")
  82. }
  83. if required&InstanceRequired != 0 && c.Instance == "" {
  84. missing = append(missing, "-instance")
  85. }
  86. if len(missing) > 0 {
  87. return fmt.Errorf("Missing %s", strings.Join(missing, " and "))
  88. }
  89. return nil
  90. }
  91. // Filename returns the filename consulted for standard configuration.
  92. func Filename() string {
  93. // TODO(dsymonds): Might need tweaking for Windows.
  94. return filepath.Join(os.Getenv("HOME"), ".cbtrc")
  95. }
  96. // Load loads a .cbtrc file.
  97. // If the file is not present, an empty config is returned.
  98. func Load() (*Config, error) {
  99. filename := Filename()
  100. data, err := ioutil.ReadFile(filename)
  101. if err != nil {
  102. // silent fail if the file isn't there
  103. if os.IsNotExist(err) {
  104. return &Config{}, nil
  105. }
  106. return nil, fmt.Errorf("Reading %s: %v", filename, err)
  107. }
  108. c := new(Config)
  109. s := bufio.NewScanner(bytes.NewReader(data))
  110. for s.Scan() {
  111. line := s.Text()
  112. i := strings.Index(line, "=")
  113. if i < 0 {
  114. return nil, fmt.Errorf("Bad line in %s: %q", filename, line)
  115. }
  116. key, val := strings.TrimSpace(line[:i]), strings.TrimSpace(line[i+1:])
  117. switch key {
  118. default:
  119. return nil, fmt.Errorf("Unknown key in %s: %q", filename, key)
  120. case "project":
  121. c.Project = val
  122. case "instance":
  123. c.Instance = val
  124. case "creds":
  125. c.Creds = val
  126. case "admin-endpoint":
  127. c.AdminEndpoint = val
  128. case "data-endpoint":
  129. c.DataEndpoint = val
  130. case "cert-file":
  131. c.CertFile = val
  132. case "user-agent":
  133. c.UserAgent = val
  134. }
  135. }
  136. return c, s.Err()
  137. }
  138. type GcloudCredential struct {
  139. AccessToken string `json:"access_token"`
  140. Expiry time.Time `json:"token_expiry"`
  141. }
  142. func (cred *GcloudCredential) Token() *oauth2.Token {
  143. return &oauth2.Token{AccessToken: cred.AccessToken, TokenType: "Bearer", Expiry: cred.Expiry}
  144. }
  145. type GcloudConfig struct {
  146. Configuration struct {
  147. Properties struct {
  148. Core struct {
  149. Project string `json:"project"`
  150. } `json:"core"`
  151. } `json:"properties"`
  152. } `json:"configuration"`
  153. Credential GcloudCredential `json:"credential"`
  154. }
  155. type GcloudCmdTokenSource struct {
  156. Command string
  157. Args []string
  158. }
  159. // Token implements the oauth2.TokenSource interface
  160. func (g *GcloudCmdTokenSource) Token() (*oauth2.Token, error) {
  161. gcloudConfig, err := LoadGcloudConfig(g.Command, g.Args)
  162. if err != nil {
  163. return nil, err
  164. }
  165. return gcloudConfig.Credential.Token(), nil
  166. }
  167. // LoadGcloudConfig retrieves the gcloud configuration values we need use via the
  168. // 'config-helper' command
  169. func LoadGcloudConfig(gcloudCmd string, gcloudCmdArgs []string) (*GcloudConfig, error) {
  170. out, err := exec.Command(gcloudCmd, gcloudCmdArgs...).Output()
  171. if err != nil {
  172. return nil, fmt.Errorf("Could not retrieve gcloud configuration")
  173. }
  174. var gcloudConfig GcloudConfig
  175. if err := json.Unmarshal(out, &gcloudConfig); err != nil {
  176. return nil, fmt.Errorf("Could not parse gcloud configuration")
  177. }
  178. return &gcloudConfig, nil
  179. }
  180. // SetFromGcloud retrieves and sets any missing config values from the gcloud
  181. // configuration if possible possible
  182. func (c *Config) SetFromGcloud() error {
  183. if c.Creds == "" {
  184. c.Creds = os.Getenv("GOOGLE_APPLICATION_CREDENTIALS")
  185. if c.Creds == "" {
  186. log.Printf("-creds flag unset, will use gcloud credential")
  187. }
  188. } else {
  189. os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", c.Creds)
  190. }
  191. if c.Project == "" {
  192. log.Printf("-project flag unset, will use gcloud active project")
  193. }
  194. if c.Creds != "" && c.Project != "" {
  195. return nil
  196. }
  197. gcloudCmd := "gcloud"
  198. if runtime.GOOS == "windows" {
  199. gcloudCmd = gcloudCmd + ".cmd"
  200. }
  201. gcloudCmdArgs := []string{"config", "config-helper",
  202. "--format=json(configuration.properties.core.project,credential)"}
  203. gcloudConfig, err := LoadGcloudConfig(gcloudCmd, gcloudCmdArgs)
  204. if err != nil {
  205. return err
  206. }
  207. if c.Project == "" && gcloudConfig.Configuration.Properties.Core.Project != "" {
  208. log.Printf("gcloud active project is \"%s\"",
  209. gcloudConfig.Configuration.Properties.Core.Project)
  210. c.Project = gcloudConfig.Configuration.Properties.Core.Project
  211. }
  212. if c.Creds == "" {
  213. c.TokenSource = oauth2.ReuseTokenSource(
  214. gcloudConfig.Credential.Token(),
  215. &GcloudCmdTokenSource{Command: gcloudCmd, Args: gcloudCmdArgs})
  216. }
  217. return nil
  218. }