package http

import (
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/lestrrat-go/jwx/jwa"

	"github.com/lestrrat-go/jwx/jwt"
)

const publicKey = `
-----BEGIN RSA PUBLIC KEY-----
MIIBCgKCAQEAl2ifzOsh6AMRHZBe8xycvk01s3EAGT12WbgV9Z7b420Dj3NXrkns
N/jvBbtO9cjQg4WM7NPZLs+ZutRkHCtMxt7vB0kjOYetLPcGdObsVBB5k1jvwvsJ
HkcfmSsZdrV0Lz2Yxuf6ADWkxBqAY3GsS0zW0A2nIMc+41ZxqsZa3ProsKJxecRX
SSZMpZtqCGt/S83Rek4eAahllcWfZpQmoEk7usLuUl5tH2TmaW3e5lo0JNfdwcq5
PMCa8WSZBFH3YzVttB8rbe7a7336wL2NJQFz6dswL5X1dECYpZ5TRtNgzQYa4V0W
AeICq+EzigaTxrjDHc5urHqEosz1le7O4QIDAQAB
-----END RSA PUBLIC KEY-----
`

const privateKey = `
-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCXaJ/M6yHoAxEd
kF7zHJy+TTWzcQAZPXZZuBX1ntvjbQOPc1euSew3+O8Fu071yNCDhYzs09kuz5m6
1GQcK0zG3u8HSSM5h60s9wZ05uxUEHmTWO/C+wkeRx+ZKxl2tXQvPZjG5/oANaTE
GoBjcaxLTNbQDacgxz7jVnGqxlrc+uiwonF5xFdJJkylm2oIa39LzdF6Th4BqGWV
xZ9mlCagSTu6wu5SXm0fZOZpbd7mWjQk193Byrk8wJrxZJkEUfdjNW20Hytt7trv
ffrAvY0lAXPp2zAvlfV0QJilnlNG02DNBhrhXRYB4gKr4TOKBpPGuMMdzm6seoSi
zPWV7s7hAgMBAAECggEAYoPOxjSP4ThtoIDZZvHNAv2V3WW/HK0jHolqsGBmznmW
AXaZLGwo6NpuG5qea8n38jupUEcfXxfw/OFJKhL6Z8OSX3k1FC+1fDZW2yWNy7zU
fg02I/XXHv5EDxM+BEFYkYxQpcs2nYBJ7tcXhpzl8DDU7JaVkfxSbPVIDEf3wyP2
k6DjYEeAj7uAsp50/32H9zhlJP/cFZaPiyFYy9/gOmDenrPVyJ/f7iQYNwYAwdbt
yfp11Wd5BpePR58+YXICE8oBtzHvB50akK6RZULC3xHVxLQQ1bSxx6vnttxw5RW+
QRHTVWtRyWiKe/l5jMvVSUo5XCLqsL2iXfR4bz6hyQKBgQDJQVWEGHyD6JyaN/6i
5M5+O/YTbzMBgt1JAuVR2c0HYE4LpgrX4cA4kT3Bsa/Z9o0uuWVuxVab5gLsjsDu
EI4o+HJQ3pl4LF9xqrndTdybwmZAT6jv3rM/VGfaCDzXCPzVx169I+WsfkyGW7Tr
Cj4KDZyykruG/9OrpN1Aeq9a3wKBgQDAmCnQHLEPvZdezJGBc34HjZntrbW67iFB
L0waCGWydyunYmzfja1FSvlSoziZdqoq0N4+uBQFPZIlERvq0zMgIzNFvxt/WlFV
kQcBV24MNa8dtd+P7GDY8TzTfYBeXwJoi5c59sWLzSwpTlw0rI+ZvuA66eEF7xih
eWw3k1jOPwKBgQCKf8DHGEbQTEtBQlmlZkrIyqDs/PCgEJwSe8Cu1HFpqxfqokkC
CiTLiQB0BMEdAbRlPEcWtQ2GWgMXIqKY8qGyhk+9YYNCFV9VjQU9zDCOrHjLt0Zu
VNcMNR0HCfY8kb3VrM+A4GxVidFGAWR+/9xz9KwqpBoTrIjRrbJphkSZBwKBgBMs
0zTqNmLH0JNasL3/vrOH0KSOYAKdhOgVinEpFt7+6HTA4vAbDf5RKaOlppP48ZZT
t1ztPOkMqUlRe8MUhgmUF53BGj7CwkhPqS/kAYvrqGS/3+NXeIkA87pmy2oZ8YZx
J3xY6nAx3Ey8hYelCqMXEwIqmQHbPUuOaEzcOcJHAoGANw0TRha2YSbUqS3HiGR/
Jm0lNfeLc3cYlEq0ggRDPLD11485rxVKnaKVHGPYW28OQ4jA+GCc7VZCfPV8VQXW
6b+jUnnBwu/KuYvGMee/xJv1c4MTG54mR9UrUt+R80S0OplpcYkcQnft2Bi+AZ1h
5aZTE7XIouXCYiMKPl4AMtI=
-----END PRIVATE KEY-----
`

func TestJWTProtect(t *testing.T) {
	tests := []struct {
		name               string
		cookie             *http.Cookie
		headers            *http.Header
		server             *httptest.Server
		authorizer         func(r *http.Request, claims map[string]interface{}) error
		expectedStatusCode int
	}{
		{
			"no cookie",
			nil,
			nil,
			nil,
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusForbidden,
		},
		{
			"bad server",
			&http.Cookie{
				Name:  "auth",
				Value: "irrelevantgibberish",
			},
			nil,
			nil,
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusInternalServerError,
		},
		{
			"non-json response",
			&http.Cookie{
				Name:  "auth",
				Value: "irrelevantgibberish",
			},
			nil,
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte("here's some bytes 4 y'all"))
			})),
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusInternalServerError,
		},
		{
			"bad key value",
			&http.Cookie{
				Name:  "auth",
				Value: "irrelevantgibberish",
			},
			nil,
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(`{"keys":["not really a key","me neither"]}`))
			})),
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusForbidden,
		},
		{
			"invalid key",
			&http.Cookie{
				Name:  "auth",
				Value: "irrelevantgibberish",
			},
			nil,
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(`{"keys":["-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCATZAMIIBCgKCAQEA2yUfHH6SRYKvBTemrefi\nHk4L4qkcc4skl4QCaHOkfgA4VcGKG2nXysYuZK7AzNOcHQVi+e4BwN+BfIZtwEU5\n7Ogctb5eg8ksxxLjS7eSRfQIvPGfAbJ12R9OoOWcue/CdUy/YMec4R/o4+tZ45S6\nQQWIMhLqYljw+s1Runda3K8Q8lOdJ4yEZckXaZr1waNJikC7oGpT7ClAgdbvWIbo\nN18G1OluRn+3WNdcN6V+vIj8c9dGs92bgTPX4cn3RmB/80BDfzeFiPMRw5xaq66F\n42zXzllkTqukQPk2wmO5m9pFy0ciRve+awfgbTtZRZOEpTSWLbbpOfd4RQ5YqDWJ\nmQIDAQAB\n-----END PUBLIC KEY-----"]}`))
			})),
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusForbidden,
		},
		{
			"valid key, bad auth token",
			&http.Cookie{
				Name:  "auth",
				Value: "irrelevantgibberish",
			},
			nil,
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(
					fmt.Sprintf(`{"keys":["%s"]}`, strings.ReplaceAll(publicKey, "\n", `\n`)),
				))
			})),
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusForbidden,
		},
		{
			"valid key, valid token",
			&http.Cookie{
				Name: "auth",
				Value: (func() string {
					token := jwt.New()
					token.Set("exp", time.Now().Add(time.Hour))
					keyBytes, _ := pem.Decode([]byte(privateKey))
					privKey, _ := x509.ParsePKCS8PrivateKey(keyBytes.Bytes)
					b, _ := token.Sign(jwa.RS256, privKey)
					return string(b)
				})(),
			},
			nil,
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(
					fmt.Sprintf(`{"keys":["%s"]}`, strings.ReplaceAll(publicKey, "\n", `\n`)),
				))
			})),
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusOK,
		},
		{
			"ok token in headers",
			nil,
			(func() *http.Header {
				token := jwt.New()
				token.Set("exp", time.Now().Add(time.Hour))
				keyBytes, _ := pem.Decode([]byte(privateKey))
				privKey, _ := x509.ParsePKCS8PrivateKey(keyBytes.Bytes)
				b, _ := token.Sign(jwa.RS256, privKey)
				return &http.Header{
					"X-RPC-Authentication": []string{string(b)},
				}
			})(),
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(
					fmt.Sprintf(`{"keys":["%s"]}`, strings.ReplaceAll(publicKey, "\n", `\n`)),
				))
			})),
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusOK,
		},
		{
			"bad token in headers",
			nil,
			(func() *http.Header {
				return &http.Header{
					"X-RPC-Authentication": []string{"nilly willy"},
				}
			})(),
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(
					fmt.Sprintf(`{"keys":["%s"]}`, strings.ReplaceAll(publicKey, "\n", `\n`)),
				))
			})),
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusForbidden,
		},
		{
			"authorizer rejects",
			&http.Cookie{
				Name: "auth",
				Value: (func() string {
					token := jwt.New()
					token.Set("exp", time.Now().Add(time.Hour))
					token.Set("priv", map[string]interface{}{"ok": false})
					keyBytes, _ := pem.Decode([]byte(privateKey))
					privKey, _ := x509.ParsePKCS8PrivateKey(keyBytes.Bytes)
					b, _ := token.Sign(jwa.RS256, privKey)
					return string(b)
				})(),
			},
			nil,
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(
					fmt.Sprintf(`{"keys":["%s"]}`, strings.ReplaceAll(publicKey, "\n", `\n`)),
				))
			})),
			func(r *http.Request, claims map[string]interface{}) error {
				if claims["ok"] == true {
					return nil
				}
				return errors.New("expected ok to be true")
			},
			http.StatusForbidden,
		},
		{
			"valid key, expired token",
			&http.Cookie{
				Name: "auth",
				Value: (func() string {
					token := jwt.New()
					token.Set("exp", time.Now().Add(-time.Hour))
					keyBytes, _ := pem.Decode([]byte(privateKey))
					privKey, _ := x509.ParsePKCS8PrivateKey(keyBytes.Bytes)
					b, _ := token.Sign(jwa.RS256, privKey)
					return string(b)
				})(),
			},
			nil,
			httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(
					fmt.Sprintf(`{"keys":["%s"]}`, strings.ReplaceAll(publicKey, "\n", `\n`)),
				))
			})),
			func(r *http.Request, claims map[string]interface{}) error { return nil },
			http.StatusForbidden,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var url string
			if test.server != nil {
				url = test.server.URL
			}
			wrappedHandler := JWTProtect(url, "auth", "X-RPC-Authentication", test.authorizer, nil)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte("OK"))
			}))
			w := httptest.NewRecorder()
			r := httptest.NewRequest(http.MethodGet, "/", nil)
			if test.cookie != nil {
				r.AddCookie(test.cookie)
			}
			if test.headers != nil {
				for key, value := range *test.headers {
					r.Header.Add(key, value[0])
				}
			}
			wrappedHandler.ServeHTTP(w, r)
			if w.Code != test.expectedStatusCode {
				t.Errorf("Unexpected status code %v", w.Code)
			}
		})
	}
}