optimize the code of the command line (#3614)

cmd/frpc/sub/admin.go

@@ -0,0 +1,117 @@
+// Copyright 2023 The frp 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 sub
+
+import (
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/rodaine/table"
+	"github.com/spf13/cobra"
+
+	"github.com/fatedier/frp/pkg/config"
+	v1 "github.com/fatedier/frp/pkg/config/v1"
+	clientsdk "github.com/fatedier/frp/pkg/sdk/client"
+)
+
+func init() {
+	rootCmd.AddCommand(NewAdminCommand(
+		"reload",
+		"Hot-Reload frpc configuration",
+		ReloadHandler,
+	))
+
+	rootCmd.AddCommand(NewAdminCommand(
+		"status",
+		"Overview of all proxies status",
+		StatusHandler,
+	))
+
+	rootCmd.AddCommand(NewAdminCommand(
+		"stop",
+		"Stop the running frpc",
+		StopHandler,
+	))
+}
+
+func NewAdminCommand(name, short string, handler func(*v1.ClientCommonConfig) error) *cobra.Command {
+	return &cobra.Command{
+		Use:   name,
+		Short: short,
+		Run: func(cmd *cobra.Command, args []string) {
+			cfg, _, _, _, err := config.LoadClientConfig(cfgFile)
+			if err != nil {
+				fmt.Println(err)
+				os.Exit(1)
+			}
+			if cfg.WebServer.Port <= 0 {
+				fmt.Println("web server port should be set if you want to use this feature")
+				os.Exit(1)
+			}
+
+			if err := handler(cfg); err != nil {
+				fmt.Println(err)
+				os.Exit(1)
+			}
+		},
+	}
+}
+
+func ReloadHandler(clientCfg *v1.ClientCommonConfig) error {
+	client := clientsdk.New(clientCfg.WebServer.Addr, clientCfg.WebServer.Port)
+	client.SetAuth(clientCfg.WebServer.User, clientCfg.WebServer.Password)
+	if err := client.Reload(); err != nil {
+		return err
+	}
+	fmt.Println("reload success")
+	return nil
+}
+
+func StatusHandler(clientCfg *v1.ClientCommonConfig) error {
+	client := clientsdk.New(clientCfg.WebServer.Addr, clientCfg.WebServer.Port)
+	client.SetAuth(clientCfg.WebServer.User, clientCfg.WebServer.Password)
+	res, err := client.GetAllProxyStatus()
+	if err != nil {
+		return err
+	}
+
+	fmt.Printf("Proxy Status...\n\n")
+	for _, typ := range proxyTypes {
+		arrs := res[typ]
+		if len(arrs) == 0 {
+			continue
+		}
+
+		fmt.Println(strings.ToUpper(typ))
+		tbl := table.New("Name", "Status", "LocalAddr", "Plugin", "RemoteAddr", "Error")
+		for _, ps := range arrs {
+			tbl.AddRow(ps.Name, ps.Status, ps.LocalAddr, ps.Plugin, ps.RemoteAddr, ps.Err)
+		}
+		tbl.Print()
+		fmt.Println("")
+	}
+	return nil
+}
+
+func StopHandler(clientCfg *v1.ClientCommonConfig) error {
+	client := clientsdk.New(clientCfg.WebServer.Addr, clientCfg.WebServer.Port)
+	client.SetAuth(clientCfg.WebServer.User, clientCfg.WebServer.Password)
+	if err := client.Stop(); err != nil {
+		return err
+	}
+	fmt.Println("stop success")
+	return nil
+}