cloud-services/pkg/tmobtokengen/encryption.go

package tmobtokengen

import (
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"strings"
	"time"

	"github.com/go-jose/go-jose/v4"
	"github.com/google/uuid"
	"github.com/pkg/errors"
	"github.com/youmark/pkcs8"
)

func buildClientSecret(clientID, secret string) string {
	return base64.StdEncoding.EncodeToString([]byte(clientID + ":" + secret))
}

var (
	ErrFailedToDecodePem = errors.New("failed to decode privateKey")
	ErrNilSigner         = errors.New("signer cannot be nil")
)

func ParsePrivateKey(pemFileBytes []byte, pwd []byte) (any, error) {
	block, _ := pem.Decode(pemFileBytes)
	if block == nil {
		return nil, ErrFailedToDecodePem
	}
	key, _, err := pkcs8.ParsePrivateKey(block.Bytes, pwd)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse pkcs8 private key")
	}

	return key, nil
}

type jwtClaims struct {
	EDTS      string `json:"edts,omitempty"`
	V         string `json:"v,omitempty"`
	Exp       int64  `json:"exp,omitempty"`
	EHTS      string `json:"ehts,omitempty"`
	IAT       int64  `json:"iat,omitempty"`
	UniqueStr string `json:"jti,omitempty"`
}

func buildClaims(
	ehts map[EhtsKey]string,
	curTime time.Time,
	expDuration time.Duration,
	uuidFunc func() string,
) ([]byte, error) {
	curTime = curTime.UTC()
	ehtsKeys, ehtsValues := ehtsToString(ehts)
	encoded := b64EncodeEHTS([]byte(ehtsValues))
	expTime := curTime.Add(expDuration)
	uniqStr := uuidFunc()

	objClaims := jwtClaims{
		EDTS:      string(encoded),
		V:         "1",
		Exp:       expTime.Unix(),
		EHTS:      ehtsKeys,
		IAT:       curTime.Unix(),
		UniqueStr: uniqStr,
	}

	c, err := json.Marshal(objClaims)
	if err != nil {
		return nil, errors.Wrap(err, "failed to marshal claims during build")
	}

	return c, nil
}

func sign(signer jose.Signer, claims []byte) (string, error) {
	if signer == nil {
		return "", ErrNilSigner
	}

	signature, err := signer.Sign(claims)
	if err != nil {
		return "", errors.Wrap(err, "failed to sign claims during pop token build")
	}

	signed, err := signature.CompactSerialize()
	if err != nil {
		return "", errors.Wrap(err, "failed to sign claims during pop token build")
	}

	return signed, nil
}

type Generator interface {
	Generate(ehts EHTSMap) (string, error)
	ClientSecretAsAuthVal() string
}

type PopTokenGenerator struct {
	signer       jose.Signer
	clientSecret string        // client secret composed from client id and secret via base64.StdEncoding
	expDuration  time.Duration // token expiration duration
	genUuid      func() string // uniq string is set as generator field for testing purposes
}

func GenerateUUID() string {
	return uuid.New().String()
}

// NewTokenGenerator creates a new PopTokenGenerator.
func NewTokenGenerator(
	clientID string,
	secret string,
	expDuration time.Duration,
	pkValue string,
	pkPwd ...[]byte,
) (*PopTokenGenerator, error) {
	// Removing the read from file. We read our key from the environment
	pkValue = strings.ReplaceAll(pkValue, "\\n", "\n")
	keyBytes := pkValue

	var pwd []byte
	if len(pkPwd) != 0 {
		pwd = pkPwd[0]
	}

	pk, err := ParsePrivateKey([]byte(keyBytes), pwd)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse private key file")
	}

	opts := new(jose.SignerOptions)
	opts.WithType("JWT")
	signer, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: pk}, opts)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create signer")
	}

	return &PopTokenGenerator{
		signer:       signer,
		clientSecret: buildClientSecret(clientID, secret),
		expDuration:  expDuration,
		genUuid:      GenerateUUID,
	}, nil
}

func (g *PopTokenGenerator) ClientSecretAsAuthVal() string {
	return "Basic " + g.clientSecret
}

// Generate generates a pop token.
// It expects EHTSMap to be populated with the following keys, order is changed intrinsically, so don't worry:
//   - ContentType
//   - Authorization
//   - URI
//   - HTTPMethod
//   - Body
func (g *PopTokenGenerator) Generate(ehts EHTSMap) (string, error) {
	return g.generate(ehts, time.Now())
}

// generate generates a pop token.
// It takes the current time as parameter for testing.
func (g *PopTokenGenerator) generate(ehts EHTSMap, curTime time.Time) (string, error) {
	c, err := buildClaims(ehts, curTime, g.expDuration, g.genUuid)
	if err != nil {
		return "", errors.Wrap(err, "failed to build claims")
	}

	signed, err := sign(g.signer, c)
	if err != nil {
		return "", errors.Wrap(err, "failed to sign claims")
	}

	return signed, nil
}