Agent-System-Backend/internal/store/login_keys.go

package store

import (
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"time"
)

// 一次性登录凭据（与长期 api_keys 分离）。明文前缀 alk_；库只存 SHA-256。

const (
	LoginKeyPurposeActivate = "activate"
	LoginKeyPurposeRecover  = "recover"

	LoginKeyStatusUnused   = "unused"
	LoginKeyStatusConsumed = "consumed"
	LoginKeyStatusRevoked  = "revoked"
	LoginKeyStatusExpired  = "expired"

	SessionScopeFull            = "full"
	SessionScopeSetPasswordOnly = "set_password_only"

	defaultLoginKeyTTL   = 30 * time.Minute
	minLoginKeyTTL       = 5 * time.Minute
	maxLoginKeyTTL       = 120 * time.Minute
	restrictedSessionTTL = 10 * time.Minute
)

// LoginKeyMeta 对外可见元数据（无明文）。
type LoginKeyMeta struct {
	ID        string    `json:"id"`
	Owner     string    `json:"owner"`
	Purpose   string    `json:"purpose"`
	Status    string    `json:"status"`
	ExpiresAt time.Time `json:"expiresAt"`
	CreatedAt time.Time `json:"createdAt"`
	CreatedBy string    `json:"createdBy,omitempty"`
}

// IssueLoginKey 签发一次性登录 Key；撤销同 owner+purpose 的 unused 旧 Key。
// 返回明文（仅此一次）与元数据。
func (s *Store) IssueLoginKey(owner, purpose, createdBy string, ttl time.Duration) (plain string, meta LoginKeyMeta, err error) {
	if purpose != LoginKeyPurposeActivate && purpose != LoginKeyPurposeRecover {
		return "", meta, fmt.Errorf("login key: invalid purpose %q", purpose)
	}
	if _, err := s.GetUser(owner); err != nil {
		return "", meta, err
	}
	if ttl <= 0 {
		ttl = defaultLoginKeyTTL
	}
	if ttl < minLoginKeyTTL {
		ttl = minLoginKeyTTL
	}
	if ttl > maxLoginKeyTTL {
		ttl = maxLoginKeyTTL
	}
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", meta, err
	}
	idRaw := make([]byte, 8)
	if _, err := rand.Read(idRaw); err != nil {
		return "", meta, err
	}
	plain = "alk_" + hex.EncodeToString(raw)
	id := "lk_" + hex.EncodeToString(idRaw)
	now := time.Now().UTC()
	exp := now.Add(ttl)
	tx, err := s.db.Begin()
	if err != nil {
		return "", meta, err
	}
	defer tx.Rollback()
	if _, err := tx.Exec(s.q(`UPDATE login_keys SET status=? WHERE owner=? AND purpose=? AND status=?`),
		LoginKeyStatusRevoked, owner, purpose, LoginKeyStatusUnused); err != nil {
		return "", meta, err
	}
	if _, err := tx.Exec(s.q(`INSERT INTO login_keys(id,owner,purpose,key_hash,status,expires_at,created_at,created_by,instance_id) VALUES(?,?,?,?,?,?,?,?,?)`),
		id, owner, purpose, hashKey(plain), LoginKeyStatusUnused, exp.Format(time.RFC3339Nano), now.Format(time.RFC3339Nano), createdBy, ""); err != nil {
		return "", meta, err
	}
	if err := tx.Commit(); err != nil {
		return "", meta, err
	}
	meta = LoginKeyMeta{ID: id, Owner: owner, Purpose: purpose, Status: LoginKeyStatusUnused, ExpiresAt: exp, CreatedAt: now, CreatedBy: createdBy}
	return plain, meta, nil
}

// RevokeLoginKey 按凭据编号撤销未使用 Key。
func (s *Store) RevokeLoginKey(id string) error {
	res, err := s.db.Exec(s.q(`UPDATE login_keys SET status=? WHERE id=? AND status=?`),
		LoginKeyStatusRevoked, id, LoginKeyStatusUnused)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// RedeemLoginKey 原子消费：unused+未过期 → consumed，并创建 set_password_only 会话。
// 明文不得走 Bearer；调用方专用 HTTP 入口。失败整体回滚。
func (s *Store) RedeemLoginKey(plain string) (token string, owner string, err error) {
	if plain == "" || len(plain) < 5 || plain[:4] != "alk_" {
		return "", "", ErrNotFound
	}
	tx, err := s.db.Begin()
	if err != nil {
		return "", "", err
	}
	defer tx.Rollback()
	var id, o, status, exp string
	err = tx.QueryRow(s.q(`SELECT id,owner,status,expires_at FROM login_keys WHERE key_hash=?`), hashKey(plain)).
		Scan(&id, &o, &status, &exp)
	if errors.Is(err, sql.ErrNoRows) {
		return "", "", ErrNotFound
	}
	if err != nil {
		return "", "", err
	}
	expT, err := time.Parse(time.RFC3339Nano, exp)
	if err != nil || time.Now().After(expT) {
		_, _ = tx.Exec(s.q(`UPDATE login_keys SET status=? WHERE id=? AND status=?`), LoginKeyStatusExpired, id, LoginKeyStatusUnused)
		_ = tx.Commit()
		return "", "", ErrNotFound
	}
	if status != LoginKeyStatusUnused {
		return "", "", ErrConflict
	}
	res, err := tx.Exec(s.q(`UPDATE login_keys SET status=?, consumed_at=? WHERE id=? AND status=?`),
		LoginKeyStatusConsumed, nowISO(), id, LoginKeyStatusUnused)
	if err != nil {
		return "", "", err
	}
	if n, _ := res.RowsAffected(); n != 1 {
		return "", "", ErrConflict
	}
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", "", err
	}
	token = "sess_" + hex.EncodeToString(raw)
	sessExp := time.Now().Add(restrictedSessionTTL)
	if _, err := tx.Exec(s.q(`INSERT INTO sessions(token_hash,owner,created_at,expires_at,last_seen_at,scope) VALUES(?,?,?,?,?,?)`),
		hashKey(token), o, nowISO(), sessExp.UTC().Format(time.RFC3339Nano), nowISO(), SessionScopeSetPasswordOnly); err != nil {
		return "", "", err
	}
	if err := tx.Commit(); err != nil {
		return "", "", err
	}
	return token, o, nil
}

// CompleteSetPassword 设密事务：写 hash、清 must_change、作废未用登录 Key、删除全部会话。
func (s *Store) CompleteSetPassword(owner, passwordHash string) error {
	_, err := s.ApplyPasswordChange(AccountWrite{Owner: owner, PasswordHash: passwordHash, Action: "account.set_password"})
	return err
}

// CountUnusedLoginKeys 未使用登录 Key 计数（list 展示）。
func (s *Store) CountUnusedLoginKeys(owner string) (int, error) {
	var n int
	err := s.db.QueryRow(s.q(`SELECT COUNT(1) FROM login_keys WHERE owner=? AND status=?`), owner, LoginKeyStatusUnused).Scan(&n)
	return n, err
}

// RevokeAllAPIKeys 失陷恢复：吊销 owner 全部长期 API Key。
func (s *Store) RevokeAllAPIKeys(owner string) error {
	_, err := s.db.Exec(s.q(`UPDATE api_keys SET revoked=1 WHERE owner=? AND revoked=0`), owner)
	return err
}