package identity import ( "context" "crypto/subtle" "encoding/json" "errors" "fmt" "strconv" "strings" "time" platformid "aigateway.local/core/internal/platform/id" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" ) type Repository struct { pool *pgxpool.Pool } func NewRepository(pool *pgxpool.Pool) *Repository { return &Repository{pool: pool} } func (r *Repository) FindAdminByLogin(ctx context.Context, login string) (Account, error) { if r.pool == nil { return Account{}, ErrUnavailable } var account Account account.Kind = KindAdmin err := r.pool.QueryRow(ctx, ` SELECT id::text, username, display_name, role, permissions, password_hash, active, failed_logins, locked_until, totp_enabled, encrypted_totp_secret, totp_kek_version, totp_last_step, totp_backup_codes FROM gateway.admin_accounts WHERE lower(username) = lower($1)`, strings.TrimSpace(login)).Scan( &account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions, &account.PasswordHash, &account.Active, &account.FailedLogins, &account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret, &account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes, ) return account, mapRepositoryError(err) } func (r *Repository) FindAdminByID(ctx context.Context, id string) (Account, error) { if r.pool == nil { return Account{}, ErrUnavailable } var account Account account.Kind = KindAdmin err := r.pool.QueryRow(ctx, ` SELECT id::text, username, display_name, role, permissions, password_hash, active, failed_logins, locked_until, totp_enabled, encrypted_totp_secret, totp_kek_version, totp_last_step, totp_backup_codes FROM gateway.admin_accounts WHERE id = $1`, id).Scan( &account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions, &account.PasswordHash, &account.Active, &account.FailedLogins, &account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret, &account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes, ) return account, mapRepositoryError(err) } func (r *Repository) FindPortalByLogin(ctx context.Context, login string) (Account, error) { if r.pool == nil { return Account{}, ErrUnavailable } var account Account account.Kind = KindPortal err := r.pool.QueryRow(ctx, ` SELECT id::text, account, name, role, permissions, COALESCE(password_hash, ''), auth_source, active, failed_logins, locked_until, totp_enabled, encrypted_totp_secret, totp_kek_version, totp_last_step, totp_backup_codes, department_id::text FROM gateway.portal_users WHERE lower(account) = lower($1)`, strings.TrimSpace(login)).Scan( &account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions, &account.PasswordHash, &account.AuthSource, &account.Active, &account.FailedLogins, &account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret, &account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes, &account.DepartmentID, ) return account, mapRepositoryError(err) } func (r *Repository) FindPortalByID(ctx context.Context, id string) (Account, error) { if r.pool == nil { return Account{}, ErrUnavailable } var account Account account.Kind = KindPortal err := r.pool.QueryRow(ctx, ` SELECT id::text, account, name, role, permissions, COALESCE(password_hash, ''), auth_source, active, failed_logins, locked_until, totp_enabled, encrypted_totp_secret, totp_kek_version, totp_last_step, totp_backup_codes, department_id::text FROM gateway.portal_users WHERE id = $1`, id).Scan( &account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions, &account.PasswordHash, &account.AuthSource, &account.Active, &account.FailedLogins, &account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret, &account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes, &account.DepartmentID, ) return account, mapRepositoryError(err) } func (r *Repository) RecordFailure(ctx context.Context, account Account, maximum int, lockDuration time.Duration) (*time.Time, error) { if r.pool == nil { return nil, ErrUnavailable } table := "gateway.admin_accounts" if account.Kind == KindPortal { table = "gateway.portal_users" } query := fmt.Sprintf(` UPDATE %s SET locked_until = CASE WHEN failed_logins + 1 >= $2 THEN clock_timestamp() + make_interval(secs => $3) ELSE locked_until END, failed_logins = CASE WHEN failed_logins + 1 >= $2 THEN 0 ELSE failed_logins + 1 END, updated_at = clock_timestamp() WHERE id = $1 RETURNING locked_until`, table) var lockedUntil *time.Time err := r.pool.QueryRow(ctx, query, account.ID, maximum, int64(lockDuration.Seconds())).Scan(&lockedUntil) return lockedUntil, mapRepositoryError(err) } func (r *Repository) CompleteLogin(ctx context.Context, account Account, upgradedHash *string) error { if r.pool == nil { return ErrUnavailable } table := "gateway.admin_accounts" if account.Kind == KindPortal { table = "gateway.portal_users" } query := fmt.Sprintf(` UPDATE %s SET failed_logins = 0, locked_until = NULL, last_login = clock_timestamp(), password_hash = COALESCE($2::text, password_hash), updated_at = clock_timestamp() WHERE id = $1`, table) result, err := r.pool.Exec(ctx, query, account.ID, upgradedHash) if err != nil { return fmt.Errorf("%w: %v", ErrUnavailable, err) } if result.RowsAffected() != 1 { return ErrNotFound } return nil } func (r *Repository) SetPassword(ctx context.Context, account Account, passwordHash string) error { query := fmt.Sprintf(`UPDATE %s SET password_hash=$2, failed_logins=0, locked_until=NULL, updated_at=clock_timestamp() WHERE id=$1`, identityTable(account.Kind)) return expectOne(r.pool, ctx, query, account.ID, passwordHash) } func (r *Repository) SetTOTPSecret(ctx context.Context, account Account, encrypted []byte, version int) error { query := fmt.Sprintf(` UPDATE %s SET encrypted_totp_secret = $2, totp_kek_version = $3, totp_enabled = false, totp_last_step = NULL, totp_backup_codes = '[]'::jsonb, totp_confirmed_at = NULL, updated_at = clock_timestamp() WHERE id = $1`, identityTable(account.Kind)) return expectOne(r.pool, ctx, query, account.ID, encrypted, version) } func (r *Repository) EnableTOTP(ctx context.Context, account Account, step int64, records []BackupCodeRecord) error { payload, err := json.Marshal(records) if err != nil { return err } query := fmt.Sprintf(` UPDATE %s SET totp_enabled = true, totp_last_step = $2, totp_backup_codes = $3::jsonb, totp_confirmed_at = clock_timestamp(), updated_at = clock_timestamp() WHERE id = $1 AND encrypted_totp_secret IS NOT NULL AND NOT totp_enabled AND (totp_last_step IS NULL OR totp_last_step < $2)`, identityTable(account.Kind)) return expectOne(r.pool, ctx, query, account.ID, step, payload) } func (r *Repository) ConsumeTOTPStep(ctx context.Context, account Account, step int64) (bool, error) { if r.pool == nil { return false, ErrUnavailable } query := fmt.Sprintf(` UPDATE %s SET totp_last_step = $2, updated_at = clock_timestamp() WHERE id = $1 AND totp_enabled AND (totp_last_step IS NULL OR totp_last_step < $2)`, identityTable(account.Kind)) result, err := r.pool.Exec(ctx, query, account.ID, step) if err != nil { return false, fmt.Errorf("%w: %v", ErrUnavailable, err) } return result.RowsAffected() == 1, nil } func (r *Repository) ConsumeBackupCode(ctx context.Context, account Account, hash string) (bool, error) { if r.pool == nil { return false, ErrUnavailable } tx, err := r.pool.BeginTx(ctx, pgx.TxOptions{}) if err != nil { return false, fmt.Errorf("%w: %v", ErrUnavailable, err) } defer func() { _ = tx.Rollback(ctx) }() query := fmt.Sprintf(`SELECT totp_backup_codes FROM %s WHERE id = $1 AND totp_enabled FOR UPDATE`, identityTable(account.Kind)) var payload []byte if err := tx.QueryRow(ctx, query, account.ID).Scan(&payload); err != nil { return false, mapRepositoryError(err) } var records []BackupCodeRecord if err := json.Unmarshal(payload, &records); err != nil { return false, fmt.Errorf("invalid stored TOTP backup codes: %w", err) } found := -1 for index := range records { if records[index].UsedAt == nil && subtle.ConstantTimeCompare([]byte(records[index].Hash), []byte(hash)) == 1 { found = index } } if found < 0 { return false, nil } now := time.Now().UTC() records[found].UsedAt = &now payload, err = json.Marshal(records) if err != nil { return false, err } update := fmt.Sprintf(`UPDATE %s SET totp_backup_codes = $2::jsonb, updated_at = clock_timestamp() WHERE id = $1`, identityTable(account.Kind)) if _, err := tx.Exec(ctx, update, account.ID, payload); err != nil { return false, fmt.Errorf("%w: %v", ErrUnavailable, err) } if err := tx.Commit(ctx); err != nil { return false, fmt.Errorf("%w: %v", ErrUnavailable, err) } return true, nil } func (r *Repository) DisableTOTP(ctx context.Context, account Account) error { query := fmt.Sprintf(` UPDATE %s SET encrypted_totp_secret = NULL, totp_kek_version = NULL, totp_enabled = false, totp_last_step = NULL, totp_backup_codes = '[]'::jsonb, totp_confirmed_at = NULL, updated_at = clock_timestamp() WHERE id = $1 AND totp_enabled`, identityTable(account.Kind)) return expectOne(r.pool, ctx, query, account.ID) } func (r *Repository) ReplaceBackupCodes(ctx context.Context, account Account, records []BackupCodeRecord) error { payload, err := json.Marshal(records) if err != nil { return err } query := fmt.Sprintf(`UPDATE %s SET totp_backup_codes = $2::jsonb, updated_at = clock_timestamp() WHERE id = $1 AND totp_enabled`, identityTable(account.Kind)) return expectOne(r.pool, ctx, query, account.ID, payload) } func identityTable(kind Kind) string { if kind == KindPortal { return "gateway.portal_users" } return "gateway.admin_accounts" } func expectOne(pool *pgxpool.Pool, ctx context.Context, query string, arguments ...any) error { if pool == nil { return ErrUnavailable } result, err := pool.Exec(ctx, query, arguments...) if err != nil { return fmt.Errorf("%w: %v", ErrUnavailable, err) } if result.RowsAffected() != 1 { return ErrNotFound } return nil } // LoginLog 是一条登录尝试记录。 type LoginLog struct { ID string `json:"id"` Kind string `json:"kind"` Login string `json:"login"` Success bool `json:"success"` IP *string `json:"ip,omitempty"` UserAgent string `json:"user_agent"` Reason string `json:"reason"` CreatedAt time.Time `json:"created_at"` } // RecordLoginLog 记录一次登录尝试(成功或失败)。 func (r *Repository) RecordLoginLog(ctx context.Context, kind Kind, login string, success bool, ip, userAgent, reason string) error { if r.pool == nil { return ErrUnavailable } id, err := platformid.NewUUID() if err != nil { return err } ipValue := strings.TrimSpace(ip) _, err = r.pool.Exec(ctx, `INSERT INTO gateway.login_logs(id,kind,login,success,ip,user_agent,reason) VALUES($1,$2,$3,$4,nullif($5,'')::inet,$6,$7)`, id, string(kind), login, success, ipValue, truncateText(userAgent, 256), truncateText(reason, 64)) if err != nil { return fmt.Errorf("%w: %v", ErrUnavailable, err) } return nil } // ListLoginLogs 查询登录记录(按时间倒序)。 func (r *Repository) ListLoginLogs(ctx context.Context, kind Kind, login string, limit int) ([]LoginLog, error) { if r.pool == nil { return nil, ErrUnavailable } if limit < 1 || limit > 500 { limit = 50 } query := `SELECT id::text,kind,login,success,ip::text,user_agent,reason,created_at FROM gateway.login_logs WHERE kind=$1` args := []any{string(kind)} if login != "" { args = append(args, login) query += ` AND login=$` + strconv.Itoa(len(args)) } query += ` ORDER BY created_at DESC LIMIT $` + strconv.Itoa(len(args)+1) args = append(args, limit) rows, err := r.pool.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("%w: %v", ErrUnavailable, err) } defer rows.Close() items := []LoginLog{} for rows.Next() { var item LoginLog var ip *string if err := rows.Scan(&item.ID, &item.Kind, &item.Login, &item.Success, &ip, &item.UserAgent, &item.Reason, &item.CreatedAt); err != nil { return nil, err } if ip != nil && *ip != "" { item.IP = ip } items = append(items, item) } return items, rows.Err() } func truncateText(value string, limit int) string { value = strings.TrimSpace(value) if len(value) <= limit { return value } return value[:limit] } // CountIdentities 统计管理员与门户账号总数(License 账号数管控用)。 func (r *Repository) CountIdentities(ctx context.Context, total *int) error { if r.pool == nil { return ErrUnavailable } err := r.pool.QueryRow(ctx, `SELECT (SELECT count(*) FROM gateway.admin_accounts) + (SELECT count(*) FROM gateway.portal_users)`).Scan(total) if err != nil { return fmt.Errorf("%w: %v", ErrUnavailable, err) } return nil } func (r *Repository) CreateAdmin(ctx context.Context, login, displayName, role, passwordHash string) (string, error) { if r.pool == nil { return "", ErrUnavailable } id, err := platformid.NewUUID() if err != nil { return "", err } result, err := r.pool.Exec(ctx, ` INSERT INTO gateway.admin_accounts (id, username, display_name, role, password_hash) VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING`, id, strings.ToLower(strings.TrimSpace(login)), displayName, role, passwordHash) if err != nil { return "", err } if result.RowsAffected() != 1 { return "", errors.New("administrator already exists") } return id, nil } func (r *Repository) CreatePortalUser(ctx context.Context, login, displayName, passwordHash string) (string, error) { if r.pool == nil { return "", ErrUnavailable } id, err := platformid.NewUUID() if err != nil { return "", err } result, err := r.pool.Exec(ctx, ` INSERT INTO gateway.portal_users (id, account, name, password_hash, auth_source) VALUES ($1, $2, $3, $4, 'local') ON CONFLICT DO NOTHING`, id, strings.ToLower(strings.TrimSpace(login)), displayName, passwordHash) if err != nil { return "", err } if result.RowsAffected() != 1 { return "", errors.New("portal user already exists") } return id, nil } func mapRepositoryError(err error) error { if errors.Is(err, pgx.ErrNoRows) { return ErrNotFound } if err != nil { return fmt.Errorf("%w: %v", ErrUnavailable, err) } return nil }