telegabber/persistence/sessions.go

package persistence

import (
	"github.com/pkg/errors"
	"io/ioutil"
	"time"

	"dev.narayana.im/narayana/telegabber/yamldb"

	log "github.com/sirupsen/logrus"
	"gopkg.in/yaml.v2"
)

var zeroLocation *time.Location

func init() {
	var err error
	zeroLocation, err = time.LoadLocation("")
	if err != nil {
		log.Fatal("Wrong hardcoded timezone")
	}
}

// SessionsYamlDB wraps YamlDB with Session
type SessionsYamlDB struct {
	yamldb.YamlDB
	Data *SessionsMap
}

// SessionsMap is for :sessions: subtree
type SessionsMap struct {
	Sessions map[string]Session `yaml:":sessions"`
}

// Session is a key-values subtree
type Session struct {
	Login       string `yaml:":login"`
	Timezone    string `yaml:":timezone"`
	KeepOnline  bool   `yaml:":keeponline"`
	RawMessages bool   `yaml:":rawmessages"`
}

var configKeys = []string{
	"timezone",
	"keeponline",
	"rawmessages",
}

var sessionDB *SessionsYamlDB

// SessionMarshaller implementation for YamlDB
func SessionMarshaller() ([]byte, error) {
	cleanedMap := SessionsMap{}
	emptySessionsMap(&cleanedMap)

	for jid, session := range sessionDB.Data.Sessions {
		if session.Login != "" {
			cleanedMap.Sessions[jid] = session
		}
	}

	return yaml.Marshal(&cleanedMap)
}

// LoadSessions restores TDlib sessions from the previous run
func LoadSessions(path string) (*SessionsYamlDB, error) {
	var sessionData SessionsMap
	var err error

	sessionDB, err = initYamlDB(path, &sessionData)
	if err != nil {
		return sessionDB, errors.Wrap(err, "Sessions restore error")
	}

	return sessionDB, nil
}

func emptySessionsMap(dataPtr *SessionsMap) {
	dataPtr.Sessions = make(map[string]Session)
}

func initYamlDB(path string, dataPtr *SessionsMap) (*SessionsYamlDB, error) {
	file, err := ioutil.ReadFile(path)
	if err == nil {
		err = yaml.Unmarshal(file, dataPtr)
		if err != nil {
			return nil, errors.Wrap(err, "YamlDB is corrupted")
		}

		if dataPtr.Sessions == nil {
			emptySessionsMap(dataPtr)
		}
		log.Debugf("Unmarshalled YAML: %#v", *dataPtr)
	} else {
		// DB file does not exist, create an empty DB
		emptySessionsMap(dataPtr)
	}

	return &SessionsYamlDB{
		YamlDB: yamldb.YamlDB{
			Path:    path,
			PathNew: path + ".new",
		},
		Data: dataPtr,
	}, nil
}

// Get retrieves a session value
func (s *Session) Get(key string) (string, error) {
	switch key {
	case "timezone":
		return s.Timezone, nil
	case "keeponline":
		return fromBool(s.KeepOnline), nil
	case "rawmessages":
		return fromBool(s.RawMessages), nil
	}

	return "", errors.New("Unknown session property")
}

// ToMap converts the session to a map
func (s *Session) ToMap() map[string]string {
	m := make(map[string]string)
	for _, configKey := range configKeys {
		value, _ := s.Get(configKey)
		m[configKey] = value
	}

	return m
}

// Set sets a session value
func (s *Session) Set(key string, value string) (string, error) {
	switch key {
	case "timezone":
		s.Timezone = value
		return value, nil
	case "keeponline":
		b, err := toBool(value)
		if err != nil {
			return "", err
		}
		s.KeepOnline = b
		return value, nil
	case "rawmessages":
		b, err := toBool(value)
		if err != nil {
			return "", err
		}
		s.RawMessages = b
		return value, nil
	}

	return "", errors.New("Unknown session property")
}

// TimezoneToLocation tries to convert config timezone to location
func (s *Session) TimezoneToLocation() *time.Location {
	time, err := time.Parse("-07:00", s.Timezone)
	if err == nil {
		return time.Location()
	}

	// default
	return zeroLocation
}

func fromBool(b bool) string {
	if b {
		return "true"
	} else {
		return "false"
	}
}

func toBool(s string) (bool, error) {
	switch s {
	case "true":
		return true, nil
	case "false":
		return false, nil
	}

	return false, errors.New("Invalid boolean value")
}