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.
 
 
 

263 lines
7.8 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. // RequiredFlags describes the flag requirements for a cbt command.
  46. type RequiredFlags uint
  47. const (
  48. // NoneRequired specifies that not flags are required.
  49. NoneRequired RequiredFlags = 0
  50. // ProjectRequired specifies that the -project flag is required.
  51. ProjectRequired RequiredFlags = 1 << iota
  52. // InstanceRequired specifies that the -instance flag is required.
  53. InstanceRequired
  54. // ProjectAndInstanceRequired specifies that both -project and -instance is required.
  55. ProjectAndInstanceRequired = ProjectRequired | InstanceRequired
  56. )
  57. // RegisterFlags registers a set of standard flags for this config.
  58. // It should be called before flag.Parse.
  59. func (c *Config) RegisterFlags() {
  60. flag.StringVar(&c.Project, "project", c.Project, "project ID, if unset uses gcloud configured project")
  61. flag.StringVar(&c.Instance, "instance", c.Instance, "Cloud Bigtable instance")
  62. flag.StringVar(&c.Creds, "creds", c.Creds, "if set, use application credentials in this file")
  63. flag.StringVar(&c.AdminEndpoint, "admin-endpoint", c.AdminEndpoint, "Override the admin api endpoint")
  64. flag.StringVar(&c.DataEndpoint, "data-endpoint", c.DataEndpoint, "Override the data api endpoint")
  65. flag.StringVar(&c.CertFile, "cert-file", c.CertFile, "Override the TLS certificates file")
  66. flag.StringVar(&c.UserAgent, "user-agent", c.UserAgent, "Override the user agent string")
  67. }
  68. // CheckFlags checks that the required config values are set.
  69. func (c *Config) CheckFlags(required RequiredFlags) error {
  70. var missing []string
  71. if c.CertFile != "" {
  72. b, err := ioutil.ReadFile(c.CertFile)
  73. if err != nil {
  74. return fmt.Errorf("Failed to load certificates from %s: %v", c.CertFile, err)
  75. }
  76. cp := x509.NewCertPool()
  77. if !cp.AppendCertsFromPEM(b) {
  78. return fmt.Errorf("Failed to append certificates from %s", c.CertFile)
  79. }
  80. c.TLSCreds = credentials.NewTLS(&tls.Config{RootCAs: cp})
  81. }
  82. if required != NoneRequired {
  83. c.SetFromGcloud()
  84. }
  85. if required&ProjectRequired != 0 && c.Project == "" {
  86. missing = append(missing, "-project")
  87. }
  88. if required&InstanceRequired != 0 && c.Instance == "" {
  89. missing = append(missing, "-instance")
  90. }
  91. if len(missing) > 0 {
  92. return fmt.Errorf("Missing %s", strings.Join(missing, " and "))
  93. }
  94. return nil
  95. }
  96. // Filename returns the filename consulted for standard configuration.
  97. func Filename() string {
  98. // TODO(dsymonds): Might need tweaking for Windows.
  99. return filepath.Join(os.Getenv("HOME"), ".cbtrc")
  100. }
  101. // Load loads a .cbtrc file.
  102. // If the file is not present, an empty config is returned.
  103. func Load() (*Config, error) {
  104. filename := Filename()
  105. data, err := ioutil.ReadFile(filename)
  106. if err != nil {
  107. // silent fail if the file isn't there
  108. if os.IsNotExist(err) {
  109. return &Config{}, nil
  110. }
  111. return nil, fmt.Errorf("Reading %s: %v", filename, err)
  112. }
  113. c := new(Config)
  114. s := bufio.NewScanner(bytes.NewReader(data))
  115. for s.Scan() {
  116. line := s.Text()
  117. i := strings.Index(line, "=")
  118. if i < 0 {
  119. return nil, fmt.Errorf("Bad line in %s: %q", filename, line)
  120. }
  121. key, val := strings.TrimSpace(line[:i]), strings.TrimSpace(line[i+1:])
  122. switch key {
  123. default:
  124. return nil, fmt.Errorf("Unknown key in %s: %q", filename, key)
  125. case "project":
  126. c.Project = val
  127. case "instance":
  128. c.Instance = val
  129. case "creds":
  130. c.Creds = val
  131. case "admin-endpoint":
  132. c.AdminEndpoint = val
  133. case "data-endpoint":
  134. c.DataEndpoint = val
  135. case "cert-file":
  136. c.CertFile = val
  137. case "user-agent":
  138. c.UserAgent = val
  139. }
  140. }
  141. return c, s.Err()
  142. }
  143. // GcloudCredential holds gcloud credential information.
  144. type GcloudCredential struct {
  145. AccessToken string `json:"access_token"`
  146. Expiry time.Time `json:"token_expiry"`
  147. }
  148. // Token creates an oauth2 token using gcloud credentials.
  149. func (cred *GcloudCredential) Token() *oauth2.Token {
  150. return &oauth2.Token{AccessToken: cred.AccessToken, TokenType: "Bearer", Expiry: cred.Expiry}
  151. }
  152. // GcloudConfig holds gcloud configuration values.
  153. type GcloudConfig struct {
  154. Configuration struct {
  155. Properties struct {
  156. Core struct {
  157. Project string `json:"project"`
  158. } `json:"core"`
  159. } `json:"properties"`
  160. } `json:"configuration"`
  161. Credential GcloudCredential `json:"credential"`
  162. }
  163. // GcloudCmdTokenSource holds the comamnd arguments. It is only intended to be set by the program.
  164. // TODO(deklerk) Can this be unexported?
  165. type GcloudCmdTokenSource struct {
  166. Command string
  167. Args []string
  168. }
  169. // Token implements the oauth2.TokenSource interface
  170. func (g *GcloudCmdTokenSource) Token() (*oauth2.Token, error) {
  171. gcloudConfig, err := LoadGcloudConfig(g.Command, g.Args)
  172. if err != nil {
  173. return nil, err
  174. }
  175. return gcloudConfig.Credential.Token(), nil
  176. }
  177. // LoadGcloudConfig retrieves the gcloud configuration values we need use via the
  178. // 'config-helper' command
  179. func LoadGcloudConfig(gcloudCmd string, gcloudCmdArgs []string) (*GcloudConfig, error) {
  180. out, err := exec.Command(gcloudCmd, gcloudCmdArgs...).Output()
  181. if err != nil {
  182. return nil, fmt.Errorf("Could not retrieve gcloud configuration")
  183. }
  184. var gcloudConfig GcloudConfig
  185. if err := json.Unmarshal(out, &gcloudConfig); err != nil {
  186. return nil, fmt.Errorf("Could not parse gcloud configuration")
  187. }
  188. return &gcloudConfig, nil
  189. }
  190. // SetFromGcloud retrieves and sets any missing config values from the gcloud
  191. // configuration if possible possible
  192. func (c *Config) SetFromGcloud() error {
  193. if c.Creds == "" {
  194. c.Creds = os.Getenv("GOOGLE_APPLICATION_CREDENTIALS")
  195. if c.Creds == "" {
  196. log.Printf("-creds flag unset, will use gcloud credential")
  197. }
  198. } else {
  199. os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", c.Creds)
  200. }
  201. if c.Project == "" {
  202. log.Printf("-project flag unset, will use gcloud active project")
  203. }
  204. if c.Creds != "" && c.Project != "" {
  205. return nil
  206. }
  207. gcloudCmd := "gcloud"
  208. if runtime.GOOS == "windows" {
  209. gcloudCmd = gcloudCmd + ".cmd"
  210. }
  211. gcloudCmdArgs := []string{"config", "config-helper",
  212. "--format=json(configuration.properties.core.project,credential)"}
  213. gcloudConfig, err := LoadGcloudConfig(gcloudCmd, gcloudCmdArgs)
  214. if err != nil {
  215. return err
  216. }
  217. if c.Project == "" && gcloudConfig.Configuration.Properties.Core.Project != "" {
  218. log.Printf("gcloud active project is \"%s\"",
  219. gcloudConfig.Configuration.Properties.Core.Project)
  220. c.Project = gcloudConfig.Configuration.Properties.Core.Project
  221. }
  222. if c.Creds == "" {
  223. c.TokenSource = oauth2.ReuseTokenSource(
  224. gcloudConfig.Credential.Token(),
  225. &GcloudCmdTokenSource{Command: gcloudCmd, Args: gcloudCmdArgs})
  226. }
  227. return nil
  228. }