package api_security

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"sort"
	"strings"

	"github.com/gin-gonic/gin"
)

type signReader struct {
	bytes.Reader
	RC io.ReadCloser
}

func (thiz *signReader) Close() error {
	return thiz.RC.Close()
}

func calcSign(kv map[string][]byte) string {
	keys := make([]string, 0, len(kv))
	for k := range kv {
		keys = append(keys, k)
	}

	sort.Strings(keys)

	h := sha256.New()
	for _, k := range keys {
		_, _ = h.Write([]byte(k))
		_, _ = h.Write(kv[k])
	}

	return strings.ToLower(hex.EncodeToString(h.Sum(nil)))
}

func ProtoSign(secret []byte,
	shouldCrypto ShouldCryptoChecker) gin.HandlerFunc {
	return func(ctx *gin.Context) {

		if shouldCrypto != nil && !shouldCrypto(ctx) {
			ctx.Next()
			return
		}

		sign := ctx.GetHeader(XProtoSign)
		nonce := ctx.GetHeader(XProtoNonce)
		ts := ctx.GetHeader(XProtoTimestamp)

		if len(sign) != 64 || len(nonce) < 10 || len(ts) < 6 {
			ctx.AbortWithStatusJSON(500, "protosign: argument invalid")
			return
		}

		data, e := ioutil.ReadAll(ctx.Request.Body)
		if e != nil {
			ctx.AbortWithStatusJSON(500, "protosign: read data failed")
			return
		}

		kv := make(map[string][]byte)
		kv[XProtoNonce] = []byte(nonce)
		kv[XProtoTimestamp] = []byte(ts)
		kv["_METHOD_"] = []byte(ctx.Request.Method)
		kv["_PATH_"] = []byte(ctx.Request.URL.Path)
		kv["_SECRET_"] = secret

		queries := map[string][]string(ctx.Request.URL.Query())
		for k, v := range queries {
			if len(v) > 0 {
				kv[fmt.Sprintf("_QUERY_%s_", strings.ToUpper(k))] = []byte(v[0])
			}
		}

		if len(data) > 0 {
			kv["_BODY_"] = []byte(data)
		}

		if calcSign(kv) != sign {
			ctx.AbortWithStatusJSON(500, "protosign: verify sign failed")
			return
		}

		ctx.Request.Body = &signReader{
			*bytes.NewReader(data),
			ctx.Request.Body,
		}
		ctx.Next()
	}
}