package server

import (
	"context"
	"errors"
	"io/ioutil"
	"log"
	"strings"
	"sync"
	"time"

	"git.autistici.org/id/go-sso"

	"git.autistici.org/id/keystore/backend"
	"git.autistici.org/id/keystore/userenckey"
)

var (
	errNoKeys       = errors.New("no keys available")
	errBadUser      = errors.New("username does not match authentication token")
	errUnauthorized = errors.New("unauthorized")
	errInvalidTTL   = errors.New("invalid ttl")
)

// Database represents the interface to the underlying backend for
// encrypted user keys.
type Database interface {
	GetPrivateKeys(context.Context, string) ([][]byte, error)
}

type userKey struct {
	pkey   []byte
	expiry time.Time
}

// Config for the KeyStore.
type Config struct {
	SSOPublicKeyFile string `yaml:"sso_public_key_file"`
	SSOService       string `yaml:"sso_service"`
	SSODomain        string `yaml:"sso_domain"`

	//Backend string `yaml:"backend"`
	LDAPConfig *backend.LDAPConfig `yaml:"ldap"`
}

func (c *Config) check() error {
	if c.SSOService == "" {
		return errors.New("sso_service is empty")
	}
	if !strings.HasSuffix(c.SSOService, "/") {
		return errors.New("sso_service is invalid (does not end with /)")
	}
	if c.SSODomain == "" {
		return errors.New("sso_domain is empty")
	}
	if c.LDAPConfig == nil {
		return errors.New("missing backend config")
	}
	return c.LDAPConfig.Valid()
}

// KeyStore holds decrypted secrets for users in memory for a short
// time (of the order of a SSO session lifespan). User secrets can be
// opened with a password (used to decrypt the key, which is stored
// encrypted in a database), queried, and closed (forgotten).
//
// The database can provide multiple versions of the encrypted key (to
// support multiple decryption passwords), in which case we'll try
// them all sequentially until one of them decrypts successfully with
// the provided password.
//
// In order to query the KeyStore, you need to present a valid SSO
// token for the user whose secrets you would like to obtain.
//
type KeyStore struct {
	mx       sync.Mutex
	userKeys map[string]userKey

	db        Database
	service   string
	validator sso.Validator
}

// NewKeyStore creates a new KeyStore with the given config and returns it.
func NewKeyStore(config *Config) (*KeyStore, error) {
	if err := config.check(); err != nil {
		return nil, err
	}

	ssoKey, err := ioutil.ReadFile(config.SSOPublicKeyFile)
	if err != nil {
		return nil, err
	}
	v, err := sso.NewValidator(ssoKey, config.SSODomain)
	if err != nil {
		return nil, err
	}

	// There is only one supported backend type, ldap.
	ldap, err := backend.NewLDAPBackend(config.LDAPConfig)
	if err != nil {
		return nil, err
	}

	s := &KeyStore{
		userKeys:  make(map[string]userKey),
		service:   config.SSOService,
		validator: v,
		db:        ldap,
	}
	go s.expire()
	return s, nil
}

func (s *KeyStore) expire() {
	for t := range time.NewTicker(600 * time.Second).C {
		s.mx.Lock()
		for u, k := range s.userKeys {
			if k.expiry.After(t) {
				wipeBytes(k.pkey)
				delete(s.userKeys, u)
			}
		}
		s.mx.Unlock()
	}
}

// Open the user's key store with the given password. If successful,
// the unencrypted user key will be stored for at most ttlSeconds, or
// until Close is called.
//
// A Context is needed because this method might issue an RPC.
func (s *KeyStore) Open(ctx context.Context, username, password string, ttlSeconds int) error {
	if ttlSeconds == 0 {
		return errInvalidTTL
	}

	encKeys, err := s.db.GetPrivateKeys(ctx, username)
	if err != nil {
		return err
	}
	if len(encKeys) == 0 {
		// No keys found. Not an error.
		return nil
	}

	// Naive and inefficient way of decrypting multiple keys: it
	// will recompute the kdf every time, which is expensive.
	pkey, err := userenckey.Decrypt(encKeys, []byte(password))
	if err != nil {
		return err
	}

	s.mx.Lock()
	s.userKeys[username] = userKey{
		pkey:   pkey,
		expiry: time.Now().Add(time.Duration(ttlSeconds) * time.Second),
	}
	s.mx.Unlock()
	return nil
}

// Get the unencrypted key for the specified user. The caller needs to
// provide a valid SSO ticket for the user.
func (s *KeyStore) Get(username, ssoTicket string) ([]byte, error) {
	// Validate the SSO ticket.
	tkt, err := s.validator.Validate(ssoTicket, "", s.service, nil)
	if err != nil {
		// Log authentication failures for debugging purposes.
		log.Printf("Validate(%s) error: %v", username, err)
		return nil, errUnauthorized
	}
	if tkt.User != username {
		log.Printf("Validate(%s) user mismatch: sso=%s", username, tkt.User)
		return nil, errBadUser
	}

	s.mx.Lock()
	defer s.mx.Unlock()
	u, ok := s.userKeys[username]
	if !ok {
		return nil, errNoKeys
	}
	return u.pkey, nil
}

// Close the user's key store and wipe the associated unencrypted key
// from memory.
func (s *KeyStore) Close(username string) {
	s.mx.Lock()
	if k, ok := s.userKeys[username]; ok {
		wipeBytes(k.pkey)
		delete(s.userKeys, username)
	}
	s.mx.Unlock()
}

func wipeBytes(b []byte) {
	for i := 0; i < len(b); i++ {
		b[i] = 0
	}
}