package saml

import (
	"crypto/rand"
	"crypto/rsa"
	"testing"
	"time"
)

func generateTestIdPKey(t *testing.T) *rsa.PrivateKey {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate rsa key: %v", err)
	}
	return key
}

func testAssertion(issuer, nameID string, roles []string) AssertionXML {
	now := time.Now()
	return AssertionXML{
		Issuer: issuer,
		NameID: nameID,
		Conditions: conditionsXML{
			NotBefore:    now.Add(-time.Minute).Format(time.RFC3339),
			NotOnOrAfter: now.Add(time.Hour).Format(time.RFC3339),
		},
		Attributes: []attributeXML{
			{Name: "roles", Values: roles},
		},
	}
}

func TestVerify_AcceptsValidSignedAssertion(t *testing.T) {
	key := generateTestIdPKey(t)
	verifier := NewVerifier(&key.PublicKey, "https://idp.example.com")

	rawResp, err := BuildResponse(key, testAssertion("https://idp.example.com", "person@example.com", []string{"sso-admins"}))
	if err != nil {
		t.Fatalf("build response: %v", err)
	}
	resp, err := Parse(rawResp)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}

	assertion, err := verifier.Verify(resp)
	if err != nil {
		t.Fatalf("verify: %v", err)
	}
	if assertion.NameID != "person@example.com" {
		t.Fatalf("nameid = %q", assertion.NameID)
	}
}

func TestVerify_RejectsTamperedAssertion(t *testing.T) {
	key := generateTestIdPKey(t)
	verifier := NewVerifier(&key.PublicKey, "https://idp.example.com")

	rawResp, err := BuildResponse(key, testAssertion("https://idp.example.com", "person@example.com", nil))
	if err != nil {
		t.Fatalf("build response: %v", err)
	}
	resp, err := Parse(rawResp)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	// Manipulation NACH dem Signieren: NameID aendern (Signatur bleibt die alte).
	resp.Assertion.NameID = "angreifer@example.com"

	if _, err := verifier.Verify(resp); err != ErrInvalidSignature {
		t.Fatalf("erwartet ErrInvalidSignature, habe %v", err)
	}
}

func TestVerify_RejectsWrongIdPKey(t *testing.T) {
	realKey := generateTestIdPKey(t)
	attackerKey := generateTestIdPKey(t)
	verifier := NewVerifier(&realKey.PublicKey, "https://idp.example.com")

	rawResp, err := BuildResponse(attackerKey, testAssertion("https://idp.example.com", "person@example.com", nil))
	if err != nil {
		t.Fatalf("build response: %v", err)
	}
	resp, _ := Parse(rawResp)

	if _, err := verifier.Verify(resp); err != ErrInvalidSignature {
		t.Fatalf("erwartet ErrInvalidSignature bei fremd signierter assertion, habe %v", err)
	}
}

func TestVerify_RejectsExpiredAssertion(t *testing.T) {
	key := generateTestIdPKey(t)
	verifier := NewVerifier(&key.PublicKey, "https://idp.example.com")

	expired := testAssertion("https://idp.example.com", "person@example.com", nil)
	expired.Conditions.NotOnOrAfter = time.Now().Add(-time.Hour).Format(time.RFC3339)

	rawResp, err := BuildResponse(key, expired)
	if err != nil {
		t.Fatalf("build response: %v", err)
	}
	resp, _ := Parse(rawResp)

	if _, err := verifier.Verify(resp); err != ErrExpired {
		t.Fatalf("erwartet ErrExpired, habe %v", err)
	}
}

func TestVerify_RejectsWrongIssuer(t *testing.T) {
	key := generateTestIdPKey(t)
	verifier := NewVerifier(&key.PublicKey, "https://idp.example.com")

	rawResp, err := BuildResponse(key, testAssertion("https://anderer-idp.example.com", "person@example.com", nil))
	if err != nil {
		t.Fatalf("build response: %v", err)
	}
	resp, _ := Parse(rawResp)

	if _, err := verifier.Verify(resp); err == nil {
		t.Fatal("erwartet fehler bei unerwartetem aussteller")
	}
}