JCS-pub/jcsctl/cmd/cmd.go

package cmd

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"os"
	"path/filepath"

	"github.com/spf13/cobra"
	"gitlink.org.cn/cloudream/jcs-pub/client/sdk/api"
	cliapi "gitlink.org.cn/cloudream/jcs-pub/client/sdk/api/v1"
)

const (
	defaultCAFileName   = "ca_cert.pem"
	defaultCertFileName = "client_cert.pem"
	defaultKeyFileName  = "client_key.pem"
)

var RootCmd = cobra.Command{}

type CommandContext struct {
	Client *cliapi.Client
	RootCA *x509.CertPool
	Cert   tls.Certificate
}

func GetCmdCtx(cmd *cobra.Command) *CommandContext {
	return cmd.Context().Value("cmdCtx").(*CommandContext)
}

func RootExecute() {
	var ca string
	var cert string
	var key string
	var endpoint string

	RootCmd.Flags().StringVar(&ca, "ca", "", "CA certificate file path")
	RootCmd.Flags().StringVar(&cert, "cert", "", "client certificate file path")
	RootCmd.Flags().StringVar(&key, "key", "", "client key file path")
	RootCmd.Flags().StringVarP(&endpoint, "endpoint", "e", "", "API endpoint")
	RootCmd.MarkFlagsRequiredTogether("ca", "cert", "key")

	if ca == "" {
		certDir := searchCertDir()
		if certDir == "" {
			fmt.Printf("cert files not found, please specify --ca, --cert and --key\n")
			os.Exit(1)
		}

		ca = filepath.Join(certDir, defaultCAFileName)
		cert = filepath.Join(certDir, defaultCertFileName)
		key = filepath.Join(certDir, defaultKeyFileName)
	}

	rootCAPool := x509.NewCertPool()
	rootCAPem, err := os.ReadFile(ca)
	if err != nil {
		fmt.Printf("reading CA file: %v\n", err)
		os.Exit(1)
	}

	if !rootCAPool.AppendCertsFromPEM(rootCAPem) {
		fmt.Printf("parsing CA failed")
		os.Exit(1)
	}

	clientCert, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		fmt.Printf("loading client cert/key: %v\n", err)
		os.Exit(1)
	}

	if endpoint == "" {
		endpoint = "https://127.0.0.1:7890"
	}

	cli := cliapi.NewClient(api.Config{
		EndPoint: endpoint,
		RootCA:   rootCAPool,
		Cert:     clientCert,
	})

	RootCmd.ExecuteContext(context.WithValue(context.Background(), "cmdCtx", &CommandContext{
		Client: cli,
		RootCA: rootCAPool,
		Cert:   clientCert,
	}))
}

func searchCertDir() string {
	execPath, err := os.Executable()
	if err == nil {
		execDir := filepath.Dir(execPath)
		ca, err := os.Stat(filepath.Join(execDir, defaultCAFileName))

		if err == nil && !ca.IsDir() {
			return execDir
		}
	}

	workDir, err := os.Getwd()
	if err == nil {
		ca, err := os.Stat(filepath.Join(workDir, defaultCAFileName))

		if err == nil && !ca.IsDir() {
			return workDir
		}
	}

	return ""
}