kanhole/cmd/frpc/sub/auth.go

package sub

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"

	"github.com/spf13/cobra"

	"github.com/fatedier/frp/pkg/util/log"
)

var (
	authServer    string
	authOutput    string
)

func init() {
	authCmd := &cobra.Command{
		Use:   "auth",
		Short: "Authenticate frpc with a frp server",
		Long: `Authenticate this frpc instance with a frp server.

One-time token:
  frpc auth token <token> --server http://server:7500

Interactive login:
  frpc auth login --server http://server:7500 --client-name myclient
`,
	}

	tokenCmd := &cobra.Command{
		Use:   "token <token>",
		Short: "Authenticate using a one-time token",
		Args:  cobra.ExactArgs(1),
		RunE:  runAuthToken,
	}
	tokenCmd.Flags().StringVarP(&authServer, "server", "s", "http://localhost:7500", "frp server admin URL")
	tokenCmd.Flags().StringVarP(&authOutput, "output", "o", "", "output config file path (default: ./frpc-<client-name>.toml)")

	loginCmd := &cobra.Command{
		Use:   "login",
		Short: "Authenticate using admin credentials",
		RunE:  runAuthLogin,
	}
	loginCmd.Flags().StringVarP(&authServer, "server", "s", "http://localhost:7500", "frp server admin URL")
	loginCmd.Flags().StringVarP(&authOutput, "output", "o", "", "output config file path (default: ./frpc-<client-name>.toml)")
	loginCmd.Flags().String("username", "", "admin username (prompts if empty)")
	loginCmd.Flags().String("password", "", "admin password (prompts if empty)")
	loginCmd.Flags().String("client-name", "", "client name (fetches list if empty)")

	rootCmd.AddCommand(authCmd)
	authCmd.AddCommand(tokenCmd)
	authCmd.AddCommand(loginCmd)
}

func runAuthToken(cmd *cobra.Command, args []string) error {
	token := args[0]

	url := authServer + "/admin/api/client/auth"
	body := map[string]string{"token": token}
	data, _ := json.Marshal(body)

	resp, err := http.Post(url, "application/json", bytes.NewReader(data))
	if err != nil {
		return fmt.Errorf("failed to connect to server: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("auth failed (HTTP %d): %s", resp.StatusCode, string(respBody))
	}

	configData, _ := io.ReadAll(resp.Body)
	return saveConfig(configData)
}

func runAuthLogin(cmd *cobra.Command, args []string) error {
	username, _ := cmd.Flags().GetString("username")
	password, _ := cmd.Flags().GetString("password")
	clientName, _ := cmd.Flags().GetString("client-name")

	if username == "" {
		fmt.Print("Admin username: ")
		fmt.Scanln(&username)
	}
	if password == "" {
		fmt.Print("Admin password: ")
		bytePassword, err := readPassword()
		if err != nil {
			return err
		}
		password = string(bytePassword)
		fmt.Println()
	}

	url := authServer + "/admin/api/client/auth"
	body := map[string]string{
		"username":    username,
		"password":    password,
		"client_name": clientName,
	}
	data, _ := json.Marshal(body)

	resp, err := http.Post(url, "application/json", bytes.NewReader(data))
	if err != nil {
		return fmt.Errorf("failed to connect to server: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("auth failed (HTTP %d): %s", resp.StatusCode, string(respBody))
	}

	// Check if server returned a client list
	contentType := resp.Header.Get("Content-Type")
	if contentType == "application/json" || len(contentType) == 0 {
		var result struct {
			Clients           []map[string]any `json:"clients"`
			RequiresClientName bool            `json:"requires_client_name"`
		}
		respBody, _ := io.ReadAll(resp.Body)
		if err := json.Unmarshal(respBody, &result); err == nil && result.RequiresClientName {
			fmt.Println("Available clients:")
			for i, c := range result.Clients {
				fmt.Printf("  %d. %s\n", i+1, c["name"])
			}
			fmt.Print("Enter client name: ")
			var name string
			fmt.Scanln(&name)

			body["client_name"] = name
			data, _ = json.Marshal(body)
			resp, err = http.Post(url, "application/json", bytes.NewReader(data))
			if err != nil {
				return fmt.Errorf("failed to connect to server: %w", err)
			}
			defer resp.Body.Close()
			if resp.StatusCode != http.StatusOK {
				respBody, _ := io.ReadAll(resp.Body)
				return fmt.Errorf("auth failed (HTTP %d): %s", resp.StatusCode, string(respBody))
			}
			configData, _ := io.ReadAll(resp.Body)
			return saveConfig(configData)
		}
		configData := respBody
		return saveConfig(configData)
	}

	configData, _ := io.ReadAll(resp.Body)
	return saveConfig(configData)
}

func saveConfig(data []byte) error {
	outputPath := authOutput
	if outputPath == "" {
		// Try to extract client name from config
		outputPath = "frpc.toml"
	}

	if err := os.MkdirAll(filepath.Dir(outputPath), 0755); err != nil {
		return fmt.Errorf("failed to create output directory: %w", err)
	}

	if err := os.WriteFile(outputPath, data, 0644); err != nil {
		return fmt.Errorf("failed to write config: %w", err)
	}

	log.Infof("config saved to %s", outputPath)
	fmt.Printf("Config saved to %s\n", outputPath)
	fmt.Printf("Run: frpc -c %s\n", outputPath)
	return nil
}

func readPassword() ([]byte, error) {
	return io.ReadAll(os.Stdin)
}