github.com/grailbio/base@v0.0.11/cmd/grail-role-group/googleclient/client.go (about)

     1  package googleclient
     2  
     3  import (
     4  	"context"
     5  	"crypto/rand"
     6  	"encoding/hex"
     7  	"encoding/json"
     8  	"fmt"
     9  	"net"
    10  	"net/http"
    11  	"os"
    12  	"path/filepath"
    13  	"strings"
    14  	"sync"
    15  
    16  	"github.com/grailbio/base/cmdutil"
    17  	"github.com/grailbio/base/web/webutil"
    18  	"golang.org/x/oauth2"
    19  	"golang.org/x/oauth2/google"
    20  	goauth2 "google.golang.org/api/oauth2/v2"
    21  	"v.io/x/lib/vlog"
    22  )
    23  
    24  // Options describes various options that can be used to create the client.
    25  type Options struct {
    26  	ClientID     string
    27  	ClientSecret string
    28  
    29  	// Scopes indicates what scope to request.
    30  	Scopes []string
    31  
    32  	// AccessType indicates what type of access is desireed. The two possible
    33  	// options are oauth2.AccessTypeOnline and oauth2.AccessTypeOffline.
    34  	AccessType oauth2.AuthCodeOption
    35  
    36  	// ConfigFile indicates where to look for and save the credentials. On Linux
    37  	// this is something like ~/.config/grail-role-group/credentials.json.
    38  	ConfigFile string
    39  
    40  	// OpenBrowser indicates that a browser should be open to obtain the proper
    41  	// credentials if they are missing.
    42  	OpenBrowser bool
    43  }
    44  
    45  // saveToken makes a best-effort attempt to save the token.
    46  func saveToken(configFile string, token *oauth2.Token) {
    47  	os.MkdirAll(filepath.Dir(configFile), 0700)
    48  	f, err := os.OpenFile(configFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
    49  	if err != nil {
    50  		vlog.Info(err)
    51  		return
    52  	}
    53  	defer f.Close()
    54  	if err := json.NewEncoder(f).Encode(token); err != nil {
    55  		vlog.Info(err)
    56  	}
    57  }
    58  
    59  func loadToken(configFile string) *oauth2.Token {
    60  	f, err := os.Open(configFile)
    61  	if err != nil {
    62  		vlog.Info(err)
    63  		return nil
    64  	}
    65  	t := &oauth2.Token{}
    66  	if err := json.NewDecoder(f).Decode(t); err != nil {
    67  		vlog.Info(err)
    68  		return nil
    69  	}
    70  	defer f.Close()
    71  	return t
    72  }
    73  
    74  // New returns a new http.Client suitable for passing to the New functions from
    75  // the packages under google.golang.org/api/. An interactive OAuth flow is
    76  // performed if the credentials don't exist in the config file. An attempt to
    77  // be open a web browser will done if opts.OpenBrowser is true.
    78  //
    79  // TODO(razvanm): add support for refresh tokens.
    80  // TODO(razvanm): add support for Application Default Credentials.
    81  func New(opts Options) (*http.Client, error) {
    82  	token := loadToken(opts.ConfigFile)
    83  	if token != nil {
    84  		// Check the validity of the access token.
    85  		config := &oauth2.Config{}
    86  		service, err := goauth2.New(http.DefaultClient)
    87  		if err != nil {
    88  			vlog.Info(err)
    89  		} else {
    90  			_, err = service.Tokeninfo().AccessToken(token.AccessToken).Do()
    91  			if err != nil {
    92  				vlog.Info(err)
    93  			} else {
    94  				return config.Client(context.Background(), token), nil
    95  			}
    96  		}
    97  	}
    98  
    99  	stateBytes := make([]byte, 16)
   100  	if _, err := rand.Read(stateBytes); err != nil {
   101  		return nil, err
   102  	}
   103  	state := hex.EncodeToString(stateBytes)
   104  
   105  	code := ""
   106  	wg := sync.WaitGroup{}
   107  	wg.Add(1)
   108  	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
   109  		if r.URL.Path != "/" {
   110  			return
   111  		}
   112  		if got, want := r.FormValue("state"), state; got != want {
   113  			cmdutil.Fatalf("Bad state: got %q, want %q", got, want)
   114  		}
   115  		code = r.FormValue("code")
   116  		w.Header().Set("Content-Type", "text/html")
   117  		// JavaScript only allows closing windows/tab that were open via JavaScript.
   118  		fmt.Fprintf(w, `<html><body>Code received. Please close this tab/window.</body></html>`)
   119  		wg.Done()
   120  	})
   121  
   122  	ln, err := net.Listen("tcp", "localhost:")
   123  	if err != nil {
   124  		return nil, err
   125  	}
   126  	vlog.Infof("listening: %v\n", ln.Addr().String())
   127  	port := strings.Split(ln.Addr().String(), ":")[1]
   128  	server := http.Server{Addr: "localhost:"}
   129  	go server.Serve(ln.(*net.TCPListener))
   130  
   131  	config := &oauth2.Config{
   132  		ClientID:     opts.ClientID,
   133  		ClientSecret: opts.ClientSecret,
   134  		Scopes:       opts.Scopes,
   135  		RedirectURL:  fmt.Sprintf("http://localhost:%s", port),
   136  		Endpoint:     google.Endpoint,
   137  	}
   138  
   139  	url := config.AuthCodeURL(state, opts.AccessType)
   140  
   141  	if opts.OpenBrowser {
   142  		fmt.Printf("Opening %q...\n", url)
   143  		if webutil.StartBrowser(url) {
   144  			wg.Wait()
   145  			server.Shutdown(context.Background())
   146  		} else {
   147  			opts.OpenBrowser = false
   148  		}
   149  	}
   150  
   151  	if !opts.OpenBrowser {
   152  		config.RedirectURL = "urn:ietf:wg:oauth:2.0:oob"
   153  		fmt.Printf("The attempt to automatically open a browser failed. Please open the following link in your browse:\n\n\t%s\n\n", url)
   154  		fmt.Printf("Paste the received code and then press enter: ")
   155  		fmt.Scanf("%s", &code)
   156  		fmt.Println("")
   157  	}
   158  
   159  	vlog.VI(1).Infof("code: %+v", code)
   160  	token, err = config.Exchange(oauth2.NoContext, code)
   161  	if err != nil {
   162  		return nil, err
   163  	}
   164  
   165  	saveToken(opts.ConfigFile, token)
   166  
   167  	return config.Client(context.Background(), token), nil
   168  }