telegabber/persistence/sessions.go

package persistence

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

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

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

// 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"`
}

var sessionDB *SessionsYamlDB

// SessionMarshaller implementation for YamlDB
func SessionMarshaller() ([]byte, error) {
	return yaml.Marshal(sessionDB.Data)
}

// 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
}