package dovecot

import (
	"context"
	"database/sql"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
	"time"

	"git.autistici.org/ai3/go-common/clientutil"
	"git.autistici.org/ai3/go-common/userenckey"
	"git.autistici.org/id/go-sso"
	"golang.org/x/crypto/ed25519"
	"gopkg.in/yaml.v3"

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

func withTestDB(t testing.TB, schema string) (func(), string) {
	dir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatal(err)
	}
	dbPath := filepath.Join(dir, "test.db")

	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		t.Fatalf("sql.Open: %v", err)
	}
	_, err = db.Exec(schema)
	if err != nil {
		t.Fatalf("sql error: %v", err)
	}
	db.Close()

	return func() {
		os.RemoveAll(dir)
	}, dbPath
}

func makeMapSlice(m map[string]interface{}) (out yaml.Node) {
	d, _ := yaml.Marshal(m)
	yaml.Unmarshal(d, &out) // nolint: errcheck
	return
}

var (
	testPw1 = []byte("password1")
	testPw2 = []byte("password2")
)

// Set up a test keystore stack, including keyproxy, along with a test
// database and valid SSO credentials.
func createTestKeyProxy(t testing.TB) (func(), *keystore.KeyStore, *KeyLookupProxy, sso.Signer) {
	// Set up a database with a simple schema and some real key
	// data. The private key is then encrypted with the two valid
	// test passwords.
	encPub, encPriv, err := userenckey.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	key1, err := userenckey.Encrypt(encPriv, testPw1)
	if err != nil {
		t.Fatal(err)
	}
	key2, err := userenckey.Encrypt(encPriv, testPw2)
	if err != nil {
		t.Fatal(err)
	}

	cleanup, dbPath := withTestDB(t, fmt.Sprintf(`
CREATE TABLE users (
    email text NOT NULL,
    public_key text
);
CREATE UNIQUE INDEX users_idx ON users(email);
INSERT INTO users (email, public_key) VALUES (
    'test@example.com', '%s'), (
    'nokeys@example.com', NULL
);
CREATE TABLE keys (
    email text NOT NULL,
    key blob NOT NULL
);
INSERT INTO keys (email, key) VALUES (
    'test@example.com', X'%s'), (
    'test@example.com', X'%s'
);
`, string(encPub), hex.EncodeToString(key1), hex.EncodeToString(key2)))

	// Backend configuration is shared by the keystore and keyproxy.
	backendConfig := &backend.Config{
		Type: "sql",
		Params: makeMapSlice(map[string]interface{}{
			"driver": "sqlite3",
			"db_uri": dbPath,
			"queries": map[string]string{
				"get_public_key":   "SELECT public_key FROM users WHERE email = ?",
				"get_private_keys": "SELECT key FROM keys WHERE email = ?",
			},
		}),
	}

	// Start a KeyStore with its own HTTP server. This requires
	// valid SSO credentials.
	ssoPub, ssoPriv, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	ssoSigner, err := sso.NewSigner(ssoPriv)
	if err != nil {
		t.Fatal(err)
	}
	pubFile := filepath.Join(filepath.Dir(dbPath), "sso.pub")
	ioutil.WriteFile(pubFile, ssoPub, 0600) // nolint
	ksConfig := keystore.Config{
		Backend:          backendConfig,
		SSOPublicKeyFile: pubFile,
		SSOService:       "service/",
		SSODomain:        "domain",
	}
	ks, err := keystore.NewKeyStore(&ksConfig)
	if err != nil {
		t.Fatalf("NewKeyStore: %v", err)
	}
	ksSrv := httptest.NewServer(keystore.NewServer(ks))

	// Create the KeyProxy, connected to the KeyStore.
	config := Config{
		Backend: backendConfig,
		Keystore: &clientutil.BackendConfig{
			URL: ksSrv.URL,
		},
	}
	kp, err := NewKeyLookupProxy(&config)
	if err != nil {
		t.Fatalf("NewKeyLookupProxy: %v", err)
	}

	return func() {
		ksSrv.Close()
		cleanup()
	}, ks, kp, ssoSigner
}

func TestKeyProxy_LookupPublicKey(t *testing.T) {
	cleanup, _, kp, _ := createTestKeyProxy(t)
	defer cleanup()

	testData := []struct {
		username string
		found    bool
		error    bool
	}{
		{"test@example.com", true, false},
		// User without a key -> no error.
		{"nokeys@example.com", false, false},
		// User does not exist -> also no error.
		{"UNKNOWN@example.com", false, false},
	}

	for _, td := range testData {
		_, found, err := kp.Lookup(context.Background(), "shared/userdb/"+td.username)
		hasError := (err != nil)
		if hasError != td.error {
			t.Errorf("Lookup(%s): err=%v", td.username, err)
		}
		if found != td.found {
			t.Errorf("Lookup(%s): found=%v, expected %v", td.username, found, td.found)
		}
	}
}

func TestKeyProxy_LookupPrivateKeys_FromDB(t *testing.T) {
	cleanup, _, kp, _ := createTestKeyProxy(t)
	defer cleanup()

	testData := []struct {
		username string
		password string
		found    bool
		error    bool
	}{
		// Unlock the key from the database with both valid passwords.
		{"test@example.com", string(testPw1), true, false},
		{"test@example.com", string(testPw2), true, false},
		// A bad password should always fail with an error, so that
		// dovecot knows not to do anything.
		{"test@example.com", "bad password", false, true},
		// These users have no encryption keys, so we should *not*
		// return an error.
		{"nokeys@example.com", "password", false, false},
		{"UNKNOWN@example.com", "password", false, false},
	}

	for _, td := range testData {
		_, found, err := kp.Lookup(context.Background(), fmt.Sprintf("shared/passdb/%s/%s", td.username, td.password))
		hasError := (err != nil)
		if hasError != td.error {
			t.Errorf("Lookup(%s): err=%v", td.username, err)
		}
		if found != td.found {
			t.Errorf("Lookup(%s): found=%v, expected %v", td.username, found, td.found)
		}
	}
}

func TestKeyProxy_LookupPrivateKeys_FromKeystore(t *testing.T) {
	cleanup, ks, kp, signer := createTestKeyProxy(t)
	defer cleanup()

	// Unlock the key in KeyStore with the correct password.
	if err := ks.Open(context.Background(), "test@example.com", string(testPw1), "session", 600); err != nil {
		t.Fatalf("ks.Open: %v", err)
	}

	// Mint a good SSO token.
	ssoToken, err := signer.Sign(sso.NewTicket("test@example.com", "service/", "domain", "", nil, 600*time.Second))
	if err != nil {
		t.Fatalf("sso.Sign: %v", err)
	}

	testData := []struct {
		username string
		password string
		found    bool
		error    bool
	}{
		// Fetch the cached key using a SSO token.
		{"test@example.com", ssoToken, true, false},
		// Even if we can't authenticate with the keystore,
		// fallthru to the database should still work.
		{"test@example.com", string(testPw1), true, false},
		{"test@example.com", string(testPw2), true, false},
		// A bad password should always fail.
		{"test@example.com", "bad password", false, true},
	}

	for _, td := range testData {
		_, found, err := kp.Lookup(context.Background(), fmt.Sprintf("shared/passdb/%s/%s", td.username, td.password))
		hasError := (err != nil)
		if hasError != td.error {
			t.Errorf("Lookup(%s): err=%v", td.username, err)
		}
		if found != td.found {
			t.Errorf("Lookup(%s): found=%v, expected %v", td.username, found, td.found)
		}
	}
}