package sso

import (
	"bytes"
	"encoding/base64"
	"encoding/xml"
	"testing"
	"time"

	dsig "github.com/russellhaering/goxmldsig"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var testMetadata = `
<?xml version="1.0" encoding="UTF-8"?>
  <md:EntityDescriptor xmlns:md="urn:oasis:names:tc:SAML:2.0:metadata" entityID="https://kolide-dev-ed.my.salesforce.com" validUntil="2027-04-29T19:22:40.750Z" xmlns:ds="http://www.w3.org/2000/09/xmldsig#">
      <md:IDPSSODescriptor WantAuthnRequestsSigned="false" protocolSupportEnumeration="urn:oasis:names:tc:SAML:2.0:protocol">
         <md:KeyDescriptor use="signing">
            <ds:KeyInfo>
               <ds:X509Data>
                  <ds:X509Certificate>MIIErDCCA5SgAwIBAgIOAVuhH3WkAAAAAB5NpvIwDQYJKoZIhvcNAQELBQAwgZAxKDAmBgNVBAMMH1NlbGZTaWduZWRDZXJ0XzI0QXByMjAxN18xODAwNDQxGDAWBgNVBAsMDzAwRDZBMDAwMDAwMTd0ODEXMBUGA1UECgwOU2FsZXNmb3JjZS5jb20xFjAUBgNVBAcMDVNhbiBGcmFuY2lzY28xCzAJBgNVBAgMAkNBMQwwCgYDVQQGEwNVU0EwHhcNMTcwNDI0MTgwMDQ1WhcNMTgwNDI0MTIwMDAwWjCBkDEoMCYGA1UEAwwfU2VsZlNpZ25lZENlcnRfMjRBcHIyMDE3XzE4MDA0NDEYMBYGA1UECwwPMDBENkEwMDAwMDAxN3Q4MRcwFQYDVQQKDA5TYWxlc2ZvcmNlLmNvbTEWMBQGA1UEBwwNU2FuIEZyYW5jaXNjbzELMAkGA1UECAwCQ0ExDDAKBgNVBAYTA1VTQTCCASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAIOR7h8BF2eFOlQHhV/1S7uOBN22Jv7PDCXMz2fU0uLc+mrv9xDGj6ElfW+9dSdXaCbQzD3+Xq4reS4pYRafJZ/27OtygXl3rpoPjSlhRiW+oYVuDcCURJpu0KuZ4I0fm5q1BDYqxcBxNPSe85OHE3+ucmKqvPozhQgYLPCregMIomC3yyANZnLCoGfCv9TpQl6/+I182tST4WPNhVPxKxijoPU4Rh6xY34Ez8+Jr8KdmzmYSNe4ukkIASplpvG7rKka824Hf8zI1BWnjWLDxb5IAxgUBbdr4x8d8C3kPfTf+3/6yC5wSOm9NSs0BA4OJNowtXZFryMzFfXzDzjl69kCAwEAAaOCAQAwgf0wHQYDVR0OBBYEFO+DkoP6qkysi9ZC74yTPuJVVg2yMA8GA1UdEwEB/wQFMAMBAf8wgcoGA1UdIwSBwjCBv4AU74OSg/qqTKyL1kLvjJM+4lVWDbKhgZakgZMwgZAxKDAmBgNVBAMMH1NlbGZTaWduZWRDZXJ0XzI0QXByMjAxN18xODAwNDQxGDAWBgNVBAsMDzAwRDZBMDAwMDAwMTd0ODEXMBUGA1UECgwOU2FsZXNmb3JjZS5jb20xFjAUBgNVBAcMDVNhbiBGcmFuY2lzY28xCzAJBgNVBAgMAkNBMQwwCgYDVQQGEwNVU0GCDgFboR91pAAAAAAeTabyMA0GCSqGSIb3DQEBCwUAA4IBAQAVhYBv5GJvhltks2j7Zc9wdFHW7yB4/hPFo05y0yiOf71tLjOlBucSyxtmXLPjrECJvIJwKhsAIgYXnVp7ditxfauCcxczJgfeL1/dxH/Ge8ePkmH6SdsO71cJL8dXEzOsoF+PAVQzUhqh8zxIipntL0wwNGTD0zIVQeTSozm0KF0SsSHIfbNy279uReGonC61i4Ouk5AMKA7Re9fVeUs6tqM2at22h9Zaj/r/OhXoDcZhzkd8Wq0ER/UKLZA1CyJHgwOC7REEZOuKrqgfWcYt4dGo5q6gqGHHPMv0N7s/MxqCvJCwGA8eJGvOO56I321vhWHQ6ZSJDWUqQFM/Ze7A</ds:X509Certificate>
               </ds:X509Data>
            </ds:KeyInfo>
         </md:KeyDescriptor>
         <md:NameIDFormat>urn:oasis:names:tc:SAML:1.1:nameid-format:unspecified</md:NameIDFormat>
         <md:SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST" Location="https://kolide-dev-ed.my.salesforce.com/idp/endpoint/HttpPost"/>
         <md:SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" Location="https://kolide-dev-ed.my.salesforce.com/idp/endpoint/HttpRedirect"/>
      </md:IDPSSODescriptor>
   </md:EntityDescriptor>
`

func TestNewValidator(t *testing.T) {
	v, err := NewValidator(testMetadata)
	assert.Nil(t, err)
	assert.NotNil(t, v)
}

var testResponse = `[encoded data, 11716 characters]`

func TestValidate(t *testing.T) {
	tm, err := time.Parse(time.UnixDate, "Sun Apr 30 22:10:00 UTC 2017")
	require.Nil(t, err)

	clock := dsig.NewFakeClockAt(tm)
	validator, err := NewValidator(testMetadata, Clock(clock))
	require.Nil(t, err)
	require.NotNil(t, validator)

	auth, err := DecodeAuthResponse(testResponse)
	require.Nil(t, err)

	signed, err := validator.ValidateSignature(auth)
	require.Nil(t, err)
	require.NotNil(t, signed)

	err = validator.ValidateResponse(auth)
	assert.Nil(t, err)
}

func tamperedResponse(original string) (string, error) {
	decoded, err := base64.StdEncoding.DecodeString(original)
	if err != nil {
		return "", err
	}
	var resp Response
	rdr := bytes.NewBuffer(decoded)
	err = xml.NewDecoder(rdr).Decode(&resp)
	if err != nil {
		return "", err
	}
	// change name
	resp.Assertion.Subject.NameID.Value = "bob@kolide.co"
	var wrtr bytes.Buffer
	err = xml.NewEncoder(&wrtr).Encode(resp)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(wrtr.Bytes()), nil
}

func TestVerfiyValidTamperedWithDocFails(t *testing.T) {
	tampered, err := tamperedResponse(testResponse)
	require.Nil(t, err)

	tm, err := time.Parse(time.UnixDate, "Sun Apr 30 22:10:00 UTC 2017")
	require.Nil(t, err)

	clock := dsig.NewFakeClockAt(tm)
	validator, err := NewValidator(testMetadata, Clock(clock))
	require.Nil(t, err)
	require.NotNil(t, validator)

	auth, err := DecodeAuthResponse(tampered)
	require.Nil(t, err)

	_, err = validator.ValidateSignature(auth)
	require.NotNil(t, err)
}

// Message hasn't been tampered with but is stale
func TestVerfiyStaleMessageFails(t *testing.T) {
	tm, err := time.Parse(time.UnixDate, "Sun Apr 30 22:14:00 UTC 2017")
	require.Nil(t, err)
	clock := dsig.NewFakeClockAt(tm)
	validator, err := NewValidator(testMetadata, Clock(clock))
	require.Nil(t, err)
	require.NotNil(t, validator)

	auth, err := DecodeAuthResponse(testResponse)
	require.Nil(t, err)

	signed, err := validator.ValidateSignature(auth)
	require.Nil(t, err)
	require.NotNil(t, signed)

	err = validator.ValidateResponse(auth)
	assert.NotNil(t, err)
}

var testGoogleMetadata = `
<?xml version="1.0" encoding="UTF-8" standalone="no"?>
<md:EntityDescriptor xmlns:md="urn:oasis:names:tc:SAML:2.0:metadata" entityID="https://accounts.google.com/o/saml2?idpid=C0171bstf" validUntil="2022-07-16T20:07:43.000Z">
  <md:IDPSSODescriptor WantAuthnRequestsSigned="false" protocolSupportEnumeration="urn:oasis:names:tc:SAML:2.0:protocol">
    <md:KeyDescriptor use="signing">
      <ds:KeyInfo xmlns:ds="http://www.w3.org/2000/09/xmldsig#">
        <ds:X509Data>
          <ds:X509Certificate>MIIDdDCCAlygAwIBAgIGAV1SKeijMA0GCSqGSIb3DQEBCwUAMHsxFDASBgNVBAoTC0dvb2dsZSBJ
bmMuMRYwFAYDVQQHEw1Nb3VudGFpbiBWaWV3MQ8wDQYDVQQDEwZHb29nbGUxGDAWBgNVBAsTD0dv
b2dsZSBGb3IgV29yazELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWEwHhcNMTcwNzE3
MjAwNzQzWhcNMjIwNzE2MjAwNzQzWjB7MRQwEgYDVQQKEwtHb29nbGUgSW5jLjEWMBQGA1UEBxMN
TW91bnRhaW4gVmlldzEPMA0GA1UEAxMGR29vZ2xlMRgwFgYDVQQLEw9Hb29nbGUgRm9yIFdvcmsx
CzAJBgNVBAYTAlVTMRMwEQYDVQQIEwpDYWxpZm9ybmlhMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8A
MIIBCgKCAQEAzLXNn7VmJBkvVNYHffTzDoow/8eSklauVeYjhELY6dtFv56wAQsFNeMovFUPxPeG
7Fci50/KStvoNZOdKqZFCwYkfI2ssXuMpBP37x2iprV7moVwGdGJb52elMNe0DesgTPbJ/IWIvzF
3GYxqYCHUlHuzJEzBYsdtvM8T/PClBxiLXRNbnjotzleFqb25w3XRfayOZg5GdQPeEmceWXDBhCa
eQyEPOrUTZ+//pZXSuKnOyaFfESNFNgvQJlYQQukjnhPtf674eWT6OdgZHyq8EBbZKfEhs5+KiAN
U43bDh9rpTJCB7rAKk1BFAW3r72pggwN9Z/sfp/C5B7uKAM5hwIDAQABMA0GCSqGSIb3DQEBCwUA
A4IBAQAZXypikbbRzichNXLdK96M/do9nGS5Q3xVgA2uxTzm/6qNkAfOSGSk8OcLrppPonbohkeZ
WVnNB5VZZava4DoSZ6OZsvKc1FM0wKvPJd83KUb7Syk1bV7TkT8DPEclfsLnn5s5g0oHlhsqkNly
0WPFTAoGHXYyOKGEARPoC/o+ZfgfvoMNyZkSQHiRboVVP2cT1ckJt4iCA65hNGXte29hSGmnX7QG
QyrBRp8n4UR9PjoeIy0tTCmG0tqu/NackFH4PkamY84Etxe9uH0StmkhID46QTT4Cv2+jqCaklg+
7VYqXbY64Wc/k0sK7WI1o3IVLWAPNb8ajV6Eo0Y8u+1N</ds:X509Certificate>
        </ds:X509Data>
      </ds:KeyInfo>
    </md:KeyDescriptor>
    <md:NameIDFormat>urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress</md:NameIDFormat>
    <md:SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" Location="https://accounts.google.com/o/saml2/idp?idpid=C0171bstf"/>
    <md:SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST" Location="https://accounts.google.com/o/saml2/idp?idpid=C0171bstf"/>
  </md:IDPSSODescriptor>
</md:EntityDescriptor>

`

func TestVerifyValidGoogleResponse(t *testing.T) {
	samlResponse :=
		`[encoded data, 5996 characters]`
	tm, err := time.Parse(time.RFC3339, "2017-07-18T14:47:08.035Z")
	require.Nil(t, err)
	clock := dsig.NewFakeClockAt(tm)
	validator, err := NewValidator(testGoogleMetadata, Clock(clock))
	require.Nil(t, err)
	require.NotNil(t, validator)
	auth, err := DecodeAuthResponse(samlResponse)
	require.Nil(t, err)
	require.NotNil(t, auth)
	signed, err := validator.ValidateSignature(auth)
	require.Nil(t, err)
	require.NotNil(t, signed)
	err = validator.ValidateResponse(auth)
	assert.Nil(t, err)
}

func TestVerifyInvalidSignatureGoogleResponse(t *testing.T) {
	samlResponse :=
		`[encoded data, 5996 characters]`
	tm, err := time.Parse(time.RFC3339, "2017-07-18T14:47:08.035Z")
	require.Nil(t, err)
	clock := dsig.NewFakeClockAt(tm)
	validator, err := NewValidator(testGoogleMetadata, Clock(clock))
	require.Nil(t, err)
	require.NotNil(t, validator)
	auth, err := DecodeAuthResponse(samlResponse)
	require.Nil(t, err)
	require.NotNil(t, auth)
	_, err = validator.ValidateSignature(auth)
	require.NotNil(t, err)
}

// validate id's are unique and that I didn't screw up my maths
func TestIDGenerator(t *testing.T) {
	idTable := make(map[string]struct{})
	for i := 0; i < 100; i++ {
		id, err := generateSAMLValidID()
		require.Nil(t, err)
		assert.Subset(t, []byte(idAlphabet), []byte(id))
		_, ok := idTable[id]
		assert.False(t, ok)
		idTable[id] = struct{}{}
	}
}

func TestIDPrefix(t *testing.T) {
	// Ensure ID comes with the appropriate prefix
	id, err := generateSAMLValidID()
	require.Nil(t, err)
	assert.Equal(t, idPrefix, id[:2])
}
