telegabber/persistence/sessions.go

package persistence

import (
	"github.com/pkg/errors"
	"io/ioutil"
	"sync"
	"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"`
	AsciiArrows     bool           `yaml:":asciiarrows"`
	OOBMode         bool           `yaml:":oobmode"`
	Carbons         bool           `yaml:":carbons"`
	HideIds         bool           `yaml:":hideids"`
	Receipts        bool           `yaml:":receipts"`
	NativeEdits     bool           `yaml:":nativeedits"`
	IgnoredChats    []int64        `yaml:":ignoredchats"`
	ignoredChatsMap map[int64]bool `yaml:"-"`
}

var configKeys = []string{
	"timezone",
	"keeponline",
	"rawmessages",
	"asciiarrows",
	"oobmode",
	"carbons",
	"hideids",
	"receipts",
	"nativeedits",
}

var sessionDB *SessionsYamlDB
var sessionsLock sync.Mutex

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

	sessionsLock.Lock()
	defer sessionsLock.Unlock()
	for jid, session := range sessionDB.Data.Sessions {
		if session.Login != "" {
			session.IgnoredChats = make([]int64, 0, len(session.ignoredChatsMap))
			for chatID := range session.ignoredChatsMap {
				session.IgnoredChats = append(session.IgnoredChats, chatID)
			}
			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)
	}

	// convert ignored users slice to map
	for jid, session := range dataPtr.Sessions {
		session.ignoredChatsMap = make(map[int64]bool)
		for _, chatID := range session.IgnoredChats {
			session.ignoredChatsMap[chatID] = true
		}
		session.IgnoredChats = nil
		dataPtr.Sessions[jid] = session
	}

	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) {
	sessionsLock.Lock()
	defer sessionsLock.Unlock()

	return s.get(key)
}

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
	case "asciiarrows":
		return fromBool(s.AsciiArrows), nil
	case "oobmode":
		return fromBool(s.OOBMode), nil
	case "carbons":
		return fromBool(s.Carbons), nil
	case "hideids":
		return fromBool(s.HideIds), nil
	case "receipts":
		return fromBool(s.Receipts), nil
	case "nativeedits":
		return fromBool(s.NativeEdits), nil
	}

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

// ToMap converts the session to a map
func (s *Session) ToMap() map[string]string {
	sessionsLock.Lock()
	defer sessionsLock.Unlock()

	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) {
	sessionsLock.Lock()
	defer sessionsLock.Unlock()

	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
	case "asciiarrows":
		b, err := toBool(value)
		if err != nil {
			return "", err
		}
		s.AsciiArrows = b
		return value, nil
	case "oobmode":
		b, err := toBool(value)
		if err != nil {
			return "", err
		}
		s.OOBMode = b
		return value, nil
	case "carbons":
		b, err := toBool(value)
		if err != nil {
			return "", err
		}
		s.Carbons = b
		return value, nil
	case "hideids":
		b, err := toBool(value)
		if err != nil {
			return "", err
		}
		s.HideIds = b
		return value, nil
	case "receipts":
		b, err := toBool(value)
		if err != nil {
			return "", err
		}
		s.Receipts = b
		return value, nil
	case "nativeedits":
		b, err := toBool(value)
		if err != nil {
			return "", err
		}
		s.NativeEdits = 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
}

// IgnoreChat adds a chat id to ignore list, returns false if already ignored
func (s *Session) IgnoreChat(chatID int64) bool {
	sessionsLock.Lock()
	defer sessionsLock.Unlock()

	if s.ignoredChatsMap == nil {
		s.ignoredChatsMap = make(map[int64]bool)
	} else if _, ok := s.ignoredChatsMap[chatID]; ok {
		return false
	}

	s.ignoredChatsMap[chatID] = true
	return true
}

// UnignoreChat removes a chat id from ignore list, returns false if not already ignored
func (s *Session) UnignoreChat(chatID int64) bool {
	sessionsLock.Lock()
	defer sessionsLock.Unlock()

	if s.ignoredChatsMap == nil {
		return false
	}

	if _, ok := s.ignoredChatsMap[chatID]; !ok {
		return false
	}

	delete(s.ignoredChatsMap, chatID)
	return true
}

// IsChatIgnored checks the chat id against the ignore list
func (s *Session) IsChatIgnored(chatID int64) bool {
	sessionsLock.Lock()
	defer sessionsLock.Unlock()

	if s.ignoredChatsMap == nil {
		return false
	}

	_, ok := s.ignoredChatsMap[chatID]
	return ok
}

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")
}