FileSource TokenSource

4 files changed, 99 insertions, 9 deletions

diff --git a/client/auth.go b/client/auth.go
index fb6852d..136ace0 100644
--- a/client/auth.go
+++ b/client/auth.go
@@ -3,12 +3,11 @@ package client
 import (
     "net/http"
     "golang.org/x/oauth2"
-    "go4.org/oauthutil"
 )
 
 type authCodeFn func(string) func() string
 
-func NewOauthClient(clientId, clientSecret, cacheFile string, authFn authCodeFn) *http.Client {
+func NewOauthClient(clientId, clientSecret, tokenFile string, authFn authCodeFn) (*http.Client, error) {
     conf := &oauth2.Config{
         ClientID:     clientId,
         ClientSecret: clientSecret,
@@ -20,13 +19,21 @@ func NewOauthClient(clientId, clientSecret, cacheFile string, authFn authCodeFn)
         },
     }
 
-    authUrl := conf.AuthCodeURL("state", oauth2.AccessTypeOffline)
+    // Read cached token
+    token, exists, err := ReadToken(tokenFile)
+    if err != nil {
+        return nil, err
+    }
 
-    tokenSource := oauthutil.TokenSource{
-        Config: conf,
-        CacheFile: cacheFile,
-        AuthCode: authFn(authUrl),
+    // Request auth code if token does not exist
+    if !exists {
+        authUrl := conf.AuthCodeURL("state", oauth2.AccessTypeOffline)
+        authCode := authFn(authUrl)()
+        token, err = conf.Exchange(oauth2.NoContext, authCode)
     }
 
-    return oauth2.NewClient(oauth2.NoContext, tokenSource)
+    return oauth2.NewClient(
+        oauth2.NoContext,
+        FileSource(tokenFile, token, conf),
+    ), nil
 }
diff --git a/client/token.go b/client/token.go
new file mode 100644
index 0000000..0931284
--- /dev/null
+++ b/client/token.go
@@ -0,0 +1,57 @@
+package client
+
+import (
+    "golang.org/x/oauth2"
+    "encoding/json"
+    "io/ioutil"
+)
+
+
+func FileSource(path string, token *oauth2.Token, conf *oauth2.Config) oauth2.TokenSource {
+    return &fileSource{
+        tokenPath: path,
+        tokenSource: conf.TokenSource(oauth2.NoContext, token),
+    }
+}
+
+type fileSource struct {
+    tokenPath string
+    tokenSource oauth2.TokenSource
+}
+
+func (self *fileSource) Token() (*oauth2.Token, error) {
+    token, err := self.tokenSource.Token()
+    if err != nil {
+        return token, err
+    }
+
+    // Save token to file
+    SaveToken(self.tokenPath, token)
+
+    return token, nil
+}
+
+func ReadToken(path string) (*oauth2.Token, bool, error) {
+    if !fileExists(path) {
+        return nil, false, nil
+    }
+
+    content, err := ioutil.ReadFile(path)
+    if err != nil {
+        return nil, true, err
+    }
+    token := &oauth2.Token{}
+    return token, true, json.Unmarshal(content, token)
+}
+
+func SaveToken(path string, token *oauth2.Token) error {
+    data, err := json.MarshalIndent(token, "", "  ")
+    if err != nil {
+        return err
+    }
+
+    if err = mkdir(path); err != nil {
+        return err
+    }
+    return ioutil.WriteFile(path, data, 0600)
+}
diff --git a/client/util.go b/client/util.go
new file mode 100644
index 0000000..b600fd6
--- /dev/null
+++ b/client/util.go
@@ -0,0 +1,22 @@
+package client
+
+import (
+    "os"
+    "path/filepath"
+)
+
+func mkdir(path string) error {
+    dir := filepath.Dir(path)
+    if fileExists(dir) {
+        return nil
+    }
+    return os.Mkdir(dir, 0700)
+}
+
+func fileExists(path string) bool {
+    _, err := os.Stat(path)
+    if err == nil {
+        return true
+    }
+    return false
+}
diff --git a/handlers.go b/handlers.go
index 69aad94..a4bb3e2 100644
--- a/handlers.go
+++ b/handlers.go
@@ -75,7 +75,11 @@ func printCommandHelp(ctx cli.Context) {
 }
 
 func newDrive() *drive.Drive {
-    oauth := client.NewOauthClient(ClientId, ClientSecret, DefaultTokenFilePath, authCodePrompt)
+    oauth, err := client.NewOauthClient(ClientId, ClientSecret, DefaultTokenFilePath, authCodePrompt)
+    if err != nil {
+        ExitF("Failed getting oauth client: %s", err.Error())
+    }
+
     client, err := client.NewClient(oauth)
     if err != nil {
         ExitF("Failed getting drive: %s", err.Error())