package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"
)

const awsAlgorithm = "AWS4-HMAC-SHA256"

func (s *Server) authorizedS3(r *http.Request) bool {
	if len(s.cfg.S3Keys) == 0 {
		return s.authorized(r)
	}

	if r.URL.Query().Get("X-Amz-Algorithm") != "" {
		return s.authorizedS3Presigned(r)
	}

	return s.authorizedS3Header(r)
}

func (s *Server) authorizedS3Presigned(r *http.Request) bool {
	if r.Method != http.MethodGet && r.Method != http.MethodHead {
		return false
	}
	query := r.URL.Query()
	if query.Get("X-Amz-Algorithm") != awsAlgorithm {
		return false
	}
	if query.Get("X-Amz-Signature") == "" {
		return false
	}

	credential := query.Get("X-Amz-Credential")
	accessKey, scope, ok := splitCredential(credential)
	key, ok := s.accessKey(accessKey)
	if !ok || !key.can(permissionForMethod(r.Method)) {
		return false
	}

	amzDate := query.Get("X-Amz-Date")
	if !s.presignedTimeValid(amzDate, query.Get("X-Amz-Expires")) {
		return false
	}
	if !s.scopeAllowed(scope, amzDate) {
		return false
	}

	signedHeaders := query.Get("X-Amz-SignedHeaders")
	canonicalRequest := canonicalRequest(r, signedHeaders, "UNSIGNED-PAYLOAD", true)
	expected := signature(key.Secret, amzDate, scope, canonicalRequest)
	return hmac.Equal([]byte(expected), []byte(query.Get("X-Amz-Signature")))
}

func (s *Server) authorizedS3Header(r *http.Request) bool {
	return s.authorizedS3HeaderPermission(r, permissionForMethod(r.Method))
}

func (s *Server) authorizedS3HeaderPermission(r *http.Request, permission string) bool {
	auth := r.Header.Get("Authorization")
	if !strings.HasPrefix(auth, awsAlgorithm+" ") {
		return false
	}
	parts := parseAuthParams(strings.TrimPrefix(auth, awsAlgorithm+" "))
	credential := parts["Credential"]
	signedHeaders := parts["SignedHeaders"]
	providedSignature := parts["Signature"]
	accessKey, scope, ok := splitCredential(credential)
	key, keyOK := s.accessKey(accessKey)
	if !ok || !keyOK || !key.can(permission) || signedHeaders == "" || providedSignature == "" {
		return false
	}

	amzDate := r.Header.Get("X-Amz-Date")
	if !headerTimeValid(amzDate) {
		return false
	}
	if !s.scopeAllowed(scope, amzDate) {
		return false
	}
	payloadHash := r.Header.Get("X-Amz-Content-Sha256")
	if payloadHash == "" {
		payloadHash = "UNSIGNED-PAYLOAD"
	}
	if permission == "write" && r.Header.Get("X-Magpie-Replication") != "1" && payloadHash == "UNSIGNED-PAYLOAD" && !s.cfg.AllowUnsignedLocal {
		return false
	}

	canonicalRequest := canonicalRequest(r, signedHeaders, payloadHash, false)
	expected := signature(key.Secret, amzDate, scope, canonicalRequest)
	return hmac.Equal([]byte(expected), []byte(providedSignature))
}

func (s *Server) accessKey(id string) (AccessKey, bool) {
	for _, key := range s.cfg.S3Keys {
		if subtle.ConstantTimeCompare([]byte(key.ID), []byte(id)) == 1 {
			return key, true
		}
	}
	return AccessKey{}, false
}

func (s *Server) requestAccessKey(r *http.Request) (AccessKey, bool) {
	credential := r.URL.Query().Get("X-Amz-Credential")
	if credential == "" {
		auth := r.Header.Get("Authorization")
		if !strings.HasPrefix(auth, awsAlgorithm+" ") {
			return AccessKey{}, false
		}
		credential = parseAuthParams(strings.TrimPrefix(auth, awsAlgorithm+" "))["Credential"]
	}
	accessKey, _, ok := splitCredential(credential)
	if !ok {
		return AccessKey{}, false
	}
	return s.accessKey(accessKey)
}

func (key AccessKey) can(permission string) bool {
	return key.Permissions["admin"] || key.Permissions[permission]
}

func permissionForMethod(method string) string {
	switch method {
	case http.MethodGet, http.MethodHead:
		return "read"
	case http.MethodPut, http.MethodPost, http.MethodDelete:
		return "write"
	default:
		return "admin"
	}
}

func signature(secret, amzDate, scope, canonicalRequest string) string {
	date := strings.Split(scope, "/")[0]
	stringToSign := strings.Join([]string{
		awsAlgorithm,
		amzDate,
		scope,
		hexSHA256(canonicalRequest),
	}, "\n")
	key := signingKey(secret, date, scope)
	sig := hmacSHA256(key, []byte(stringToSign))
	return hex.EncodeToString(sig)
}

func signingKey(secret, date, scope string) []byte {
	parts := strings.Split(scope, "/")
	region := "auto"
	service := "s3"
	if len(parts) >= 3 {
		region = parts[1]
		service = parts[2]
	}
	kDate := hmacSHA256([]byte("AWS4"+secret), []byte(date))
	kRegion := hmacSHA256(kDate, []byte(region))
	kService := hmacSHA256(kRegion, []byte(service))
	return hmacSHA256(kService, []byte("aws4_request"))
}

func (s *Server) scopeAllowed(scope string, amzDate string) bool {
	parts := strings.Split(scope, "/")
	if len(parts) != 4 {
		return false
	}
	if len(amzDate) < 8 || parts[0] != amzDate[:8] {
		return false
	}
	if parts[1] != s.cfg.S3Region || parts[2] != "s3" || parts[3] != "aws4_request" {
		return false
	}
	return true
}

func canonicalRequest(r *http.Request, signedHeaders, payloadHash string, presigned bool) string {
	return strings.Join([]string{
		r.Method,
		canonicalURI(r.URL.Path),
		canonicalQuery(r.URL.Query(), presigned),
		canonicalHeaders(r, signedHeaders),
		signedHeaders,
		payloadHash,
	}, "\n")
}

func canonicalURI(path string) string {
	if path == "" {
		return "/"
	}
	segments := strings.Split(path, "/")
	for i, segment := range segments {
		segments[i] = awsEscape(segment)
	}
	return strings.Join(segments, "/")
}

func canonicalQuery(values url.Values, presigned bool) string {
	type pair struct{ key, value string }
	pairs := make([]pair, 0)
	for key, vals := range values {
		if presigned && key == "X-Amz-Signature" {
			continue
		}
		for _, value := range vals {
			pairs = append(pairs, pair{awsEscape(key), awsEscape(value)})
		}
	}
	sort.Slice(pairs, func(i, j int) bool {
		if pairs[i].key == pairs[j].key {
			return pairs[i].value < pairs[j].value
		}
		return pairs[i].key < pairs[j].key
	})
	parts := make([]string, len(pairs))
	for i, pair := range pairs {
		parts[i] = pair.key + "=" + pair.value
	}
	return strings.Join(parts, "&")
}

func canonicalHeaders(r *http.Request, signedHeaders string) string {
	headers := strings.Split(signedHeaders, ";")
	var b strings.Builder
	for _, header := range headers {
		name := strings.ToLower(strings.TrimSpace(header))
		value := ""
		if name == "host" {
			value = r.Host
		} else {
			value = strings.Join(r.Header.Values(name), ",")
		}
		b.WriteString(name)
		b.WriteString(":")
		b.WriteString(compressSpaces(value))
		b.WriteString("\n")
	}
	return b.String()
}

func splitCredential(credential string) (string, string, bool) {
	parts := strings.SplitN(credential, "/", 2)
	if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
		return "", "", false
	}
	scopeParts := strings.Split(parts[1], "/")
	if len(scopeParts) != 4 || scopeParts[3] != "aws4_request" {
		return "", "", false
	}
	return parts[0], parts[1], true
}

func parseAuthParams(value string) map[string]string {
	params := map[string]string{}
	for _, part := range strings.Split(value, ",") {
		keyValue := strings.SplitN(strings.TrimSpace(part), "=", 2)
		if len(keyValue) == 2 {
			params[keyValue[0]] = keyValue[1]
		}
	}
	return params
}

func (s *Server) presignedTimeValid(amzDate, expiresRaw string) bool {
	created, err := time.Parse("20060102T150405Z", amzDate)
	if err != nil {
		return false
	}
	expires, err := time.ParseDuration(expiresRaw + "s")
	if err != nil || expires <= 0 || expires > s.cfg.PresignedMaxExpiry {
		return false
	}
	now := time.Now().UTC()
	return now.After(created.Add(-15*time.Minute)) && now.Before(created.Add(expires))
}

func headerTimeValid(amzDate string) bool {
	created, err := time.Parse("20060102T150405Z", amzDate)
	if err != nil {
		return false
	}
	now := time.Now().UTC()
	return now.After(created.Add(-15*time.Minute)) && now.Before(created.Add(15*time.Minute))
}

func awsEscape(value string) string {
	escaped := url.QueryEscape(value)
	escaped = strings.ReplaceAll(escaped, "+", "%20")
	escaped = strings.ReplaceAll(escaped, "%7E", "~")
	return escaped
}

func compressSpaces(value string) string {
	return strings.Join(strings.Fields(strings.TrimSpace(value)), " ")
}

func hexSHA256(value string) string {
	sum := sha256.Sum256([]byte(value))
	return hex.EncodeToString(sum[:])
}

func hmacSHA256(key []byte, value []byte) []byte {
	mac := hmac.New(sha256.New, key)
	_, _ = mac.Write(value)
	return mac.Sum(nil)
}