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

package store

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"database/sql"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"strings"
	"time"

	"golang.org/x/crypto/argon2"
)

// ---- 认证与租户（阶段 1：API key 认证，user 即租户 owner）----
// key 明文仅创建时返回一次；库中只存 SHA-256 摘要（数据可靠性标准：不落明文凭证）。

// User 是租户账户（owner 全局唯一；id 供后续扩展用户目录）。
type User struct {
	ID        string    `json:"id"`
	Owner     string    `json:"owner"`
	Name      string    `json:"name"`
	CreatedAt time.Time `json:"createdAt"`
}

// APIKeyInfo 是密钥的展示视图（不含摘要）。
type APIKeyInfo struct {
	Name       string     `json:"name"`
	Owner      string     `json:"owner"`
	CreatedAt  time.Time  `json:"createdAt"`
	LastUsedAt *time.Time `json:"lastUsedAt,omitempty"`
	Revoked    bool       `json:"revoked"`
}

func hashKey(plain string) string {
	sum := sha256.Sum256([]byte(plain))
	return hex.EncodeToString(sum[:])
}

// EnsureUser 幂等创建租户账户（owner 唯一）。
func (s *Store) EnsureUser(owner, name string) error {
	_, err := s.db.Exec(s.q(`INSERT INTO users(id,owner,name,created_at) VALUES(?,?,?,?)
		ON CONFLICT(owner) DO NOTHING`), newID("usr"), owner, name, nowISO())
	return err
}

func (s *Store) GetUser(owner string) (User, error) {
	var u User
	var ca string
	err := s.db.QueryRow(s.q(`SELECT id,owner,name,created_at FROM users WHERE owner=?`), owner).Scan(&u.ID, &u.Owner, &u.Name, &ca)
	if errors.Is(err, sql.ErrNoRows) {
		return u, ErrNotFound
	}
	if err != nil {
		return u, err
	}
	u.CreatedAt, _ = parseTime(ca)
	return u, nil
}

func (s *Store) ListUsers() ([]User, error) {
	rows, err := s.db.Query(s.q(`SELECT id,owner,name,created_at FROM users ORDER BY created_at`))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []User{}
	for rows.Next() {
		var u User
		var ca string
		if err := rows.Scan(&u.ID, &u.Owner, &u.Name, &ca); err != nil {
			return nil, err
		}
		u.CreatedAt, _ = parseTime(ca)
		out = append(out, u)
	}
	return out, rows.Err()
}

// CreateAPIKey 创建密钥：明文 `ask_<hex32>` 仅此一次返回。
func (s *Store) CreateAPIKey(owner, name string) (string, error) {
	if err := s.EnsureUser(owner, ""); err != nil {
		return "", err
	}
	raw := make([]byte, 16)
	if _, err := rand.Read(raw); err != nil {
		return "", err
	}
	plain := "ask_" + hex.EncodeToString(raw)
	if name == "" {
		name = "key-" + nowISO()[:10]
	}
	_, err := s.db.Exec(s.q(`INSERT INTO api_keys(key_hash,owner,name,created_at,revoked) VALUES(?,?,?,?,0)`),
		hashKey(plain), owner, name, nowISO())
	if err != nil {
		return "", err
	}
	return plain, nil
}

// RevokeAPIKeyByName 撤销指定名称的未撤销密钥（返回是否命中）。
func (s *Store) RevokeAPIKeyByName(owner, name string) (bool, error) {
	res, err := s.db.Exec(s.q(`UPDATE api_keys SET revoked=1 WHERE owner=? AND name=? AND revoked=0`), owner, name)
	if err != nil {
		return false, err
	}
	n, _ := res.RowsAffected()
	return n > 0, nil
}

func (s *Store) ListAPIKeys(owner string) ([]APIKeyInfo, error) {
	rows, err := s.db.Query(s.q(`SELECT name,owner,created_at,last_used_at,revoked FROM api_keys WHERE owner=? ORDER BY created_at`), owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []APIKeyInfo{}
	for rows.Next() {
		var k APIKeyInfo
		var ca string
		var last sql.NullString
		if err := rows.Scan(&k.Name, &k.Owner, &ca, &last, &k.Revoked); err != nil {
			return nil, err
		}
		k.CreatedAt, _ = parseTime(ca)
		if last.Valid && last.String != "" {
			t, _ := parseTime(last.String)
			k.LastUsedAt = &t
		}
		out = append(out, k)
	}
	return out, rows.Err()
}

// AuthOwner 校验 API key：命中返回 owner；未命中/已撤销返回 false。
// last_used_at 尽力刷新（失败不影响认证结果）。
func (s *Store) AuthOwner(plainKey string) (string, bool) {
	if plainKey == "" {
		return "", false
	}
	var owner string
	var revoked int
	err := s.db.QueryRow(s.q(`SELECT owner,revoked FROM api_keys WHERE key_hash=?`), hashKey(plainKey)).Scan(&owner, &revoked)
	if err != nil || revoked == 1 {
		return "", false
	}
	s.db.Exec(s.q(`UPDATE api_keys SET last_used_at=? WHERE key_hash=?`), nowISO(), hashKey(plainKey))
	return owner, true
}

// ---- 账号密码与会话（阶段 2：管理员发放制，2026-09-30 裁定）----
// 无自助注册：账号由管理员（或 CLI 引导）创建；初始/重置密码后 must_change_password=1，
// 首次登录强制改密。会话为服务端状态：库存 SHA-256 摘要，cookie 只发一次明文。

const sessionTTL = 7 * 24 * time.Hour

// AdminUser 管理员用户列表的展示视图（不含哈希）。
type AdminUser struct {
	Owner              string     `json:"owner"`
	Name               string     `json:"name"`
	IsAdmin            bool       `json:"isAdmin"`
	MustChangePassword bool       `json:"mustChangePassword"`
	HasPassword        bool       `json:"hasPassword"`
	PasswordSetAt      *time.Time `json:"passwordSetAt,omitempty"`
	CreatedAt          time.Time  `json:"createdAt"`
}

// HashPassword 生成 argon2id PHC 串（m=64MiB,t=2,p=2；随机盐）。
func HashPassword(password string) (string, error) {
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return "", err
	}
	key := argon2.IDKey([]byte(password), salt, 2, 64*1024, 2, 32)
	return "$argon2id$v=19$m=65536,t=2,p=2$" +
		base64.RawStdEncoding.EncodeToString(salt) + "$" +
		base64.RawStdEncoding.EncodeToString(key), nil
}

// VerifyPassword 校验 PHC 串（常量时间比较；坏格式按不匹配处理）。
func VerifyPassword(hash, password string) bool {
	parts := strings.Split(hash, "$")
	if len(parts) != 6 || parts[1] != "argon2id" {
		return false
	}
	var m uint32 = 65536
	var t uint32 = 2
	var p uint8 = 2
	for _, kv := range strings.Split(parts[3], ",") {
		kv2 := strings.SplitN(kv, "=", 2)
		if len(kv2) != 2 {
			continue
		}
		var v int
		for _, c := range kv2[1] {
			v = v*10 + int(c-'0')
		}
		switch kv2[0] {
		case "m":
			m = uint32(v)
		case "t":
			t = uint32(v)
		case "p":
			p = uint8(v)
		}
	}
	salt, err := base64.RawStdEncoding.DecodeString(parts[4])
	if err != nil {
		return false
	}
	want, err := base64.RawStdEncoding.DecodeString(parts[5])
	if err != nil {
		return false
	}
	got := argon2.IDKey([]byte(password), salt, t, m, p, uint32(len(want)))
	return subtle.ConstantTimeCompare(got, want) == 1
}

// CreateAccount 管理员建号：owner 唯一；初始密码 must_change=1。
func (s *Store) CreateAccount(owner, name, passwordHash string, isAdmin bool) error {
	if _, err := s.GetUser(owner); err == nil {
		return ErrConflict
	}
	if err := s.EnsureUser(owner, name); err != nil {
		return err
	}
	_, err := s.db.Exec(s.q(`UPDATE users SET name=?, password_hash=?, password_set_at=?, must_change_password=1, is_admin=? WHERE owner=?`),
		name, passwordHash, nowISO(), boolInt(isAdmin), owner)
	return err
}

// SetUserPassword 设置/重置密码：must_change=1（首登强制改密）。
func (s *Store) SetUserPassword(owner, passwordHash string) error {
	res, err := s.db.Exec(s.q(`UPDATE users SET password_hash=?, password_set_at=?, must_change_password=1 WHERE owner=?`),
		passwordHash, nowISO(), owner)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// ClearMustChangePassword 改密完成后清标志（改密者会话保留，其余吊销由调用方处理）。
func (s *Store) ClearMustChangePassword(owner string) error {
	_, err := s.db.Exec(s.q(`UPDATE users SET must_change_password=0 WHERE owner=?`), owner)
	return err
}

// UserAuthState 登录与身份视图。
type UserAuthState struct {
	Owner              string
	Name               string
	IsAdmin            bool
	MustChangePassword bool
	HasPassword        bool
}

// GetUserAuthState 按 owner 取认证态（含会话/APIOwner 路径共用）。
func (s *Store) GetUserAuthState(owner string) (UserAuthState, error) {
	var st UserAuthState
	var hash, ca string
	var mustChange, isAdmin int
	err := s.db.QueryRow(s.q(`SELECT owner,COALESCE(name,''),password_hash,must_change_password,is_admin,created_at FROM users WHERE owner=?`), owner).
		Scan(&st.Owner, &st.Name, &hash, &mustChange, &isAdmin, &ca)
	if errors.Is(err, sql.ErrNoRows) {
		return st, ErrNotFound
	}
	if err != nil {
		return st, err
	}
	st.MustChangePassword = mustChange == 1
	st.IsAdmin = isAdmin == 1
	st.HasPassword = hash != ""
	return st, nil
}

// VerifyLogin 校验用户名密码。
func (s *Store) VerifyLogin(owner, password string) (UserAuthState, bool) {
	var hash string
	err := s.db.QueryRow(s.q(`SELECT password_hash FROM users WHERE owner=?`), owner).Scan(&hash)
	if err != nil || hash == "" || !VerifyPassword(hash, password) {
		return UserAuthState{}, false
	}
	st, err := s.GetUserAuthState(owner)
	if err != nil {
		return UserAuthState{}, false
	}
	return st, true
}

// ---- 会话 ----

// CreateSession 建会话：明文 token 仅此一次返回；库存 SHA-256 摘要与到期时间。
func (s *Store) CreateSession(owner string) (string, error) {
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", err
	}
	token := "sess_" + hex.EncodeToString(raw)
	exp := time.Now().Add(sessionTTL)
	_, err := s.db.Exec(s.q(`INSERT INTO sessions(token_hash,owner,created_at,expires_at,last_seen_at) VALUES(?,?,?,?,?)`),
		hashKey(token), owner, nowISO(), exp.UTC().Format(time.RFC3339Nano), nowISO())
	return token, err
}

// SessionOwner 校验会话：命中且未过期返回 owner（尽力刷新 last_seen）。
func (s *Store) SessionOwner(token string) (string, bool) {
	if token == "" {
		return "", false
	}
	var owner, exp string
	err := s.db.QueryRow(s.q(`SELECT owner,expires_at FROM sessions WHERE token_hash=?`), hashKey(token)).Scan(&owner, &exp)
	if err != nil {
		return "", false
	}
	if expT, err := time.Parse(time.RFC3339Nano, exp); err != nil || time.Now().After(expT) {
		s.db.Exec(s.q(`DELETE FROM sessions WHERE token_hash=?`), hashKey(token))
		return "", false
	}
	s.db.Exec(s.q(`UPDATE sessions SET last_seen_at=? WHERE token_hash=?`), nowISO(), hashKey(token))
	return owner, true
}

// DeleteSession 注销当前会话。
func (s *Store) DeleteSession(token string) {
	s.db.Exec(s.q(`DELETE FROM sessions WHERE token_hash=?`), hashKey(token))
}

// DeleteUserSessions 吊销用户全部会话（改密/重置后）；except 非空时保留该会话。
func (s *Store) DeleteUserSessions(owner, exceptToken string) {
	if exceptToken != "" {
		s.db.Exec(s.q(`DELETE FROM sessions WHERE owner=? AND token_hash<>?`), owner, hashKey(exceptToken))
		return
	}
	s.db.Exec(s.q(`DELETE FROM sessions WHERE owner=?`), owner)
}

// ListUsersAdmin 管理员用户列表（按创建序）。
func (s *Store) ListUsersAdmin() ([]AdminUser, error) {
	rows, err := s.db.Query(s.q(`SELECT owner,COALESCE(name,''),password_hash,must_change_password,is_admin,password_set_at,created_at FROM users ORDER BY created_at`))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []AdminUser{}
	for rows.Next() {
		var u AdminUser
		var hash, psa, ca string
		var mustChange, isAdmin int
		if err := rows.Scan(&u.Owner, &u.Name, &hash, &mustChange, &isAdmin, &psa, &ca); err != nil {
			return nil, err
		}
		u.MustChangePassword = mustChange == 1
		u.IsAdmin = isAdmin == 1
		u.HasPassword = hash != ""
		if psa != "" {
			t, _ := time.Parse(time.RFC3339Nano, psa)
			u.PasswordSetAt = &t
		}
		u.CreatedAt, _ = time.Parse(time.RFC3339Nano, ca)
		out = append(out, u)
	}
	return out, rows.Err()
}