package identity

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"math/big"
	"testing"
	"time"

	"github.com/beevik/etree"
	"github.com/crewjam/saml"
	dsig "github.com/russellhaering/goxmldsig"
)

func TestValidateSAMLMetadataRequiresLiveSigningCertificate(t *testing.T) {
	now := time.Unix(1_800_000_000, 0)
	certificate := makeTestCertificate(t, now.Add(-time.Hour), now.Add(time.Hour))
	metadata := testIDPMetadata(certificate)
	if err := validateSAMLMetadata(context.Background(), metadata, true, now); err != nil {
		t.Fatalf("valid metadata rejected: %v", err)
	}

	metadata.IDPSSODescriptors[0].KeyDescriptors = nil
	if err := validateSAMLMetadata(context.Background(), metadata, true, now); err == nil {
		t.Fatal("metadata without signing certificate was accepted")
	}

	metadata = testIDPMetadata(makeTestCertificate(t, now.Add(-2*time.Hour), now.Add(-time.Hour)))
	if err := validateSAMLMetadata(context.Background(), metadata, true, now); err == nil {
		t.Fatal("expired signing certificate was accepted")
	}

	metadata = testIDPMetadata(certificate)
	metadata.IDPSSODescriptors[0].SingleSignOnServices = nil
	if err := validateSAMLMetadata(context.Background(), metadata, true, now); err == nil {
		t.Fatal("metadata without redirect SSO endpoint was accepted")
	}
}

func TestSAMLEntityIDAndAttributes(t *testing.T) {
	if got, err := validateSAMLEntityID(" urn:example:gateway "); err != nil || got != "urn:example:gateway" {
		t.Fatalf("valid URN entity ID rejected: %q %v", got, err)
	}
	for _, raw := range []string{"", "relative", "https://user:secret@example.com/sp", "urn:example:sp#fragment"} {
		if _, err := validateSAMLEntityID(raw); err == nil {
			t.Fatalf("invalid entity ID accepted: %q", raw)
		}
	}
	assertion := &saml.Assertion{AttributeStatements: []saml.AttributeStatement{{Attributes: []saml.Attribute{{
		FriendlyName: "mail", Name: "urn:oid:0.9.2342.19200300.100.1.3",
		Values: []saml.AttributeValue{{Value: " user@example.com "}},
	}}}}}
	if got := samlAttribute(assertion, "mail"); got != "user@example.com" {
		t.Fatalf("unexpected friendly-name attribute %q", got)
	}
	if got := samlAttribute(assertion, "urn:oid:0.9.2342.19200300.100.1.3"); got != "user@example.com" {
		t.Fatalf("unexpected named attribute %q", got)
	}
}

func TestSAMLSignatureAlgorithmPolicyRejectsSHA1(t *testing.T) {
	document := etree.NewDocument()
	if err := document.ReadFromString(`<Response><Signature><SignedInfo><SignatureMethod Algorithm="` + dsig.RSASHA256SignatureMethod + `"/><Reference><DigestMethod Algorithm="http://www.w3.org/2001/04/xmlenc#sha256"/></Reference></SignedInfo></Signature></Response>`); err != nil {
		t.Fatal(err)
	}
	if err := validateSAMLSignatureAlgorithms(document.Root()); err != nil {
		t.Fatalf("SHA-256 signature policy rejected: %v", err)
	}
	document.Root().FindElement("./Signature/SignedInfo/SignatureMethod").SelectAttr("Algorithm").Value = dsig.RSASHA1SignatureMethod
	if err := validateSAMLSignatureAlgorithms(document.Root()); err == nil {
		t.Fatal("SHA-1 signature algorithm was accepted")
	}
	document.Root().FindElement("./Signature/SignedInfo/SignatureMethod").SelectAttr("Algorithm").Value = dsig.RSASHA256SignatureMethod
	document.Root().FindElement("./Signature/SignedInfo/Reference/DigestMethod").SelectAttr("Algorithm").Value = "http://www.w3.org/2000/09/xmldsig#sha1"
	if err := validateSAMLSignatureAlgorithms(document.Root()); err == nil {
		t.Fatal("SHA-1 digest algorithm was accepted")
	}
}

func testIDPMetadata(certificate *x509.Certificate) *saml.EntityDescriptor {
	return &saml.EntityDescriptor{
		EntityID: "https://idp.example/metadata",
		IDPSSODescriptors: []saml.IDPSSODescriptor{{
			SSODescriptor: saml.SSODescriptor{RoleDescriptor: saml.RoleDescriptor{KeyDescriptors: []saml.KeyDescriptor{{
				Use: "signing", KeyInfo: saml.KeyInfo{X509Data: saml.X509Data{X509Certificates: []saml.X509Certificate{{Data: base64.StdEncoding.EncodeToString(certificate.Raw)}}}},
			}}}},
			SingleSignOnServices: []saml.Endpoint{{Binding: saml.HTTPRedirectBinding, Location: "http://127.0.0.1:9091/sso"}},
		}},
	}
}

func makeTestCertificate(t *testing.T, notBefore, notAfter time.Time) *x509.Certificate {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	template := &x509.Certificate{
		SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "SAML test"},
		NotBefore: notBefore, NotAfter: notAfter, KeyUsage: x509.KeyUsageDigitalSignature,
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	certificate, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatal(err)
	}
	return certificate
}