cmd/token.go

package cmd

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"github.com/spf13/cobra"
	"os"
)

const (
	DefaultTokenFile = "/var/run/kmstoken"
)

var tokenCmd = &cobra.Command{
	Use:   "token",
	Short: "generate alibabacloud kms agent ssrf token",
	Run: func(cmd *cobra.Command, args []string) {
		var output string

		if len(args) > 0 {
			output = args[0]
		}
		if output == "" {
			output = DefaultTokenFile
		}

		if fileExists(output) {
			_, _ = fmt.Fprintf(os.Stderr, "ssrf token file already exist: %v\n", output)
			return
		}

		token, err := generateToken(32)
		if err != nil {
			_, _ = fmt.Fprintf(os.Stderr, "genetate rrsf token error: %v\n", err)
			os.Exit(1)
		}

		err = saveTokenToFile(token, output)
		if err != nil {
			_, _ = fmt.Fprintf(os.Stderr, "save rrsf token to file[%s] error: %v\n", output, err)
			os.Exit(1)
		}
	},
}

func generateToken(length int) (string, error) {
	bytes := make([]byte, length)
	_, err := rand.Read(bytes)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(bytes), nil
}

func saveTokenToFile(content, filePath string) error {
	file, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0640)
	if err != nil {
		return fmt.Errorf("create file error:%w", err)
	}
	defer file.Close()

	_, err = file.WriteString(content)
	if err != nil {
		return fmt.Errorf("write file error:%w", err)
	}
	return nil
}