gdrive-filereplacer/main.go

package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"github.com/andrew-d/go-termutil"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/drive/v3"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/option"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
)

type TokenCache struct {
	CredentialPath   string
	CredentialScopes []string
	ConfigPath       string
	TokenPath        string
	configLoaded     bool
	CredentialData   *google.Credentials
	ConfigData       oauth2.Config
}

func (c *TokenCache) Credentials() (*google.Credentials, error) {
	if c.CredentialData == nil {
		f, err := os.Open(c.CredentialPath)
		if err != nil {
			return nil, fmt.Errorf("could not open credentials: %v", err)
		}
		defer func() { _ = f.Close() }()
		data, err := ioutil.ReadAll(f)
		if err != nil {
			return nil, fmt.Errorf("could not read credentials: %v", err)
		}
		c.CredentialData, err = google.CredentialsFromJSON(context.Background(), data, c.CredentialScopes...)
		if err != nil {
			return nil, fmt.Errorf("could not decode credentials: %v", err)
		}
	}
	return c.CredentialData, nil
}

func (c *TokenCache) Config() (*oauth2.Config, error) {
	if !c.configLoaded {
		f, err := os.Open(c.ConfigPath)
		if err != nil {
			return nil, fmt.Errorf("could not open config: %v", err)
		}
		defer func() { _ = f.Close() }()
		err = json.NewDecoder(f).Decode(&c.ConfigData)
		if err != nil {
			return nil, fmt.Errorf("could not decode config: %v", err)
		}
		c.configLoaded = true
	}
	return &c.ConfigData, nil
}

func (c *TokenCache) CredentialToken() (*oauth2.Token, error) {
	log.Printf("Getting token from credentials file %s", c.CredentialPath)
	creds, err := c.Credentials()
	if err != nil {
		return nil, err
	}
	return creds.TokenSource.Token()
}

func (c *TokenCache) Token() (*oauth2.Token, error) {
	if c.CredentialPath != "" {
		return c.CredentialToken()
	} else {
		return c.StandardToken()
	}
}

// Retrieves a token and saves the token.
func (c *TokenCache) StandardToken() (*oauth2.Token, error) {
	// The file token.json stores the user's access and refresh tokens, and is
	// created automatically when the authorization flow completes for the first
	// time.
	log.Printf("Getting token from secret file %s with cache in %s", c.ConfigPath, c.TokenPath)
	tok, err := c.CachedToken()
	if err == nil && tok.Valid() {
		return tok, nil
	}
	if err != nil {
		log.Printf("Could not get token from cache: %v", err)
	} else if !tok.Valid() {
		log.Print("Cached token is invalid, refreshing...")
		tok, err = c.RefreshToken(tok)
		if err != nil {
			log.Printf("Cached token failed to refresh: %v", err)
		}
	}
	if err != nil {
		if !termutil.Isatty(os.Stdin.Fd()) {
			return nil, errors.New("cached token is invalid, and can't get token from web without stdin being a terminal")
		}
		log.Print("Could not use cached token. Getting a new token from web...")
		tok, err = c.NewToken()
		if err != nil {
			return nil, err
		}
	}
	err = c.SaveToken(tok)
	if err != nil {
		log.Printf("Failed caching token: %v", err)
	}
	return tok, nil
}

// Request a token from the web, then returns the retrieved token.
func (c *TokenCache) NewToken() (*oauth2.Token, error) {
	config, err := c.Config()
	if err != nil {
		return nil, fmt.Errorf("failed loading config: %v", err)
	}
	authURL := config.AuthCodeURL("state-token", oauth2.AccessTypeOffline)
	fmt.Printf("Go to the following link in your browser then type the "+
		"authorization code: \n%v\n", authURL)

	var authCode string
	if _, err := fmt.Scan(&authCode); err != nil {
		return nil, fmt.Errorf("unable to read authorization code: %v", err)
	}

	tok, err := config.Exchange(context.TODO(), authCode)
	if err != nil {
		return nil, fmt.Errorf("unable to retrieve token from web: %v", err)
	}
	return tok, nil
}

func (c *TokenCache) RefreshToken(token *oauth2.Token) (*oauth2.Token, error) {
	config, err := c.Config()
	if err != nil {
		return nil, fmt.Errorf("failed loading config: %v", err)
	}

	ts := config.TokenSource(context.Background(), token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	if token.AccessToken == newToken.AccessToken {
		return nil, fmt.Errorf("refreshed token is the same as the original token")
	}
	return newToken, nil
}

// Retrieves a token from a local file.
func (c *TokenCache) CachedToken() (*oauth2.Token, error) {
	f, err := os.Open(c.TokenPath)
	if err != nil {
		return nil, err
	}
	defer func() { _ = f.Close() }()
	tok := &oauth2.Token{}
	err = json.NewDecoder(f).Decode(tok)
	return tok, err
}

// Saves a token to a file path.
func (c *TokenCache) SaveToken(token *oauth2.Token) error {
	f, err := os.OpenFile(c.TokenPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return fmt.Errorf("unable to write to cache for OAuth2 token: %v", err)
	}
	defer func() { _ = f.Close() }()
	err = json.NewEncoder(f).Encode(token)
	if err != nil {
		return fmt.Errorf("unable to write OAuth2 token to cache: %v", err)
	}
	return err
}

func main() {
	tokenCacheSource := TokenCache{
		CredentialScopes: []string{drive.DriveScope},
		ConfigData: oauth2.Config{
			Endpoint:    google.Endpoint,
			RedirectURL: "urn:ietf:wg:oauth:2.0:oob",
			Scopes:      []string{drive.DriveFileScope},
		},
	}
	flag.StringVar(
		&tokenCacheSource.CredentialPath,
		"credentials",
		"",
		"If set, overrides secrets.json and token.json. The path where the credentials file should be read from.")
	flag.StringVar(
		&tokenCacheSource.TokenPath,
		"token",
		"./token.json",
		"The path where the cached token should be stored and read from.")
	flag.StringVar(
		&tokenCacheSource.ConfigPath,
		"secrets",
		"./secrets.json",
		"The path to the json file containing the client ID and secret.")
	fileId := flag.String(
		"file-id",
		"",
		"The ID of the file that should be overwritten. If not provided, a new file will be uploaded.")
	stdin := flag.String(
		"use-stdin-as",
		"",
		"Accepts input from stdin instead of a file, and uses the provided value as the file's name.")
	checkToken := flag.Bool(
		"check-token",
		false,
		"True to exit after getting the token.")
	contentType := flag.String(
		"content-type",
		"",
		"The content type to upload as. If omitted, no explicit content type is set.")
	flag.Parse()
	var dataFile *os.File
	var dataFilename string
	switch flag.NArg() {
	case 0:
		if *stdin == "" && !*checkToken {
			flag.Usage()
			log.Fatal("The data file path must be given, or --use-stdin-as or --check-token set.")
		}
		dataFile = os.Stdin
		dataFilename = *stdin
	case 1:
		if *stdin != "" || *checkToken {
			flag.Usage()
			log.Fatal("The data file path is not used if --use-stdin-as or --check-token is set.")
		}
		dataPath := flag.Arg(0)
		dataFilename = filepath.Base(dataPath)
		var err error
		dataFile, err = os.Open(dataPath)
		if err != nil {
			log.Fatalf("The data file %s could not be opened: %v", dataPath, err)
		}
		defer func() { _ = dataFile.Close() }()
	default:
		flag.Usage()
		log.Fatal("The only argument accepted is the data file path, if any.")
	}

	token, err := tokenCacheSource.Token()
	if err != nil {
		log.Fatalf("Failed getting token: %v", err)
	}
	if *checkToken {
		fmt.Println(token.Expiry)
		return
	}
	ts := oauth2.ReuseTokenSource(token, &tokenCacheSource)

	srv, err := drive.NewService(
		context.Background(),
		option.WithTokenSource(ts))
	if err != nil {
		log.Fatalf("Unable to initialize Drive service: %v", err)
	}

	metadata := drive.File{
		Name: dataFilename,
	}

	var mediaOptions []googleapi.MediaOption
	if *contentType != "" {
		mediaOptions = append(mediaOptions, googleapi.ContentType(*contentType))
	}

	if *fileId != "" {
		fileOut, err := srv.Files.Update(*fileId, &metadata).Media(dataFile, mediaOptions...).Do()
		if err != nil {
			log.Fatalf("Failed updating file %s: %v", *fileId, err)
		}
		fmt.Println(fileOut.Id)
	} else {
		fileOut, err := srv.Files.Create(&metadata).Media(dataFile, mediaOptions...).Do()
		if err != nil {
			log.Fatalf("Failed uploading file: %v", err)
		}
		fmt.Println(fileOut.Id)
	}
}