template-monorepo/cmd/totp-test/main.go

// Command totp-test runs an interactive TOTP enrollment and verification flow
// against local Redis + in-memory profile (single process).
//
// Prerequisites:
//
//	make deps-up
//	make setup-dev   # ensure Member.TOTP.SecretKEK is set
//
// Usage:
//
//	make totp-test
//	go run ./cmd/totp-test -f etc/gateway.dev.yaml
package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"net/url"
	"os"
	"strings"

	"gateway/internal/config"
	redislib "gateway/internal/library/redis"
	domusecase "gateway/internal/model/member/domain/usecase"
	memberusecase "gateway/internal/model/member/usecase"

	"github.com/skip2/go-qrcode"
	"github.com/zeromicro/go-zero/core/conf"
)

const (
	stepFlow    = "flow"
	stepEnroll  = "enroll"
	stepConfirm = "confirm"
	stepVerify  = "verify"
	stepStatus  = "status"
	stepDisable = "disable"
)

var (
	configFile = flag.String("f", "etc/gateway.dev.yaml", "config file")
	stepFlag   = flag.String("step", stepFlow, "step: flow, enroll, confirm, verify, status, disable")
	tenantID   = flag.String("tenant", "totp-test", "tenant_id")
	uidFlag    = flag.String("uid", "totp-test-uid", "uid")
	account    = flag.String("account", "totp-test@example.com", "account label shown in Authenticator")
	codeFlag   = flag.String("code", "", "TOTP code (non-interactive confirm/verify)")
	kekFlag    = flag.String("kek", "", "override Member.TOTP.SecretKEK (hex or base64)")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: totp-test [options]\n\n")
		fmt.Fprintf(os.Stderr, "Interactive TOTP test for Google Authenticator / Authy.\n")
		fmt.Fprintf(os.Stderr, "Default step=flow guides enroll → confirm → verify in one process.\n\n")
		fmt.Fprintf(os.Stderr, "Examples:\n")
		fmt.Fprintf(os.Stderr, "  totp-test\n")
		fmt.Fprintf(os.Stderr, "  totp-test -step status\n")
		fmt.Fprintf(os.Stderr, "  totp-test -step confirm -code 482913\n")
		flag.PrintDefaults()
	}
	flag.Parse()

	code, err := run()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
	if code != 0 {
		os.Exit(code)
	}
}

func run() (int, error) {
	step := strings.TrimSpace(*stepFlag)
	if step == "" {
		step = stepFlow
	}

	var c config.Config
	conf.MustLoad(*configFile, &c)
	if c.Redis.Host == "" {
		return 1, fmt.Errorf("totp-test: Redis.Host is empty (run: make deps-up)")
	}

	kek := strings.TrimSpace(*kekFlag)
	if kek == "" {
		kek = strings.TrimSpace(c.Member.TOTP.SecretKEK)
	}
	if kek == "" {
		return 1, fmt.Errorf("totp-test: Member.TOTP.SecretKEK is empty; set it in %s or pass -kek", *configFile)
	}
	c.Member.TOTP.SecretKEK = kek

	ctx := context.Background()
	rds, err := redislib.NewClient(c.Redis)
	if err != nil {
		return 1, fmt.Errorf("totp-test: redis: %w", err)
	}

	mod, err := memberusecase.NewModuleFromParam(memberusecase.ModuleParam{
		Redis:  rds,
		Config: c.Member,
	})
	if err != nil {
		return 1, fmt.Errorf("totp-test: member: %w", err)
	}
	if mod.TOTP == nil {
		return 1, fmt.Errorf("totp-test: TOTP usecase not wired (invalid SecretKEK?)")
	}

	env := &session{
		ctx:    ctx,
		tenant: *tenantID,
		uid:    *uidFlag,
		totp:   mod.TOTP,
		in:     bufio.NewReader(os.Stdin),
	}

	switch step {
	case stepFlow:
		if err := env.runFlow(); err != nil {
			return 1, fmt.Errorf("FAIL: %w", err)
		}
	case stepEnroll:
		if err := env.doEnroll(); err != nil {
			return 1, fmt.Errorf("FAIL: %w", err)
		}
	case stepConfirm:
		if err := env.doConfirm(strings.TrimSpace(*codeFlag)); err != nil {
			return 1, fmt.Errorf("FAIL: %w", err)
		}
	case stepVerify:
		if err := env.doVerify(strings.TrimSpace(*codeFlag)); err != nil {
			return 1, fmt.Errorf("FAIL: %w", err)
		}
	case stepStatus:
		if err := env.doStatus(); err != nil {
			return 1, fmt.Errorf("FAIL: %w", err)
		}
	case stepDisable:
		if err := env.doDisable(); err != nil {
			return 1, fmt.Errorf("FAIL: %w", err)
		}
	default:
		flag.Usage()
		return 2, fmt.Errorf("totp-test: unknown step %q", step)
	}

	fmt.Println("OK")
	return 0, nil
}

type session struct {
	ctx    context.Context
	tenant string
	uid    string
	totp   domusecase.TOTPUseCase
	in     *bufio.Reader
}

func (s *session) runFlow() error {
	status, err := s.totp.Status(s.ctx, s.tenant, s.uid)
	if err != nil {
		return err
	}
	if status.Enrolled {
		fmt.Println("Already enrolled for this tenant/uid.")
		if err := s.doStatus(); err != nil {
			return err
		}
		fmt.Println()
		fmt.Println("Proceeding to verify-only (skip enroll/confirm).")
	} else {
		if err := s.doEnroll(); err != nil {
			return err
		}
		fmt.Println()
		confirmCode, err := s.readCode("Enter the 6-digit code from Google Authenticator to confirm enrollment: ")
		if err != nil {
			return err
		}
		if err := s.doConfirm(confirmCode); err != nil {
			return err
		}
		fmt.Println()
		fmt.Println("Wait for the Authenticator code to refresh (up to 30s), then verify step-up.")
	}

	verifyCode, err := s.readCode("Enter a fresh 6-digit code to verify (step-up): ")
	if err != nil {
		return err
	}
	if err := s.doVerify(verifyCode); err != nil {
		return err
	}

	fmt.Println()
	fmt.Println("Testing replay protection with the same code (should fail)...")
	if err := s.doVerify(verifyCode); err == nil {
		return fmt.Errorf("expected replay failure but verify succeeded")
	}
	fmt.Println("Replay correctly rejected.")
	return nil
}

func (s *session) doEnroll() error {
	start, err := s.totp.StartEnroll(s.ctx, s.tenant, s.uid, *account)
	if err != nil {
		return err
	}

	secret, err := secretFromOtpauthURL(start.OtpauthURL)
	if err != nil {
		return err
	}

	fmt.Println("=== TOTP Enrollment ===")
	fmt.Printf("tenant=%s uid=%s\n", s.tenant, s.uid)
	fmt.Printf("issuer=%s account=%s digits=%d period=%ds expires_in=%ds\n",
		start.Issuer, start.Account, start.Digits, start.PeriodSec, start.ExpiresIn)
	fmt.Println()
	fmt.Println("Option A — scan QR code with Google Authenticator:")
	fmt.Println()
	if err := printTerminalQR(start.OtpauthURL); err != nil {
		fmt.Fprintf(os.Stderr, "warn: QR render failed: %v\n", err)
	}
	fmt.Println()
	fmt.Println("Option B — enter setup key manually in Google Authenticator:")
	fmt.Printf("  Type:       Time based\n")
	fmt.Printf("  Account:    %s\n", start.Account)
	fmt.Printf("  Issuer:     %s\n", start.Issuer)
	fmt.Printf("  Secret key: %s\n", secret)
	fmt.Println()
	fmt.Println("otpauth URL (for debugging):")
	fmt.Println(start.OtpauthURL)
	fmt.Println()
	fmt.Printf("Complete enrollment within %d seconds.\n", start.ExpiresIn)
	return nil
}

func (s *session) doConfirm(code string) error {
	if code == "" {
		var err error
		code, err = s.readCode("Enter the 6-digit code from Google Authenticator: ")
		if err != nil {
			return err
		}
	}
	backup, err := s.totp.ConfirmEnroll(s.ctx, s.tenant, s.uid, code)
	if err != nil {
		return err
	}
	fmt.Println("Enrollment confirmed.")
	fmt.Printf("backup_codes (%d, save these — shown once):\n", len(backup))
	for i, c := range backup {
		fmt.Printf("  [%02d] %s\n", i+1, c)
	}
	return s.doStatus()
}

func (s *session) doVerify(code string) error {
	if code == "" {
		var err error
		code, err = s.readCode("Enter a 6-digit TOTP code (or backup code): ")
		if err != nil {
			return err
		}
	}
	if err := s.totp.VerifyCode(s.ctx, s.tenant, s.uid, code); err != nil {
		return err
	}
	fmt.Println("VerifyCode: success")
	return nil
}

func (s *session) doStatus() error {
	status, err := s.totp.Status(s.ctx, s.tenant, s.uid)
	if err != nil {
		return err
	}
	fmt.Printf("status enrolled=%t backup_codes_remaining=%d", status.Enrolled, status.BackupCodesRemaining)
	if status.Enrolled {
		fmt.Printf(" enrolled_at=%d", status.EnrolledAt)
	}
	fmt.Println()
	return nil
}

func (s *session) doDisable() error {
	if err := s.totp.Disable(s.ctx, s.tenant, s.uid); err != nil {
		return err
	}
	fmt.Println("TOTP disabled.")
	return s.doStatus()
}

func (s *session) readCode(prompt string) (string, error) {
	fmt.Print(prompt)
	line, err := s.in.ReadString('\n')
	if err != nil {
		return "", fmt.Errorf("read code: %w", err)
	}
	code := strings.TrimSpace(line)
	if code == "" {
		return "", fmt.Errorf("code is empty")
	}
	return code, nil
}

func secretFromOtpauthURL(raw string) (string, error) {
	u, err := url.Parse(raw)
	if err != nil {
		return "", fmt.Errorf("parse otpauth url: %w", err)
	}
	secret := strings.TrimSpace(u.Query().Get("secret"))
	if secret == "" {
		return "", fmt.Errorf("otpauth url missing secret parameter")
	}
	return secret, nil
}

func printTerminalQR(content string) error {
	qr, err := qrcode.New(content, qrcode.Medium)
	if err != nil {
		return err
	}
	fmt.Print(qr.ToSmallString(false))
	return nil
}