package autorest

// Copyright 2017 Microsoft Corporation

//

//  Licensed under the Apache License, Version 2.0 (the "License");

//  you may not use this file except in compliance with the License.

//  You may obtain a copy of the License at

//

//      http://www.apache.org/licenses/LICENSE-2.0

//

//  Unless required by applicable law or agreed to in writing, software

//  distributed under the License is distributed on an "AS IS" BASIS,

//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

//  See the License for the specific language governing permissions and

//  limitations under the License.

import (

	"bytes"

	"crypto/hmac"

	"crypto/sha256"

	"encoding/base64"

	"fmt"

	"net/http"

	"net/url"

	"sort"

	"strings"

	"time"

)

// SharedKeyType defines the enumeration for the various shared key types.

// See https://docs.microsoft.com/en-us/rest/api/storageservices/authorize-with-shared-key for details on the shared key types.

type SharedKeyType string

const (

	// SharedKey is used to authorize against blobs, files and queues services.

	SharedKey SharedKeyType = "sharedKey"

	// SharedKeyForTable is used to authorize against the table service.

	SharedKeyForTable SharedKeyType = "sharedKeyTable"

	// SharedKeyLite is used to authorize against blobs, files and queues services.  It's provided for

	// backwards compatibility with API versions before 2009-09-19.  Prefer SharedKey instead.

	SharedKeyLite SharedKeyType = "sharedKeyLite"

	// SharedKeyLiteForTable is used to authorize against the table service.  It's provided for

	// backwards compatibility with older table API versions.  Prefer SharedKeyForTable instead.

	SharedKeyLiteForTable SharedKeyType = "sharedKeyLiteTable"

)

const (

	headerAccept            = "Accept"

	headerAcceptCharset     = "Accept-Charset"

	headerContentEncoding   = "Content-Encoding"

	headerContentLength     = "Content-Length"

	headerContentMD5        = "Content-MD5"

	headerContentLanguage   = "Content-Language"

	headerIfModifiedSince   = "If-Modified-Since"

	headerIfMatch           = "If-Match"

	headerIfNoneMatch       = "If-None-Match"

	headerIfUnmodifiedSince = "If-Unmodified-Since"

	headerDate              = "Date"

	headerXMSDate           = "X-Ms-Date"

	headerXMSVersion        = "x-ms-version"

	headerRange             = "Range"

)

const storageEmulatorAccountName = "devstoreaccount1"

// SharedKeyAuthorizer implements an authorization for Shared Key

// this can be used for interaction with Blob, File and Queue Storage Endpoints

type SharedKeyAuthorizer struct {

	accountName string

	accountKey  []byte

	keyType     SharedKeyType

}

// NewSharedKeyAuthorizer creates a SharedKeyAuthorizer using the provided credentials and shared key type.

func NewSharedKeyAuthorizer(accountName, accountKey string, keyType SharedKeyType) (*SharedKeyAuthorizer, error) {

	key, err := base64.StdEncoding.DecodeString(accountKey)

	if err != nil {

		return nil, fmt.Errorf("malformed storage account key: %v", err)

	}

	return &SharedKeyAuthorizer{

		accountName: accountName,

		accountKey:  key,

		keyType:     keyType,

	}, nil

}

// WithAuthorization returns a PrepareDecorator that adds an HTTP Authorization header whose

// value is "<SharedKeyType> " followed by the computed key.

// This can be used for the Blob, Queue, and File Services

//

// from: https://docs.microsoft.com/en-us/rest/api/storageservices/authorize-with-shared-key

// You may use Shared Key authorization to authorize a request made against the

// 2009-09-19 version and later of the Blob and Queue services,

// and version 2014-02-14 and later of the File services.

func (sk *SharedKeyAuthorizer) WithAuthorization() PrepareDecorator {

	return func(p Preparer) Preparer {

		return PreparerFunc(func(r *http.Request) (*http.Request, error) {

			r, err := p.Prepare(r)

			if err != nil {

				return r, err

			}

			sk, err := buildSharedKey(sk.accountName, sk.accountKey, r, sk.keyType)

			if err != nil {

				return r, err

			}

			return Prepare(r, WithHeader(headerAuthorization, sk))

		})

	}

}

func buildSharedKey(accName string, accKey []byte, req *http.Request, keyType SharedKeyType) (string, error) {

	canRes, err := buildCanonicalizedResource(accName, req.URL.String(), keyType)

	if err != nil {

		return "", err

	}

	if req.Header == nil {

		req.Header = http.Header{}

	}

	// ensure date is set

	if req.Header.Get(headerDate) == "" && req.Header.Get(headerXMSDate) == "" {

		date := time.Now().UTC().Format(http.TimeFormat)

		req.Header.Set(headerXMSDate, date)

	}

	canString, err := buildCanonicalizedString(req.Method, req.Header, canRes, keyType)

	if err != nil {

		return "", err

	}

	return createAuthorizationHeader(accName, accKey, canString, keyType), nil

}

func buildCanonicalizedResource(accountName, uri string, keyType SharedKeyType) (string, error) {

	errMsg := "buildCanonicalizedResource error: %s"

	u, err := url.Parse(uri)

	if err != nil {

		return "", fmt.Errorf(errMsg, err.Error())

	}

	cr := bytes.NewBufferString("")

	if accountName != storageEmulatorAccountName {

		cr.WriteString("/")

		cr.WriteString(getCanonicalizedAccountName(accountName))

	}

	if len(u.Path) > 0 {

		// Any portion of the CanonicalizedResource string that is derived from

		// the resource's URI should be encoded exactly as it is in the URI.

		// -- https://msdn.microsoft.com/en-gb/library/azure/dd179428.aspx

		cr.WriteString(u.EscapedPath())

	}

	params, err := url.ParseQuery(u.RawQuery)

	if err != nil {

		return "", fmt.Errorf(errMsg, err.Error())

	}

	// See https://github.com/Azure/azure-storage-net/blob/master/Lib/Common/Core/Util/AuthenticationUtility.cs#L277

	if keyType == SharedKey {

		if len(params) > 0 {

			cr.WriteString("\n")

			keys := []string{}

			for key := range params {

				keys = append(keys, key)

			}

			sort.Strings(keys)

			completeParams := []string{}

			for _, key := range keys {

				if len(params[key]) > 1 {

					sort.Strings(params[key])

				}

				completeParams = append(completeParams, fmt.Sprintf("%s:%s", key, strings.Join(params[key], ",")))

			}

			cr.WriteString(strings.Join(completeParams, "\n"))

		}

	} else {

		// search for "comp" parameter, if exists then add it to canonicalizedresource

		if v, ok := params["comp"]; ok {

			cr.WriteString("?comp=" + v[0])

		}

	}

	return string(cr.Bytes()), nil

}

func getCanonicalizedAccountName(accountName string) string {

	// since we may be trying to access a secondary storage account, we need to

	// remove the -secondary part of the storage name

	return strings.TrimSuffix(accountName, "-secondary")

}

func buildCanonicalizedString(verb string, headers http.Header, canonicalizedResource string, keyType SharedKeyType) (string, error) {

	contentLength := headers.Get(headerContentLength)

	if contentLength == "0" {

		contentLength = ""

	}

	date := headers.Get(headerDate)

	if v := headers.Get(headerXMSDate); v != "" {

		if keyType == SharedKey || keyType == SharedKeyLite {

			date = ""

		} else {

			date = v

		}

	}

	var canString string

	switch keyType {

	case SharedKey:

		canString = strings.Join([]string{

			verb,

			headers.Get(headerContentEncoding),

			headers.Get(headerContentLanguage),

			contentLength,

			headers.Get(headerContentMD5),

			headers.Get(headerContentType),

			date,

			headers.Get(headerIfModifiedSince),

			headers.Get(headerIfMatch),

			headers.Get(headerIfNoneMatch),

			headers.Get(headerIfUnmodifiedSince),

			headers.Get(headerRange),

			buildCanonicalizedHeader(headers),

			canonicalizedResource,

		}, "\n")

	case SharedKeyForTable:

		canString = strings.Join([]string{

			verb,

			headers.Get(headerContentMD5),

			headers.Get(headerContentType),

			date,

			canonicalizedResource,

		}, "\n")

	case SharedKeyLite:

		canString = strings.Join([]string{

			verb,

			headers.Get(headerContentMD5),

			headers.Get(headerContentType),

			date,

			buildCanonicalizedHeader(headers),

			canonicalizedResource,

		}, "\n")

	case SharedKeyLiteForTable:

		canString = strings.Join([]string{

			date,

			canonicalizedResource,

		}, "\n")

	default:

		return "", fmt.Errorf("key type '%s' is not supported", keyType)

	}

	return canString, nil

}

func buildCanonicalizedHeader(headers http.Header) string {

	cm := make(map[string]string)

	for k := range headers {

		headerName := strings.TrimSpace(strings.ToLower(k))

		if strings.HasPrefix(headerName, "x-ms-") {

			cm[headerName] = headers.Get(k)

		}

	}

	if len(cm) == 0 {

		return ""

	}

	keys := []string{}

	for key := range cm {

		keys = append(keys, key)

	}

	sort.Strings(keys)

	ch := bytes.NewBufferString("")

	for _, key := range keys {

		ch.WriteString(key)

		ch.WriteRune(':')

		ch.WriteString(cm[key])

		ch.WriteRune('\n')

	}

	return strings.TrimSuffix(string(ch.Bytes()), "\n")

}

func createAuthorizationHeader(accountName string, accountKey []byte, canonicalizedString string, keyType SharedKeyType) string {

	h := hmac.New(sha256.New, accountKey)

	h.Write([]byte(canonicalizedString))

	signature := base64.StdEncoding.EncodeToString(h.Sum(nil))

	var key string

	switch keyType {

	case SharedKey, SharedKeyForTable:

		key = "SharedKey"

	case SharedKeyLite, SharedKeyLiteForTable:

		key = "SharedKeyLite"

	}

	return fmt.Sprintf("%s %s:%s", key, getCanonicalizedAccountName(accountName), signature)

}