package security

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/binary"
	"fmt"
	"strings"
	"testing"
	"time"
)

func TestGenerateSecretLength(t *testing.T) {
	secret, err := GenerateSecret()
	if err != nil {
		t.Fatalf("GenerateSecret: %v", err)
	}
	if len(secret) != totpSecretSize {
		t.Fatalf("secret length = %d, want %d", len(secret), totpSecretSize)
	}
}

func TestTOTPValidateAcceptsCurrentStep(t *testing.T) {
	secret, _ := GenerateSecret()
	now := time.Unix(1_700_000_000, 0)
	step := CurrentStep(now)
	current := code(secret, uint64(step))

	var totp TOTP
	ok, used := totp.Validate(secret, current, now, 0)
	if !ok {
		t.Fatal("Validate rejected a freshly computed current code")
	}
	if used != step {
		t.Fatalf("usedStep = %d, want %d", used, step)
	}
}

func TestTOTPMonotonicReplayGuard(t *testing.T) {
	secret, _ := GenerateSecret()
	now := time.Unix(1_700_000_000, 0)
	step := CurrentStep(now)
	var totp TOTP

	// Consume the current step.
	ok, used := totp.Validate(secret, code(secret, uint64(step)), now, 0)
	if !ok || used != step {
		t.Fatalf("first use failed: ok=%v used=%d", ok, used)
	}

	// Replaying the consumed code is rejected once lastUsedStep is the consumed step.
	if ok, _ := totp.Validate(secret, code(secret, uint64(step)), now, step); ok {
		t.Fatal("replay of the consumed code was accepted")
	}
	// An earlier in window step (step minus one) is also rejected after the consume.
	if ok, _ := totp.Validate(secret, code(secret, uint64(step-1)), now, step); ok {
		t.Fatal("an earlier in window code was accepted after a later success")
	}
	// The next code (step plus one) is the legitimate advance and is accepted, moving the marker.
	ok, used = totp.Validate(secret, code(secret, uint64(step+1)), now, step)
	if !ok || used != step+1 {
		t.Fatalf("legitimate next code rejected: ok=%v used=%d", ok, used)
	}
}

func TestTOTPRejectsSHA1Code(t *testing.T) {
	// The server is pinned to SHA 256. A code derived with SHA 1 for the same secret and step must be
	// rejected, which is what catches an authenticator that silently defaulted to SHA 1.
	secret, _ := GenerateSecret()
	now := time.Unix(1_700_000_000, 0)
	step := CurrentStep(now)

	sha1Code := func(secret []byte, counter uint64) string {
		var buf [8]byte
		binary.BigEndian.PutUint64(buf[:], counter)
		mac := hmac.New(sha1.New, secret)
		mac.Write(buf[:])
		sum := mac.Sum(nil)
		offset := sum[len(sum)-1] & 0x0f
		bin := (uint32(sum[offset])&0x7f)<<24 | uint32(sum[offset+1])<<16 | uint32(sum[offset+2])<<8 | uint32(sum[offset+3])
		return fmt.Sprintf("%06d", bin%totpModulo)
	}

	var totp TOTP
	bad := sha1Code(secret, uint64(step))
	good := code(secret, uint64(step))
	if bad == good {
		t.Skip("sha1 and sha256 codes collided for this secret, skipping")
	}
	if ok, _ := totp.Validate(secret, bad, now, 0); ok {
		t.Fatal("Validate accepted a SHA 1 derived code")
	}
}

func TestProvisioningURIPinsAlgorithm(t *testing.T) {
	secret, _ := GenerateSecret()
	uri := ProvisioningURI(secret, "owner@example.com", "missionMATCH")
	for _, want := range []string{"otpauth://totp/", "algorithm=SHA256", "digits=6", "period=30", "issuer=missionMATCH"} {
		if !strings.Contains(uri, want) {
			t.Fatalf("provisioning URI %q missing %q", uri, want)
		}
	}
}

func TestTOTPRejectsWrongLength(t *testing.T) {
	secret, _ := GenerateSecret()
	var totp TOTP
	if ok, _ := totp.Validate(secret, "123", time.Unix(1_700_000_000, 0), 0); ok {
		t.Fatal("Validate accepted a code of the wrong length")
	}
}