NixMsg/test/harness/admin.go

package harness

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"time"
)

// ErrAdminInitUnsupported 表示当前二进制还没有 admin init 子命令。
var ErrAdminInitUnsupported = errors.New("admin init unsupported")

// AdminInitializer 在启动 serve 前初始化管理员密码。
// P1 实现 admin init 后，默认实现会解析终端输出中的密码。
type AdminInitializer interface {
	Init(binPath, configPath string) (password string, err error)
}

// CLIAdminInit 调用 `nixmsg admin init`；若命令不存在则返回 ErrAdminInitUnsupported。
type CLIAdminInit struct{}

// Init 执行 admin init。
func (CLIAdminInit) Init(binPath, configPath string) (string, error) {
	cmd := exec.Command(binPath, "admin", "init")
	cmd.Env = append(os.Environ(), "NIXMSG_CONFIG="+configPath)
	out, err := cmd.CombinedOutput()
	text := string(out)
	if err != nil {
		if isAdminInitUnsupported(text) {
			return "", ErrAdminInitUnsupported
		}
		return "", fmt.Errorf("admin init: %w\n%s", err, text)
	}
	pass := parseAdminPassword(text)
	if pass == "" {
		return "", fmt.Errorf("admin init succeeded but password not found in output:\n%s", text)
	}
	return pass, nil
}

func isAdminInitUnsupported(text string) bool {
	lower := strings.ToLower(text)
	return strings.Contains(lower, "unknown command: admin") ||
		strings.Contains(lower, "unknown command") && strings.Contains(lower, "admin")
}

func parseAdminPassword(out string) string {
	// 约定：P1 实现后密码单独占一行，或出现在 "password:" 之后。
	lines := strings.Split(out, "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		lower := strings.ToLower(line)
		switch {
		case strings.HasPrefix(lower, "password:"):
			return strings.TrimSpace(line[len("password:"):])
		case strings.HasPrefix(lower, "admin password:"):
			return strings.TrimSpace(line[len("admin password:"):])
		}
	}
	for i := len(lines) - 1; i >= 0; i-- {
		line := strings.TrimSpace(lines[i])
		if len(line) >= 12 && !strings.Contains(strings.ToLower(line), "error") {
			fields := strings.Fields(line)
			if len(fields) == 1 {
				return fields[0]
			}
		}
	}
	return ""
}

// AdminClient 管理接口 HTTP 客户端。
// Cookie 会话模式下，改变状态的方法会自动加 X-Nixmsg-Request: 1。
// 设置 APIToken 后走 Bearer，不再要求该头。
type AdminClient struct {
	BaseURL  string
	HTTP     *http.Client
	APIToken string
}

// NewAdminClient 创建带 CookieJar 的管理客户端。baseURL 形如 http://127.0.0.1:12345。
func NewAdminClient(baseURL string) (*AdminClient, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, err
	}
	return &AdminClient{
		BaseURL: strings.TrimRight(baseURL, "/"),
		HTTP: &http.Client{
			Timeout: 30 * time.Second,
			Jar:     jar,
		},
	}, nil
}

// SetSessionCookie 手动写入管理员会话 Cookie（测试辅助）。
func (c *AdminClient) SetSessionCookie(value string) error {
	u, err := url.Parse(c.BaseURL)
	if err != nil {
		return err
	}
	c.HTTP.Jar.SetCookies(u, []*http.Cookie{{
		Name:  "nixmsg_admin",
		Value: value,
		Path:  "/",
	}})
	return nil
}

// Do 发送请求。method 为 POST/PUT/PATCH/DELETE 且未使用 API 令牌时，自动加 X-Nixmsg-Request: 1。
func (c *AdminClient) Do(method, path string, body []byte, contentType string) (*http.Response, error) {
	if !strings.HasPrefix(path, "/") {
		path = "/" + path
	}
	var rdr io.Reader
	if body != nil {
		rdr = bytes.NewReader(body)
	}
	req, err := http.NewRequest(method, c.BaseURL+path, rdr)
	if err != nil {
		return nil, err
	}
	if contentType != "" {
		req.Header.Set("Content-Type", contentType)
	}
	if c.APIToken != "" {
		req.Header.Set("Authorization", "Bearer "+c.APIToken)
	} else if isMutatingMethod(method) {
		req.Header.Set("X-Nixmsg-Request", "1")
	}
	return c.HTTP.Do(req)
}

// Get JSON GET。
func (c *AdminClient) Get(path string) (*http.Response, error) {
	return c.Do(http.MethodGet, path, nil, "")
}

// PostJSON POST application/json。
func (c *AdminClient) PostJSON(path string, body []byte) (*http.Response, error) {
	return c.Do(http.MethodPost, path, body, "application/json")
}

func isMutatingMethod(method string) bool {
	switch strings.ToUpper(method) {
	case http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete:
		return true
	default:
		return false
	}
}