// Copyright 2021 The Sigstore Authors.
//
// 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.

package pkcs11key

import (
	"errors"
	"fmt"
	"net/url"
	"strconv"
	"strings"

	"github.com/sigstore/cosign/v3/pkg/cosign/env"
)

const (
	ReferenceScheme = "pkcs11:"
)

var pathAttrValueChars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~:[]@!$'()*+,=&"
var queryAttrValueChars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~:[]@!$'()*+,=/?|"

func percentEncode(input []byte) string {
	if len(input) == 0 {
		return ""
	}

	var stringBuilder strings.Builder
	for i := 0; i < len(input); i++ {
		stringBuilder.WriteByte('%')
		fmt.Fprintf(&stringBuilder, "%.2x", input[i])
	}

	return stringBuilder.String()
}

func EncodeURIComponent(uriString string, isForPath bool, usePercentEncoding bool) (string, error) {
	var stringBuilder strings.Builder
	var allowedChars string

	if isForPath {
		allowedChars = pathAttrValueChars
	} else {
		allowedChars = queryAttrValueChars
	}

	for i := 0; i < len(uriString); i++ {
		allowedChar := false

		for j := 0; j < len(allowedChars); j++ {
			if uriString[i] == allowedChars[j] {
				allowedChar = true
				break
			}
		}

		if allowedChar {
			stringBuilder.WriteByte(uriString[i])
		} else {
			if usePercentEncoding {
				stringBuilder.WriteString(percentEncode([]byte{uriString[i]}))
			} else {
				return "", errors.New("string contains an invalid character")
			}
		}
	}

	return stringBuilder.String(), nil
}

type Pkcs11UriConfig struct {
	uriPathAttributes  url.Values
	uriQueryAttributes url.Values

	ModulePath string
	SlotID     *int
	TokenLabel string
	KeyLabel   []byte
	KeyID      []byte
	Pin        string
}

func NewPkcs11UriConfig() *Pkcs11UriConfig {
	return &Pkcs11UriConfig{
		uriPathAttributes:  make(url.Values),
		uriQueryAttributes: make(url.Values),
	}
}

func NewPkcs11UriConfigFromInput(modulePath string, slotID *int, tokenLabel string, keyLabel []byte, keyID []byte, pin string) *Pkcs11UriConfig {
	return &Pkcs11UriConfig{
		uriPathAttributes:  make(url.Values),
		uriQueryAttributes: make(url.Values),
		ModulePath:         modulePath,
		SlotID:             slotID,
		TokenLabel:         tokenLabel,
		KeyLabel:           keyLabel,
		KeyID:              keyID,
		Pin:                pin,
	}
}

func (conf *Pkcs11UriConfig) Parse(uriString string) error {
	var slotID *int
	var pin string

	uri, err := url.Parse(uriString)
	if err != nil {
		return fmt.Errorf("parse uri: %w", err)
	}
	if uri.Scheme != "pkcs11" {
		return errors.New("invalid uri: not a PKCS11 uri")
	}

	// Semicolons are no longer valid separators, therefore,
	// we need to replace all occurrences of ";" with "&"
	// in uri.Opaque and uri.RawQuery before passing them to url.ParseQuery().
	uri.Opaque = strings.ReplaceAll(uri.Opaque, ";", "&")
	uriPathAttributes, err := url.ParseQuery(uri.Opaque)
	if err != nil {
		return fmt.Errorf("parse uri path: %w", err)
	}
	uri.RawQuery = strings.ReplaceAll(uri.RawQuery, ";", "&")
	uriQueryAttributes, err := url.ParseQuery(uri.RawQuery)
	if err != nil {
		return fmt.Errorf("parse uri query: %w", err)
	}
	modulePath := uriQueryAttributes.Get("module-path")
	pinValue := uriQueryAttributes.Get("pin-value")
	tokenLabel := uriPathAttributes.Get("token")
	slotIDStr := uriPathAttributes.Get("slot-id")
	keyLabel := uriPathAttributes.Get("object")
	keyID := uriPathAttributes.Get("id")

	// At least one of token and slot-id must be specified.
	if tokenLabel == "" && slotIDStr == "" {
		return errors.New("invalid uri: one of token and slot-id must be set")
	}

	// slot-id, if specified, should be a number.
	if slotIDStr != "" {
		slot, err := strconv.Atoi(slotIDStr)
		if err != nil {
			return fmt.Errorf("invalid uri: slot-id '%s' is not a valid number", slotIDStr)
		}
		slotID = &slot
	}

	// If pin-value is specified, take it as it is.
	if pinValue != "" {
		pin = pinValue
	}

	// module-path should be specified and should point to the absolute path of the PKCS11 module.
	// If it is not, COSIGN_PKCS11_MODULE_PATH environment variable must be set.
	if modulePath == "" {
		modulePath = env.Getenv(env.VariablePKCS11ModulePath)
		if modulePath == "" {
			return errors.New("invalid uri: module-path or COSIGN_PKCS11_MODULE_PATH must be set to the absolute path of the PKCS11 module")
		}
	}

	// At least one of object and id must be specified.
	if keyLabel == "" && keyID == "" {
		return errors.New("invalid uri: one of object and id must be set")
	}

	conf.uriPathAttributes = uriPathAttributes
	conf.uriQueryAttributes = uriQueryAttributes
	conf.ModulePath = modulePath
	conf.TokenLabel = tokenLabel
	conf.SlotID = slotID
	conf.KeyLabel = []byte(keyLabel)
	conf.KeyID = []byte(keyID) // url.ParseQuery() already calls url.QueryUnescape() on the id, so we only need to cast the result into byte array
	conf.Pin = pin

	return nil
}

func (conf *Pkcs11UriConfig) Construct() (string, error) {
	var modulePath, pinValue, tokenLabel, slotID, keyID, keyLabel string
	var err error

	uriString := "pkcs11:"

	// module-path should be specified and should point to the absolute path of the PKCS11 module.
	if conf.ModulePath == "" {
		return "", errors.New("module path must be set to the absolute path of the PKCS11 module")
	}

	// At least one of keyLabel and keyID must be specified.
	if len(conf.KeyLabel) == 0 && len(conf.KeyID) == 0 {
		return "", errors.New("one of keyLabel and keyID must be set")
	}

	// At least one of tokenLabel and slotID must be specified.
	if conf.TokenLabel == "" && conf.SlotID == nil {
		return "", errors.New("one of tokenLabel and slotID must be set")
	}

	// Construct the URI.
	if conf.TokenLabel != "" {
		tokenLabel, err = EncodeURIComponent(conf.TokenLabel, true, true)
		if err != nil {
			return "", fmt.Errorf("encode token label: %w", err)
		}
		uriString += "token=" + tokenLabel
	}
	if conf.SlotID != nil {
		slotID = fmt.Sprintf("%d", *conf.SlotID)
		uriString += ";slot-id=" + slotID
	}
	if len(conf.KeyID) != 0 {
		keyID = percentEncode(conf.KeyID)
		uriString += ";id=" + keyID
	}
	if len(conf.KeyLabel) != 0 {
		keyLabel, err = EncodeURIComponent(string(conf.KeyLabel), true, true)
		if err != nil {
			return "", fmt.Errorf("encode key label: %w", err)
		}
		uriString += ";object=" + keyLabel
	}
	modulePath, err = EncodeURIComponent(conf.ModulePath, false, true)
	if err != nil {
		return "", fmt.Errorf("encode module path: %w", err)
	}
	uriString += "?module-path=" + modulePath
	if conf.Pin != "" {
		pinValue, err = EncodeURIComponent(conf.Pin, false, true)
		if err != nil {
			return "", fmt.Errorf("encode pin: %w", err)
		}
		uriString += "&pin-value=" + pinValue
	}

	return uriString, nil
}