姚章浩
90 lines
3.2 KiB
Go
90 lines
3.2 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
cryptorand "crypto/rand"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"math/big"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/settings"
|
|
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/store/postgres"
|
|
)
|
|
|
|
const bootstrapAdminTTL = 10 * time.Minute
|
|
|
|
type bootstrapAdminRepository interface {
|
|
AdministratorCount(context.Context, string) (int, error)
|
|
CreateBootstrapChallenge(context.Context, string, string, time.Time) error
|
|
ConfirmBootstrap(context.Context, string) error
|
|
}
|
|
|
|
var createBootstrapChallenge = func(ctx context.Context, repo bootstrapAdminRepository, tenant, code string, expiry time.Time) error {
|
|
return repo.CreateBootstrapChallenge(ctx, tenant, code, expiry)
|
|
}
|
|
|
|
func runBootstrapAdminCommand(ctx context.Context, output io.Writer) error {
|
|
tenant := strings.TrimSpace(os.Getenv("DOBOTAI_TENANT_KEY"))
|
|
store, err := postgres.Open(ctx, os.Getenv("DATABASE_URL"))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer store.Close()
|
|
repository, err := postgres.NewSettingsRepository(store)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return runBootstrapAdmin(ctx, output, repository, tenant, cryptorand.Reader, time.Now, time.Second)
|
|
}
|
|
|
|
func runBootstrapAdmin(ctx context.Context, output io.Writer, repository bootstrapAdminRepository, tenant string, random io.Reader, now func() time.Time, pollInterval time.Duration) error {
|
|
if repository == nil || strings.TrimSpace(tenant) == "" || random == nil || now == nil || pollInterval <= 0 {
|
|
return errors.New("bootstrap administrator command requires repository, tenant, random source, clock, and poll interval")
|
|
}
|
|
count, err := repository.AdministratorCount(ctx, tenant)
|
|
if err != nil {
|
|
return fmt.Errorf("check active administrators: %w", err)
|
|
}
|
|
if count > 0 {
|
|
return errors.New("an active administrator already exists; use administrator management instead")
|
|
}
|
|
codeNumber, err := cryptorand.Int(random, big.NewInt(1_000_000))
|
|
if err != nil {
|
|
return fmt.Errorf("generate bootstrap code: %w", err)
|
|
}
|
|
code := fmt.Sprintf("%06d", codeNumber.Int64())
|
|
expiresAt := now().Add(bootstrapAdminTTL)
|
|
if err := createBootstrapChallenge(ctx, repository, tenant, code, expiresAt); err != nil {
|
|
return fmt.Errorf("create bootstrap challenge: %w", err)
|
|
}
|
|
fmt.Fprintf(output, "首位管理员绑定已开始。\n绑定码：%s\n请在%s前私聊钉钉机器人发送：/setting %s\n终端正在等待绑定完成；按 Ctrl+C 取消。\n", code, expiresAt.Local().Format("2006-01-02 15:04:05"), code)
|
|
|
|
waitFor := expiresAt.Sub(now())
|
|
if waitFor <= 0 {
|
|
return errors.New("bootstrap binding expiry must be in the future")
|
|
}
|
|
waitCtx, cancel := context.WithTimeout(ctx, waitFor)
|
|
defer cancel()
|
|
ticker := time.NewTicker(pollInterval)
|
|
defer ticker.Stop()
|
|
for {
|
|
err := repository.ConfirmBootstrap(waitCtx, tenant)
|
|
if err == nil {
|
|
fmt.Fprintln(output, "首位管理员绑定成功：候选用户已设为 super_admin。")
|
|
return nil
|
|
}
|
|
if !errors.Is(err, settings.ErrBootstrapCandidatePending) {
|
|
return fmt.Errorf("confirm bootstrap candidate: %w", err)
|
|
}
|
|
select {
|
|
case <-waitCtx.Done():
|
|
return errors.New("bootstrap binding expired or cancelled before a candidate was verified")
|
|
case <-ticker.C:
|
|
}
|
|
}
|
|
}
|