package pack

import (
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"net/http"
	"net/url"
	"sort"
	"strings"

	"github.com/appist/appy/support"
	"github.com/gorilla/securecookie"
)

var (
	mdwCSRFTokenLength                 = 32
	mdwCSRFSecureCookie                *securecookie.SecureCookie
	mdwCSRFSkipCheckCtxKey             = ContextKey("csrfSkipCheck")
	mdwCSRFAuthenticityFieldNameCtxKey = ContextKey("csrfAuthenticityFieldName")
	mdwCSRFAuthenticityTokenCtxKey     = ContextKey("csrfAuthenticityToken")
	mdwCSRFSafeMethods                 = []string{"GET", "HEAD", "OPTIONS", "TRACE"}
	errCSRFNoReferer                   = errors.New("the request referer is missing")
	errCSRFBadReferer                  = errors.New("the request referer is invalid")
	errCSRFNoToken                     = errors.New("the CSRF token is missing")
	errCSRFBadToken                    = errors.New("the CSRF token is invalid")
	generateRandomBytes                = func(n int) ([]byte, error) {
		b := make([]byte, n)
		_, err := rand.Read(b)
		return b, err
	}
)

// CSRFSkipCheck skips the CSRF check for the request.
func CSRFSkipCheck() HandlerFunc {
	return func(c *Context) {
		c.Set(mdwCSRFSkipCheckCtxKey.String(), true)
		c.Next()
	}
}

func mdwCSRF(config *support.Config, logger *support.Logger) HandlerFunc {
	mdwCSRFSecureCookie = securecookie.New(config.HTTPCSRFSecret, nil)
	mdwCSRFSecureCookie.SetSerializer(securecookie.JSONEncoder{})
	mdwCSRFSecureCookie.MaxAge(config.HTTPCSRFCookieMaxAge)

	return func(c *Context) {
		mdwCSRFHandler(c, config, logger)
	}
}

func mdwCSRFHandler(c *Context, config *support.Config, logger *support.Logger) {
	if c.IsAPIOnly() {
		c.Set(mdwCSRFSkipCheckCtxKey.String(), true)
	}

	skipCheck, exists := c.Get(mdwCSRFSkipCheckCtxKey.String())
	if exists && skipCheck.(bool) {
		c.Next()
		return
	}

	csrfToken, err := getCSRFTokenFromCookie(c, config)
	if err != nil || len(csrfToken) != mdwCSRFTokenLength {
		csrfToken, err = generateRandomBytes(mdwCSRFTokenLength)
		if err != nil {
			logger.Error(err)
			c.AbortWithError(http.StatusForbidden, err)
			return
		}

		err = saveCSRFTokenIntoCookie(csrfToken, c, config)
		if err != nil {
			logger.Error(err)
			c.AbortWithError(http.StatusForbidden, err)
			return
		}
	}

	newAuthenticityToken, err := generateAuthenticityToken(csrfToken)
	if err != nil {
		logger.Error(err)
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}

	saveAuthenticityTokenIntoCookie(newAuthenticityToken, c, config)

	c.Set(mdwCSRFAuthenticityTokenCtxKey.String(), newAuthenticityToken)
	c.Set(mdwCSRFAuthenticityFieldNameCtxKey.String(), strings.ToLower(config.HTTPCSRFAuthenticityFieldName))

	r := c.Request
	if !support.ArrayContains(mdwCSRFSafeMethods, r.Method) {
		// Enforce an origin check for HTTPS connections. As per the Django CSRF implementation (https://goo.gl/vKA7GE)
		// the Referer header is almost always present for same-domain HTTP requests.
		if r.TLS != nil {
			referer, err := url.Parse(r.Referer())
			if err != nil || referer.String() == "" {
				logger.Error(errCSRFNoReferer)
				c.AbortWithError(http.StatusForbidden, errCSRFNoReferer)
				return
			}

			if !(referer.Scheme == "https" && referer.Host == r.Host) {
				logger.Error(errCSRFBadReferer)
				c.AbortWithError(http.StatusForbidden, errCSRFBadReferer)
				return
			}
		}

		if csrfToken == nil {
			logger.Error(errCSRFNoToken)
			c.AbortWithError(http.StatusForbidden, errCSRFNoToken)
			return
		}

		csrfTokenFromRequest := getCSRFTokenFromRequest(c, config)
		if !compareTokens(csrfTokenFromRequest, csrfToken) {
			logger.Error(errCSRFBadToken)
			c.AbortWithError(http.StatusForbidden, errCSRFBadToken)
			return
		}
	}

	c.Writer.Header().Add("Vary", "Cookie")
	c.Next()
}

func compareTokens(a, b []byte) bool {
	return subtle.ConstantTimeCompare(a, b) == 1
}

// generateAuthenticityToken returns a unique-per-request token to mitigate the BREACH
// attack as per http://breachattack.com/#mitigations.
//
// The token is generated by XOR'ing a one-time-pad and the base (session) CSRF
// token and returning them together as a 64-byte slice. This effectively
// randomises the token on a per-request basis without breaking multiple browser
// tabs/windows.
func generateAuthenticityToken(csrfToken []byte) (string, error) {
	otp, err := generateRandomBytes(mdwCSRFTokenLength)
	if err != nil {
		return "", err
	}

	// XOR the OTP with the real token to generate a masked token. Append the
	// OTP to the front of the masked token to allow unmasking in the subsequent
	// request.
	return base64.StdEncoding.EncodeToString(append(otp, xorToken(otp, csrfToken)...)), nil
}

// getCSRFTokenFromAuthenticityToken splits the issued token (one-time-pad + masked token)
// and returns the unmasked request CSRF token for comparison.
func getCSRFTokenFromAuthenticityToken(issued []byte) []byte {
	// Issued tokens are always masked and combined with the pad.
	if len(issued) != mdwCSRFTokenLength*2 {
		return nil
	}

	// We now know the length of the byte slice.
	otp := issued[mdwCSRFTokenLength:]
	masked := issued[:mdwCSRFTokenLength]

	// Unmask the token by XOR'ing it against the OTP used to mask it.
	return xorToken(otp, masked)
}

func getCSRFTokenFromCookie(c *Context, config *support.Config) ([]byte, error) {
	encodedToken, err := c.Cookie(config.HTTPCSRFCookieName)
	if err != nil {
		return nil, err
	}

	token := make([]byte, mdwCSRFTokenLength)
	err = mdwCSRFSecureCookie.Decode(config.HTTPCSRFCookieName, encodedToken, &token)
	if err != nil {
		return nil, err
	}

	return token, nil
}

func getCSRFTokenFromRequest(c *Context, config *support.Config) []byte {
	r := c.Request
	fieldName := mdwCSRFAuthenticityTemplateFieldName(c)

	// 1. Check the HTTP header first.
	issued := r.Header.Get(http.CanonicalHeaderKey(config.HTTPCSRFRequestHeader))

	// 2. Fallback to the POST (form) value.
	if issued == "" {
		issued = r.PostFormValue(fieldName)
	}

	// 3. Finally, fallback to the multipart form (if set).
	if issued == "" && r.MultipartForm != nil {
		vals := r.MultipartForm.Value[fieldName]

		if len(vals) > 0 {
			issued = vals[0]
		}
	}

	// Decode the "issued" (pad + masked) token sent in the request. Return a nil
	// byte slice on a decoding error.
	decoded, err := base64.StdEncoding.DecodeString(issued)
	if err != nil {
		return nil
	}

	return getCSRFTokenFromAuthenticityToken(decoded)
}

func mdwCSRFAuthenticityTemplateFieldName(c *Context) string {
	fieldName, exists := c.Get(mdwCSRFAuthenticityFieldNameCtxKey.String())

	if fieldName == "" || !exists {
		fieldName = "authenticity_token"
	}

	return strings.ToLower(fieldName.(string))
}

func saveAuthenticityTokenIntoCookie(authenticityToken string, c *Context, config *support.Config) {
	c.SetSameSite(config.HTTPCSRFCookieSameSite)
	c.SetCookie(
		config.HTTPCSRFAuthenticityFieldName,
		authenticityToken,
		config.HTTPCSRFCookieMaxAge,
		config.HTTPCSRFCookiePath,
		config.HTTPCSRFCookieDomain,
		config.HTTPCSRFCookieSecure,
		false,
	)
}

func saveCSRFTokenIntoCookie(csrfToken []byte, c *Context, config *support.Config) error {
	encodedCSRFToken, err := mdwCSRFSecureCookie.Encode(config.HTTPCSRFCookieName, csrfToken)
	if err != nil {
		return err
	}

	c.SetSameSite(config.HTTPCSRFCookieSameSite)
	c.SetCookie(
		config.HTTPCSRFCookieName,
		encodedCSRFToken,
		config.HTTPCSRFCookieMaxAge,
		config.HTTPCSRFCookiePath,
		config.HTTPCSRFCookieDomain,
		config.HTTPCSRFCookieSecure,
		config.HTTPCSRFCookieHTTPOnly,
	)

	return nil
}

// xorToken XORs tokens ([]byte) to provide unique-per-request CSRF tokens. It
// will return a masked token if the base token is XOR'ed with a one-time-pad.
// An unmasked token will be returned if a masked token is XOR'ed with the
// one-time-pad used to mask it.
func xorToken(a, b []byte) []byte {
	lens := []int{len(a), len(b)}
	sort.Ints(lens)

	n := lens[0]
	res := make([]byte, n)

	for i := 0; i < n; i++ {
		res[i] = a[i] ^ b[i]
	}

	return res
}
