package jwt

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"regexp"
	"time"

	. "gopkg.in/check.v1"
)

type JWSSuite struct {
	signers []Signer
	hmacKey []byte
	rsaKey  *rsa.PrivateKey
}

var _ = Suite(&JWSSuite{})

func (s *JWSSuite) SetUpSuite(c *C) {
	var err error
	s.hmacKey = NewHMACKey(24)

	s.rsaKey, err = CachedRSAkey()

	c.Assert(err, IsNil)

	s.signers = []Signer{
		NewHMAC256Signer(s.hmacKey),
		NewHMAC384Signer(s.hmacKey),
		NewHMAC512Signer(s.hmacKey),
		NewRSA256Signer(s.rsaKey),
		NewRSA384Signer(s.rsaKey),
		NewRSA512Signer(s.rsaKey),
	}
}

func (s *JWSSuite) TestCanEncodeSigned(c *C) {
	type foo struct {
		A string
		B int
	}

	j := &JOSE{}

	for _, signer := range s.signers {
		decoded := foo{A: "blah", B: 42}
		res, err := EncodeJWS(j, decoded, signer)
		c.Check(string(res), Matches, `\A[0-9A-Za-z_\-]+\.[0-9A-Za-z_\-]+.[0-9A-Za-z_\-]+\z`)
		if c.Check(err, IsNil) == false {
			continue
		}
		redecoded := foo{}
		err = DecodeJWS(res, &redecoded, signer)
		if c.Check(err, IsNil, Commentf("Algo is: %s", signer.Algorithm())) == true {
			c.Check(redecoded, DeepEquals, decoded)
		}
	}
}

func (s *JWSSuite) TestCanEncodeUnprotected(c *C) {

	type foo struct {
		A string
		B int
	}

	j := &JOSE{}

	decoded := foo{A: "blah", B: 42}
	res, err := EncodeJWS(j, decoded, nil)
	c.Check(string(res), Matches, `\A[0-9A-Za-z_\-]+\.[0-9A-Za-z_\-]+.\z`)
	c.Assert(err, IsNil)

	redecoded := foo{}
	err = DecodeJWS(res, &redecoded, nil)
	c.Assert(err, IsNil)
	c.Check(redecoded, DeepEquals, decoded)

}

func (s *JWSSuite) TestAttackerTriesToForgeAToken(c *C) {
	type ident struct {
		Iat int64
		Adm bool
	}

	//--------------------------------------------------------------------------
	// Server side
	//--------------------------------------------------------------------------

	//we create a token and sign it

	validToken, err := EncodeJWS(&JOSE{}, ident{Iat: time.Now().Unix(), Adm: false}, s.signers[0])
	c.Assert(err, IsNil)

	//--------------------------------------------------------------------------
	// Untrusted side
	//--------------------------------------------------------------------------

	// now an attacker takes our valid token, unmarshall it, and create a new one with no signature
	t := ident{}

	tokenRx := regexp.MustCompile(`\A([a-zA-Z0-9_\-]+)\.([a-zA-Z0-9_\-]+).([a-zA-Z0-9_\-]+)\z`)

	matches := tokenRx.FindSubmatch(validToken)
	c.Assert(matches, NotNil)

	// extract JOSE
	j, err := DecodeJOSE(matches[1])
	c.Assert(err, IsNil)

	// extract payload
	payload := make([]byte, Base64DecodedLenFromStripped(len(matches[2])))
	Base64Decode(payload, matches[2])
	err = json.Unmarshal(payload, &t)
	c.Assert(err, IsNil)

	//we forge an admin token
	t.Adm = true

	//we un-sign the token, by using the none algorithm
	forgedToken, err := EncodeJWS(j, t, nil)

	//--------------------------------------------------------------------------
	// Server side
	//--------------------------------------------------------------------------

	//we  verify the forgedToken, it should fail
	received := ident{Adm: false}
	// please note that we speficy ecplitely the signer, and therefore
	// the algorithm to use.
	err = DecodeJWS(forgedToken, &received, s.signers[0])

	//we got a wrong signature for our chosen algorithm
	c.Assert(err, ErrorMatches, fmt.Sprintf(`jws: .* %s .*`, s.signers[0].Algorithm()))
	// of course, our token remain Adm:false
	c.Assert(received.Adm, Equals, false)

}

func (s *JWSSuite) TestAttackerTriesToForgeAHMACToken(c *C) {
	type ident struct {
		Iat int64
		Adm bool
	}

	//--------------------------------------------------------------------------
	// Server side
	//--------------------------------------------------------------------------

	//we create a token and sign it using RSA256
	serverSigner := NewRSA256Signer(s.rsaKey)
	validToken, err := EncodeJWS(&JOSE{}, ident{Iat: time.Now().Unix(), Adm: false}, serverSigner)
	c.Assert(err, IsNil)

	//--------------------------------------------------------------------------
	// Untrusted side
	//--------------------------------------------------------------------------

	// now an attacker takes our valid token, unmarshall it, and create a new one with no signature
	t := ident{}

	tokenRx := regexp.MustCompile(`\A([a-zA-Z0-9_\-]+)\.([a-zA-Z0-9_\-]+).([a-zA-Z0-9_\-]+)\z`)

	matches := tokenRx.FindSubmatch(validToken)
	c.Assert(matches, NotNil)

	// extract JOSE
	j, err := DecodeJOSE(matches[1])
	c.Assert(err, IsNil)

	// extract payload
	payload := make([]byte, Base64DecodedLenFromStripped(len(matches[2])))
	Base64Decode(payload, matches[2])
	err = json.Unmarshal(payload, &t)
	c.Assert(err, IsNil)

	//we forge an admin token
	t.Adm = true

	//now here is the trick, we sign an HMAC token using the serverPublicKey

	publicKeyAsByte, err := x509.MarshalPKIXPublicKey(s.rsaKey.Public())
	c.Assert(err, IsNil, Commentf("While marshalling key"))
	//we create a new signing algorithm expecting the server will
	//still use the assymetric key on the HMAC.
	attackerSigner := NewHMAC256Signer(publicKeyAsByte)
	//we un-sign the token, by using another symmetric algorithm, and
	//the publickKey
	forgedToken, err := EncodeJWS(j, t, attackerSigner)

	//--------------------------------------------------------------------------
	// Server side
	//--------------------------------------------------------------------------

	//we  verify the forgedToken, it should fail
	received := ident{Adm: false}
	// please note that we speficy ecplitely the signer, and therefore
	// the algorithm to use.
	err = DecodeJWS(forgedToken, &received, serverSigner)

	//we got a wrong signature for our chosen algorithm, we still are
	//using the rsa algorithm
	c.Assert(err, ErrorMatches, `crypto/rsa: verification error`)
	// of course, our token remain Adm:false
	c.Assert(received.Adm, Equals, false)

}