Add support for user-provided refresh token

jcruzcruz · Feb 20, 2016 · 21cc148 · 21cc148
1 parent 28c4eb9
commit 21cc148
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 13 deletions.
diff --git a/auth/oauth.go b/auth/oauth.go
@@ -2,23 +2,15 @@ package auth
 
 import (
     "fmt"
+    "time"
     "net/http"
     "golang.org/x/oauth2"
 )
 
 type authCodeFn func(string) func() string
 
-func NewOauthClient(clientId, clientSecret, tokenFile string, authFn authCodeFn) (*http.Client, error) {
-    conf := &oauth2.Config{
-        ClientID:     clientId,
-        ClientSecret: clientSecret,
-        Scopes:       []string{"https://www.googleapis.com/auth/drive"},
-        RedirectURL:  "urn:ietf:wg:oauth:2.0:oob",
-        Endpoint: oauth2.Endpoint{
-            AuthURL:  "https://accounts.google.com/o/oauth2/auth",
-            TokenURL: "https://accounts.google.com/o/oauth2/token",
-        },
-    }
+func NewFileSourceClient(clientId, clientSecret, tokenFile string, authFn authCodeFn) (*http.Client, error) {
+    conf := getConfig(clientId, clientSecret)
 
     // Read cached token
     token, exists, err := ReadToken(tokenFile)
@@ -42,3 +34,31 @@ func NewOauthClient(clientId, clientSecret, tokenFile string, authFn authCodeFn)
         FileSource(tokenFile, token, conf),
     ), nil
 }
+
+func NewRefreshTokenClient(clientId, clientSecret, refreshToken string) *http.Client {
+    conf := getConfig(clientId, clientSecret)
+
+    token := &oauth2.Token{
+        TokenType: "Bearer",
+        RefreshToken: refreshToken,
+        Expiry: time.Now(),
+    }
+
+    return oauth2.NewClient(
+        oauth2.NoContext,
+        conf.TokenSource(oauth2.NoContext, token),
+    )
+}
+
+func getConfig(clientId, clientSecret string) *oauth2.Config {
+    return &oauth2.Config{
+        ClientID:     clientId,
+        ClientSecret: clientSecret,
+        Scopes:       []string{"https://www.googleapis.com/auth/drive"},
+        RedirectURL:  "urn:ietf:wg:oauth:2.0:oob",
+        Endpoint: oauth2.Endpoint{
+            AuthURL:  "https://accounts.google.com/o/oauth2/auth",
+            TokenURL: "https://accounts.google.com/o/oauth2/token",
+        },
+    }
+}
diff --git a/gdrive.go b/gdrive.go
@@ -28,6 +28,11 @@ func main() {
             Description: fmt.Sprintf("Application path, default: %s", DefaultConfigDir),
             DefaultValue: DefaultConfigDir,
         },
+        cli.StringFlag{
+            Name: "refreshToken",
+            Patterns: []string{"--refresh-token"},
+            Description: "Oauth refresh token used to get access token (for advanced users)",
+        },
     }
 
     handlers := []*cli.Handler{

diff --git a/handlers_drive.go b/handlers_drive.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"io/ioutil"
 	"path/filepath"
+	"net/http"
     "./cli"
 	"./auth"
 	"./drive"
@@ -309,10 +310,18 @@ func aboutExportHandler(ctx cli.Context) {
     checkErr(err)
 }
 
-func newDrive(args cli.Arguments) *drive.Drive {
+func getOauthClient(args cli.Arguments) (*http.Client, error) {
+    if args.String("refreshToken") != "" {
+        return auth.NewRefreshTokenClient(ClientId, ClientSecret, args.String("refreshToken")), nil
+    }
+
     configDir := args.String("configDir")
     tokenPath := ConfigFilePath(configDir, TokenFilename)
-    oauth, err := auth.NewOauthClient(ClientId, ClientSecret, tokenPath, authCodePrompt)
+    return auth.NewFileSourceClient(ClientId, ClientSecret, tokenPath, authCodePrompt)
+}
+
+func newDrive(args cli.Arguments) *drive.Drive {
+    oauth, err := getOauthClient(args)
     if err != nil {
         ExitF("Failed getting oauth client: %s", err.Error())
     }