oauth_token_store.go

package integram

import (
	"fmt"
	"time"

	"gopkg.in/mgo.v2/bson"
)

type OAuthTokenStore interface {
	Name() string
	GetOAuthAccessToken(user *User) (token string, expireDate *time.Time, err error)
	SetOAuthAccessToken(user *User, token string, expireDate *time.Time) error
	GetOAuthRefreshToken(user *User) (string, error)
	SetOAuthRefreshToken(user *User, token string) error
}

type DefaultOAuthTokenMongoStore struct {
}

var oauthTokenStore OAuthTokenStore = &DefaultOAuthTokenMongoStore{}

func SetOAuthTokenStore(store OAuthTokenStore) {
	oauthTokenStore = store
}

func MigrateOAuthFromTo(c *Context, oldTS OAuthTokenStore, newTS OAuthTokenStore, onlyValid bool) (total int, migrated int, expired int, err error) {
	keyPrefix := "protected." + c.ServiceName

	query := bson.M{
		keyPrefix + ".oauthstore": oldTS.Name(),
	}

	if onlyValid {
		query[keyPrefix+".oauthvalid"] = true
	}

	users, err := c.FindUsers(query)
	if err != nil {
		return
	}

	total = len(users)
	expiredOlderThan := time.Now().Add((-1) * time.Hour * 24 * 30)
	for i, userData := range users {
		ctxCopy := *userData.ctx
		ctxCopy.User = userData.User
		ctxCopy.User.ctx = &ctxCopy
		ctxCopy.Chat = Chat{ID: userData.ID, ctx: &ctxCopy}
		userData.ctx = &ctxCopy
		user := userData.User
		user.data = &userData

		if i%100 == 0 {
			fmt.Printf("MigrateOAuthFromTo: %d/%d users transfered\n", i, len(users))
		}

		token, expiry, err := oldTS.GetOAuthAccessToken(&user)
		if err != nil {
			c.Log().Errorf("MigrateOAuthFromTo got error on GetOAuthAccessToken: %s", err.Error())
			continue
		}

		if onlyValid && token == "" {
			expired++
			continue
		}

		if onlyValid && expiry != nil && expiry.Before(expiredOlderThan) {
			expired++
			continue
		}

		err = newTS.SetOAuthAccessToken(&user, token, expiry)
		if err != nil {
			c.Log().Errorf("MigrateOAuthFromTo got error on SetOAuthAccessToken: %s", err.Error())
			continue
		}

		refreshToken, err := oldTS.GetOAuthRefreshToken(&user)
		if err != nil {
			c.Log().Errorf("MigrateOAuthFromTo got error on GetOAuthRefreshToken: %s", err.Error())
			continue
		}

		err = newTS.SetOAuthRefreshToken(&user, refreshToken)
		if err != nil {
			c.Log().Errorf("MigrateOAuthFromTo got error on SetOAuthRefreshToken: %s", err.Error())
			continue
		}

		err = c.db.C("users").UpdateId(user.ID, bson.M{"$set": bson.M{keyPrefix + ".oauthstore": newTS.Name(), keyPrefix + ".oauthvalid": true}})
		if err != nil {
			c.Log().Errorf("MigrateOAuthFromTo got error: %s", err.Error())
			continue
		}

		migrated++
	}

	fmt.Printf("MigrateOAuthFromTo: %d/%d users transfered\n", len(users), len(users))

	return
}

func (d *DefaultOAuthTokenMongoStore) GetOAuthAccessToken(user *User) (token string, expireDate *time.Time, err error) {
	ps, err := user.protectedSettings()
	if err != nil {
		return "", nil, err
	}

	return ps.OAuthToken, ps.OAuthExpireDate, nil
}

func (d *DefaultOAuthTokenMongoStore) GetOAuthRefreshToken(user *User) (string, error) {
	ps, err := user.protectedSettings()

	if err != nil {
		return "", err
	}

	return ps.OAuthRefreshToken, nil
}

func (d *DefaultOAuthTokenMongoStore) SetOAuthAccessToken(user *User, token string, expireDate *time.Time) error {
	ps, err := user.protectedSettings()
	if err != nil {
		return err
	}

	ps.OAuthStore = d.Name()
	ps.OAuthToken = token
	ps.OAuthExpireDate = expireDate

	return user.saveProtectedSettings()
}

func (d *DefaultOAuthTokenMongoStore) SetOAuthRefreshToken(user *User, refreshToken string) error {
	ps, err := user.protectedSettings()
	if err != nil {
		return err
	}

	ps.OAuthRefreshToken = refreshToken

	return user.saveProtectedSetting("OAuthRefreshToken", refreshToken)
}

func (d *DefaultOAuthTokenMongoStore) Name() string {
	return "default"
}