package dovecot

import (
	"context"
	"encoding/base64"
	"errors"
	"log"
	"strings"

	"git.autistici.org/ai3/go-common/clientutil"

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

// Config for the dovecot-keystore daemon.
type Config struct {
	Shard      string                    `yaml:"shard"`
	LDAPConfig *backend.LDAPConfig       `yaml:"ldap"`
	Keystore   *clientutil.BackendConfig `yaml:"keystore"`

	// Set this to true if the keys obtained from the backend need
	// to be base64-encoded before being sent to Dovecot.
	Base64Encode bool `yaml:"base64_encode_results"`
}

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

func (c *Config) check() error {
	if c.Keystore == nil {
		return errors.New("missing keystore config")
	}
	if c.LDAPConfig == nil {
		return errors.New("missing backend config")
	}
	return c.LDAPConfig.Valid()
}

type userdbResponse struct {
	PublicKey string `json:"mail_crypt_global_public_key"`
}

type passdbResponse struct {
	PrivateKey string `json:"mail_crypt_global_private_key"`
}

var passwordSep = "/"

// KeyLookupProxy interfaces Dovecot with the user encryption key database.
type KeyLookupProxy struct {
	config   *Config
	keystore client.Client
	db       Database
}

// NewKeyLookupProxy returns a KeyLookupProxy with the specified configuration.
func NewKeyLookupProxy(config *Config) (*KeyLookupProxy, error) {
	if err := config.check(); err != nil {
		return nil, err
	}

	ksc, err := client.New(config.Keystore)
	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
	}

	return &KeyLookupProxy{
		config:   config,
		keystore: ksc,
		db:       ldap,
	}, nil
}

const (
	namespace    = "shared"
	namespaceLen = len(namespace)
)

// Lookup a key using the dovecot dict proxy interface.
//
// We can be sent a userdb lookup, or a passdb lookup, and we can tell
// them apart with the key prefix (passdb/ or userdb/).
func (s *KeyLookupProxy) Lookup(ctx context.Context, key string) (interface{}, bool, error) {
	switch {
	case strings.HasPrefix(key, namespace+"/passdb/"):
		kparts := strings.SplitN(key[namespaceLen+8:], passwordSep, 2)
		return s.lookupPassdb(ctx, kparts[0], kparts[1])
	case strings.HasPrefix(key, namespace+"/userdb/"):
		return s.lookupUserdb(ctx, key[namespaceLen+8:])
	default:
		log.Printf("unknown key %s", key)
		return nil, false, nil
	}
}

func (s *KeyLookupProxy) lookupUserdb(ctx context.Context, username string) (interface{}, bool, error) {
	pub, err := s.db.GetPublicKey(ctx, username)
	if err != nil {
		return nil, false, err
	}
	if pub == nil {
		log.Printf("userdb lookup for %s (no keys)", username)
		return nil, false, nil
	}
	log.Printf("userdb lookup for %s", username)
	return &userdbResponse{PublicKey: s.b64encode(pub)}, true, nil
}

func (s *KeyLookupProxy) lookupPassdb(ctx context.Context, username, password string) (interface{}, bool, error) {
	// If the password is a SSO token, try to fetch the
	// unencrypted key from the keystore daemon.
	priv, err := s.keystore.Get(ctx, s.config.Shard, username, password)
	if err == nil {
		log.Printf("passdb lookup for %s (from keystore)", username)
		return &passdbResponse{PrivateKey: s.b64encode(priv)}, true, nil
	}

	// Otherwise, fetch encrypted keys from the db and attempt to
	// decrypt them.
	encKeys, err := s.db.GetPrivateKeys(ctx, username)
	if err != nil {
		return nil, false, err
	}
	if len(encKeys) == 0 {
		log.Printf("failed passdb lookup for %s (no keys)", username)
		return nil, false, nil
	}
	priv, err = userenckey.Decrypt(encKeys, []byte(password))
	if err != nil {
		log.Printf("failed passdb lookup for %s (could not decrypt key)", username)
		return nil, false, err
	}
	log.Printf("passdb lookup for %s (decrypted)", username)
	return &passdbResponse{PrivateKey: s.b64encode(priv)}, true, nil
}

func (s *KeyLookupProxy) b64encode(b []byte) string {
	if s.config.Base64Encode {
		return base64.StdEncoding.EncodeToString(b)
	}
	return string(b)
}