package console

import (
	"crypto/hmac"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"path/filepath"
	"testing"
	"time"
)

// testFlow is the flow hash the store tests' sign-in challenges belong to.
var testFlow = hashToken("a-test-flow-token-of-some-length-000000")

func newSignInChallenge(t *testing.T, s *Store, now time.Time, flowHash string) (link, code string) {
	t.Helper()
	link, code, err := s.CreateEmailChallenge(EmailChallenge{
		Purpose: ChallengeSignIn, Address: "bob@example.com", Next: "/auth/cli", FlowHash: flowHash,
		ExpiresAt: now.Add(ConfirmationTTL),
	}, now)
	if err != nil {
		t.Fatal(err)
	}
	return link, code
}

// live is the flow's newest live sign-in challenge, failing the test when
// there is none.
func live(t *testing.T, s *Store, flowHash string, now time.Time) *EmailChallenge {
	t.Helper()
	c, ok := s.LiveEmailChallenge(ChallengeSignIn, flowHash, now)
	if !ok {
		t.Fatal("the flow has no live sign-in challenge")
	}
	return c
}

// wrongCode returns a six-digit code that is not code.
func wrongCode(code string) string {
	if code == "000000" {
		return "000001"
	}
	return "000000"
}

func TestEmailChallengeLinkPeekDoesNotConsume(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()
	link, _ := newSignInChallenge(t, s, now, testFlow)
	for i := 0; i < 2; i++ {
		c, ok := s.PeekEmailLink(link, now)
		if !ok || c.Address != "bob@example.com" || c.Next != "/auth/cli" || c.Purpose != ChallengeSignIn || c.FlowHash != testFlow {
			t.Fatalf("peek %d = %+v, %v", i, c, ok)
		}
	}
	if c, ok := s.TakeEmailLink(link, now); !ok || c.Address != "bob@example.com" || c.Next != "/auth/cli" {
		t.Fatalf("take = %+v, %v", c, ok)
	}
	if _, ok := s.TakeEmailLink(link, now); ok {
		t.Fatal("a link was taken twice")
	}
	if _, ok := s.PeekEmailLink(link, now); ok {
		t.Fatal("a taken link still peeks")
	}
}

func TestEmailChallengeLiveByFlow(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()
	_, code := newSignInChallenge(t, s, now, testFlow)
	for i := 0; i < 2; i++ {
		c := live(t, s, testFlow, now)
		if c.Address != "bob@example.com" || c.Next != "/auth/cli" || c.Purpose != ChallengeSignIn {
			t.Fatalf("live read %d = %+v", i, c)
		}
	}
	for _, tc := range []struct {
		why, purpose, flow string
		at                 time.Time
	}{
		{"expired", ChallengeSignIn, testFlow, now.Add(ConfirmationTTL + time.Minute)},
		{"empty flow", ChallengeSignIn, "", now},
		{"another flow", ChallengeSignIn, hashToken("other"), now},
		{"another purpose", ChallengeStepUp, testFlow, now},
	} {
		if _, ok := s.LiveEmailChallenge(tc.purpose, tc.flow, tc.at); ok {
			t.Errorf("%s: a challenge was found", tc.why)
		}
	}
	if _, res := s.TakeEmailCode(live(t, s, testFlow, now), code, now); res != CodeOK {
		t.Fatalf("reading spent the challenge: %v", res)
	}
	if _, ok := s.LiveEmailChallenge(ChallengeSignIn, testFlow, now); ok {
		t.Fatal("a taken challenge is still live")
	}
	// A step-up is found by its session's hash.
	if _, _, err := s.CreateEmailChallenge(EmailChallenge{
		Purpose: ChallengeStepUp, Address: "bob@example.com", SessionHash: "sh", PendingHash: "ph",
	}, now); err != nil {
		t.Fatal(err)
	}
	if c, ok := s.LiveEmailChallenge(ChallengeStepUp, "sh", now); !ok || c.PendingHash != "ph" {
		t.Fatalf("step-up by session = %+v, %v", c, ok)
	}
}

func TestEmailChallengeCode(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()

	link, code := newSignInChallenge(t, s, now, testFlow)
	c := live(t, s, testFlow, now)
	if got, res := s.TakeEmailCode(c, code, now); res != CodeOK || got == nil || got.Address != "bob@example.com" {
		t.Fatalf("right code = %+v, %v", got, res)
	}
	if _, ok := s.TakeEmailLink(link, now); ok {
		t.Fatal("the link outlived a successful code")
	}
	if _, res := s.TakeEmailCode(c, code, now); res != CodeGone {
		t.Fatalf("a code was taken twice: %v", res)
	}
	if _, res := s.TakeEmailCode(nil, code, now); res != CodeGone {
		t.Fatalf("no challenge = %v, want CodeGone", res)
	}

	link, code = newSignInChallenge(t, s, now, testFlow)
	c = live(t, s, testFlow, now)
	bad := wrongCode(code)
	for i := 1; i <= 4; i++ {
		if _, res := s.TakeEmailCode(c, bad, now); res != CodeWrong {
			t.Fatalf("wrong code %d = %v, want CodeWrong", i, res)
		}
	}
	if _, res := s.TakeEmailCode(c, bad, now); res != CodeBurned {
		t.Fatalf("5th wrong code = %v, want CodeBurned", res)
	}
	if _, res := s.TakeEmailCode(c, code, now); res != CodeGone {
		t.Fatalf("right code after burn = %v, want CodeGone", res)
	}
	if _, ok := s.TakeEmailLink(link, now); ok {
		t.Fatal("the link outlived a burned challenge")
	}

	_, code = newSignInChallenge(t, s, now, testFlow)
	c = live(t, s, testFlow, now)
	other := *c
	other.Purpose = ChallengeStepUp
	if _, res := s.TakeEmailCode(&other, code, now); res != CodeGone {
		t.Fatalf("wrong purpose = %v, want CodeGone", res)
	}
	if _, res := s.TakeEmailCode(c, code, now); res != CodeGone {
		t.Fatalf("a wrong-purpose attempt left the challenge standing: %v", res)
	}
}

func countChallenges(t *testing.T, s *Store) int {
	t.Helper()
	var n int
	if err := s.db.QueryRow(`SELECT COUNT(*) FROM email_challenges`).Scan(&n); err != nil {
		t.Fatal(err)
	}
	return n
}

func TestEmailChallengeExpiry(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()
	late := now.Add(ConfirmationTTL + time.Second)

	link, code := newSignInChallenge(t, s, now, testFlow)
	c := live(t, s, testFlow, now)
	if _, ok := s.PeekEmailLink(link, late); ok {
		t.Fatal("an expired link peeks")
	}
	if _, ok := s.TakeEmailLink(link, late); ok {
		t.Fatal("an expired link was taken")
	}
	if n := countChallenges(t, s); n != 0 {
		t.Fatalf("expired link attempt left %d rows", n)
	}
	if _, res := s.TakeEmailCode(c, code, late); res != CodeGone {
		t.Fatalf("code after expiry = %v", res)
	}

	link, code = newSignInChallenge(t, s, now, testFlow)
	c = live(t, s, testFlow, now)
	if _, res := s.TakeEmailCode(c, code, late); res != CodeGone {
		t.Fatalf("expired code = %v, want CodeGone", res)
	}
	if n := countChallenges(t, s); n != 0 {
		t.Fatalf("expired code attempt left %d rows", n)
	}
	if _, ok := s.TakeEmailLink(link, now); ok {
		t.Fatal("the link survived an expired code attempt")
	}
}

func TestEmailChallengeCodeKeyIsProcessLocal(t *testing.T) {
	path := filepath.Join(t.TempDir(), "console.db")
	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now()
	link, code := newSignInChallenge(t, s, now, testFlow)

	// The stored MAC is exactly HMAC-SHA256 under this process's key over the
	// challenge's binding (its link hash) and the code, behind the epoch.
	var mac string
	if err := s.db.QueryRow(`SELECT code_mac FROM email_challenges WHERE link_hash = ?`, hashToken(link)).Scan(&mac); err != nil {
		t.Fatal(err)
	}
	macUnder := func(key []byte) string {
		m := hmac.New(sha256.New, key)
		m.Write([]byte(hashToken(link) + ":" + code))
		return s.codeEpoch + ":" + hex.EncodeToString(m.Sum(nil))
	}
	if want := macUnder(s.codeKey); mac != want {
		t.Fatalf("code_mac = %q, want HMAC under the process key %q", mac, want)
	}
	otherKey := append([]byte(nil), s.codeKey...)
	otherKey[0] ^= 0xff
	if mac == macUnder(otherKey) {
		t.Fatal("a different key gives the same code_mac: the key does not bind it")
	}
	plain := sha256.Sum256([]byte(hashToken(link) + ":" + code))
	if mac == s.codeEpoch+":"+hex.EncodeToString(plain[:]) {
		t.Fatal("code_mac is a plain hash of the binding and code, recoverable from a database read")
	}
	if err := s.Close(); err != nil {
		t.Fatal(err)
	}

	s2, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	defer s2.Close()
	// The link is 32 random bytes hashed in the database, so it survives.
	if c, ok := s2.PeekEmailLink(link, now); !ok || c.Address != "bob@example.com" {
		t.Fatalf("link after restart = %+v, %v", c, ok)
	}
	// The code's key died with the first process: the code reads as used up.
	if _, res := s2.TakeEmailCode(live(t, s2, testFlow, now), code, now); res != CodeGone {
		t.Fatalf("code after restart = %v, want CodeGone (used up)", res)
	}
	if _, ok := s2.TakeEmailLink(link, now); !ok {
		t.Fatal("the link did not survive a restart")
	}
}

// A new challenge replaces the flow's earlier one, link and code alike, and
// leaves other flows alone.
func TestEmailChallengeReplace(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()
	linkA, codeA := newSignInChallenge(t, s, now, testFlow)
	cA := live(t, s, testFlow, now)
	otherFlow := hashToken("another-browser")
	linkOther, _ := newSignInChallenge(t, s, now, otherFlow)
	linkB, _ := newSignInChallenge(t, s, now, testFlow)
	if _, ok := s.PeekEmailLink(linkA, now); ok {
		t.Fatal("the replaced link still works")
	}
	if _, res := s.TakeEmailCode(cA, codeA, now); res != CodeGone {
		t.Fatalf("the replaced code = %v, want CodeGone", res)
	}
	if n := countChallenges(t, s); n != 2 {
		t.Fatalf("%d challenges after a replace, want 2 (this flow's newest and the other flow's)", n)
	}
	if _, ok := s.PeekEmailLink(linkOther, now); !ok {
		t.Fatal("a replace in one flow ended another flow's challenge")
	}
	// B's own link signs in, once.
	if c, ok := s.TakeEmailLink(linkB, now); !ok || c.Address != "bob@example.com" {
		t.Fatalf("the new link = %+v, %v", c, ok)
	}
	if _, ok := s.LiveEmailChallenge(ChallengeSignIn, testFlow, now); ok {
		t.Fatal("B's challenge is still live after its link was used")
	}
}

// A refused request's challenge is recorded only when the flow has no live
// one, and never displaces the email already sent.
func TestEmailChallengeHold(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()
	link, code := newSignInChallenge(t, s, now, testFlow)
	held := EmailChallenge{Purpose: ChallengeSignIn, Address: "carol@example.com", FlowHash: testFlow}
	if err := s.HoldEmailChallenge(held, now); err != nil {
		t.Fatal(err)
	}
	if n := countChallenges(t, s); n != 1 {
		t.Fatalf("%d challenges after a hold beside a live one, want 1", n)
	}
	if c := live(t, s, testFlow, now); c.Address != "bob@example.com" {
		t.Fatalf("the hold displaced the sent challenge: %+v", c)
	}
	if _, ok := s.PeekEmailLink(link, now); !ok {
		t.Fatal("the hold ended the sent link")
	}
	if _, res := s.TakeEmailCode(live(t, s, testFlow, now), code, now); res != CodeOK {
		t.Fatalf("the sent code after a hold = %v", res)
	}
	// A flow with none gets an unsent challenge for the refused address.
	if err := s.HoldEmailChallenge(held, now); err != nil {
		t.Fatal(err)
	}
	c := live(t, s, testFlow, now)
	if c.Address != "carol@example.com" {
		t.Fatalf("held challenge = %+v", c)
	}
	if _, res := s.TakeEmailCode(c, "000000", now); res == CodeGone {
		t.Fatal("a code against the held challenge found nothing, unlike a sent one")
	}
	// An expired one counts as none.
	if err := s.HoldEmailChallenge(held, now.Add(ConfirmationTTL+time.Second)); err != nil {
		t.Fatal(err)
	}
	if n := countChallenges(t, s); n != 1 {
		t.Fatalf("%d challenges after holding past an expired one, want 1", n)
	}
}

// A successful link or code ends every other challenge of the same flow.
func TestEmailChallengeTakeEndsFlow(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()
	for _, by := range []string{"link", "code"} {
		t.Run(by, func(t *testing.T) {
			link, code := newSignInChallenge(t, s, now, testFlow)
			// A second row in the same flow, as an older build might leave.
			if _, err := s.db.Exec(`INSERT INTO email_challenges (id_hash, link_hash, code_mac, purpose, address, flow_hash, created_at, expires_at)
VALUES (?, ?, 'x', ?, 'old@example.com', ?, ?, ?)`, hashToken(by+"-row"), hashToken(by+"-link"), ChallengeSignIn, testFlow, now.Unix()-1, now.Add(time.Minute).Unix()); err != nil {
				t.Fatal(err)
			}
			if by == "link" {
				if _, ok := s.TakeEmailLink(link, now); !ok {
					t.Fatal("link refused")
				}
			} else if _, res := s.TakeEmailCode(live(t, s, testFlow, now), code, now); res != CodeOK {
				t.Fatalf("code = %v", res)
			}
			if n := countChallenges(t, s); n != 0 {
				t.Fatalf("%d of the flow's challenges outlived a sign-in", n)
			}
		})
	}
}

func TestEmailChallengePurgedWithExpired(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()
	newSignInChallenge(t, s, now, testFlow)
	if err := s.PurgeExpired(now); err != nil {
		t.Fatal(err)
	}
	if n := countChallenges(t, s); n != 1 {
		t.Fatalf("PurgeExpired removed a live challenge (%d left)", n)
	}
	if err := s.PurgeExpired(now.Add(ConfirmationTTL + time.Second)); err != nil {
		t.Fatal(err)
	}
	if n := countChallenges(t, s); n != 0 {
		t.Fatalf("PurgeExpired left %d expired challenges", n)
	}
}

func TestSessionMethod(t *testing.T) {
	s := openTestStore(t)
	now := time.Now()

	raw, err := s.CreateSession("bob@example.com", "tok", now)
	if err != nil {
		t.Fatal(err)
	}
	if a, m, ok := s.SessionInfo(raw, now); !ok || a != "bob@example.com" || m != MethodKeymail {
		t.Fatalf("CreateSession info = %q %q %v", a, m, ok)
	}
	raw, err = s.CreateSessionMethod("bob@example.com", "", MethodEmail, now)
	if err != nil {
		t.Fatal(err)
	}
	if a, m, ok := s.SessionInfo(raw, now); !ok || a != "bob@example.com" || m != MethodEmail {
		t.Fatalf("email info = %q %q %v", a, m, ok)
	}
	if _, _, ok := s.SessionInfo(raw, now.Add(SessionTTL+time.Second)); ok {
		t.Fatal("an expired session resolved")
	}
	if _, _, ok := s.SessionInfo("nope", now); ok {
		t.Fatal("an unknown session resolved")
	}
}

func TestSessionMethodPreMigrationRowReadsKeymail(t *testing.T) {
	path := filepath.Join(t.TempDir(), "console.db")
	old, err := sql.Open("sqlite", path)
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now()
	raw := "an-old-session-token"
	for _, q := range []string{
		`CREATE TABLE sessions (
  token_hash    TEXT PRIMARY KEY,
  address       TEXT NOT NULL,
  keymail_token TEXT NOT NULL DEFAULT '',
  created_at    INTEGER NOT NULL,
  expires_at    INTEGER NOT NULL
)`,
	} {
		if _, err := old.Exec(q); err != nil {
			t.Fatal(err)
		}
	}
	if _, err := old.Exec(`INSERT INTO sessions (token_hash, address, keymail_token, created_at, expires_at) VALUES (?, ?, '', ?, ?)`,
		hashToken(raw), "old@example.com", now.Unix(), now.Add(time.Hour).Unix()); err != nil {
		t.Fatal(err)
	}
	if err := old.Close(); err != nil {
		t.Fatal(err)
	}

	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	if a, m, ok := s.SessionInfo(raw, now); !ok || a != "old@example.com" || m != MethodKeymail {
		t.Fatalf("pre-migration session = %q %q %v, want keymail", a, m, ok)
	}
}