NixMsg/cmd/nixmsg/admin.go

package main

import (
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	"git.asio.asia/nixevol/NixMsg/internal/auth"
	"git.asio.asia/nixevol/NixMsg/internal/config"
	"git.asio.asia/nixevol/NixMsg/internal/store"
)

const (
	adminPasswordLen    = 20
	minAdminPasswordLen = 12
	adminPasswordChars  = "ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz23456789"
)

func cmdAdmin(args []string) error {
	if len(args) < 1 {
		return errors.New("usage: nixmsg admin <init|set-password>")
	}
	switch args[0] {
	case "init":
		return cmdAdminInit(args[1:])
	case "set-password":
		return cmdAdminSetPassword(args[1:])
	default:
		return fmt.Errorf("unknown admin command: %s", args[0])
	}
}

func cmdAdminInit(_ []string) error {
	cfg, err := loadAndValidateConfig()
	if err != nil {
		return err
	}
	ctx := context.Background()
	db, err := store.Open(cfg.DataDir, cfg.SQLiteSynchronous)
	if err != nil {
		return err
	}
	defer func() { _ = db.Close() }()

	ok, err := store.HasAdminPassword(ctx, db.Write)
	if err != nil {
		return err
	}
	if ok {
		return errors.New("admin already initialized; use admin set-password")
	}

	password, err := generateAdminPassword(adminPasswordLen)
	if err != nil {
		return err
	}
	phc, err := auth.HashPassword(password)
	if err != nil {
		return err
	}
	if err := store.SetAdminPasswordHash(ctx, db.Write, phc); err != nil {
		return err
	}
	// 密码只打印到终端一次，不进日志。
	fmt.Printf("admin password: %s\n", password)
	return nil
}

func cmdAdminSetPassword(args []string) error {
	password, err := parseSetPasswordArgs(args, os.Stdin)
	if err != nil {
		return err
	}
	if vErr := validateAdminPassword(password); vErr != nil {
		return vErr
	}
	cfg, err := loadAndValidateConfig()
	if err != nil {
		return err
	}
	ctx := context.Background()
	db, err := store.Open(cfg.DataDir, cfg.SQLiteSynchronous)
	if err != nil {
		return err
	}
	defer func() { _ = db.Close() }()

	phc, err := auth.HashPassword(password)
	if err != nil {
		return err
	}
	if err := store.SetAdminPasswordHash(ctx, db.Write, phc); err != nil {
		return err
	}
	fmt.Fprintln(os.Stderr, "admin password updated")
	return nil
}

func parseSetPasswordArgs(args []string, stdin io.Reader) (string, error) {
	var password string
	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "--password":
			if i+1 >= len(args) {
				return "", errors.New("usage: nixmsg admin set-password --password <password>")
			}
			password = args[i+1]
			i++
		default:
			if strings.HasPrefix(args[i], "-") {
				return "", fmt.Errorf("unknown flag: %s", args[i])
			}
			if password != "" {
				return "", errors.New("usage: nixmsg admin set-password --password <password>")
			}
			password = args[i]
		}
	}
	if password == "" {
		b, err := io.ReadAll(io.LimitReader(stdin, 4096))
		if err != nil {
			return "", err
		}
		password = strings.TrimSpace(string(b))
	}
	if password == "" {
		return "", errors.New("password required (pass --password or stdin)")
	}
	return password, nil
}

func validateAdminPassword(password string) error {
	if len(password) < minAdminPasswordLen {
		return fmt.Errorf("admin password must be at least %d characters", minAdminPasswordLen)
	}
	return nil
}

func generateAdminPassword(n int) (string, error) {
	buf := make([]byte, n)
	charset := []byte(adminPasswordChars)
	for i := 0; i < n; {
		var b [1]byte
		if _, err := rand.Read(b[:]); err != nil {
			return "", err
		}
		if int(b[0]) >= 256-(256%len(charset)) {
			continue
		}
		buf[i] = charset[int(b[0])%len(charset)]
		i++
	}
	return string(buf), nil
}

func loadAndValidateConfig() (config.Config, error) {
	cfgPath := config.PathFromEnv()
	cfg, err := config.Load(cfgPath)
	if err != nil {
		return config.Config{}, err
	}
	if err := cfg.Validate(); err != nil {
		return config.Config{}, err
	}
	return cfg, nil
}