client_token.go

package shush

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"time"
)

// RegisterToken registers a one-time-use secret token in the agent.
// Returns the token string and a cleanup function.
func RegisterToken(socketPath, secret string, ttl time.Duration, maxUses int) (string, func(), error) {
	return RegisterTokenWithCapability(socketPath, ResolveCapability(nil), secret, ttl, maxUses)
}

// RegisterTokenWithCapability registers a one-time-use secret token with capability.
func RegisterTokenWithCapability(socketPath, capability, secret string, ttl time.Duration, maxUses int) (string, func(), error) {
	if strings.TrimSpace(secret) == "" {
		return "", nil, errors.New("secret must not be empty")
	}
	if ttl <= 0 {
		ttl = defaultTokenTTL
	}
	if maxUses <= 0 {
		maxUses = defaultTokenMaxUses
	}

	if err := startProcess(socketPath, nil, capability); err != nil {
		return "", nil, err
	}

	token, err := generateToken()
	if err != nil {
		return "", nil, err
	}
	client := &Client{SocketPath: socketPath, Capability: capability}
	_, err = client.do(request{
		Action:    actionSetToken,
		Token:     token,
		SecretB64: base64.StdEncoding.EncodeToString([]byte(secret)),
		ExpiresAt: time.Now().Add(ttl).Unix(),
		MaxUses:   maxUses,
	})
	if err != nil {
		return "", nil, err
	}

	cleanup := func() {
		_, _ = client.do(request{
			Action: actionClearToken,
			Token:  token,
		})
	}
	return token, cleanup, nil
}

// ResolveToken retrieves and consumes a one-time token.
func ResolveToken(socketPath, token string) (string, error) {
	return ResolveTokenWithCapability(socketPath, ResolveCapability(nil), token)
}

// ResolveTokenWithCapability retrieves and consumes a token with capability.
func ResolveTokenWithCapability(socketPath, capability, token string) (string, error) {
	client := &Client{SocketPath: socketPath, Capability: capability}
	resp, err := client.do(request{
		Action: actionGetToken,
		Token:  token,
	})
	if err != nil {
		return "", err
	}
	if !resp.Found {
		return "", errors.New("token not found or expired")
	}
	secret, err := base64.StdEncoding.DecodeString(resp.SecretB64)
	if err != nil || len(secret) == 0 {
		return "", errors.New("invalid token secret payload")
	}
	return string(secret), nil
}

func generateToken() (string, error) {
	buf := make([]byte, tokenBytes)
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("generate token: %w", err)
	}
	return hex.EncodeToString(buf), nil
}