BYWCloudreve/inventory/loginaudit.go

package inventory

import (
	"context"
	"time"

	"github.com/cloudreve/Cloudreve/v4/ent"
	"github.com/cloudreve/Cloudreve/v4/ent/loginaudit"
	"github.com/cloudreve/Cloudreve/v4/ent/schema"
)

type (
	LoginAuditStatus string

	LoginAuditListArgs struct {
		Page        int
		PageSize    int
		UserID      *int
		IP          string
		Status      *LoginAuditStatus
		Username    string
		StartTime   *time.Time
		EndTime     *time.Time
	}

	LoginAuditDailyItem struct {
		Date     string `json:"date"`
		Success  int    `json:"success"`
		Fail     int    `json:"fail"`
		Total    int    `json:"total"`
	}

	LoginAuditListResult struct {
		Total int              `json:"total"`
		Items []*ent.LoginAudit `json:"items"`
	}

	LoginAuditClient interface {
		TxOperator
		// RecordLogin 记录一次登录尝试（成功或失败）
		RecordLogin(ctx context.Context, userID *int, ip, userAgent string, success bool, failReason, usernameTried *string) (*ent.LoginAudit, error)
		// MarkLogout 标记登出（找到该用户+IP最近一条成功且未登出的记录）
		MarkLogout(ctx context.Context, uid int, ip string, logoutAt time.Time) error
		// List 分页查询登录审计
		List(ctx context.Context, args *LoginAuditListArgs) (*LoginAuditListResult, error)
		// DailyStats 每日登录统计（按 created_at 日期分组）
		DailyStats(ctx context.Context, start, end time.Time) ([]LoginAuditDailyItem, error)
	}
)

const (
	LoginAuditStatusSuccess LoginAuditStatus = "success"
	LoginAuditStatusFail    LoginAuditStatus = "fail"
)

func NewLoginAuditClient(client *ent.Client) LoginAuditClient {
	return &loginAuditClient{client: client}
}

type loginAuditClient struct {
	client *ent.Client
}

func (c *loginAuditClient) SetClient(newClient *ent.Client) TxOperator {
	return &loginAuditClient{client: newClient}
}

func (c *loginAuditClient) GetClient() *ent.Client {
	return c.client
}

func (c *loginAuditClient) RecordLogin(ctx context.Context, userID *int, ip, userAgent string, success bool, failReason, usernameTried *string) (*ent.LoginAudit, error) {
	status := loginaudit.StatusFail
	if success {
		status = loginaudit.StatusSuccess
	}
	b := c.client.LoginAudit.Create().
		SetIP(ip).
		SetUserAgent(userAgent).
		SetStatus(status)

	if userID != nil {
		b = b.SetUserID(*userID)
	}
	if !success && failReason != nil {
		b = b.SetFailReason(*failReason)
	}
	if usernameTried != nil {
		b = b.SetUsernameTried(*usernameTried)
	}

	return b.Save(ctx)
}

func (c *loginAuditClient) MarkLogout(ctx context.Context, uid int, ip string, logoutAt time.Time) error {
	// 查找该用户+IP 最近一条登录成功且未登出的记录
	latest, err := c.client.LoginAudit.Query().
		Where(
			loginaudit.UserID(uid),
			loginaudit.IP(ip),
			loginaudit.StatusEQ(loginaudit.StatusSuccess),
			loginaudit.LogoutAtIsNil(),
		).
		Order(ent.Desc(loginaudit.FieldCreatedAt)).
		First(ctx)
	if err != nil {
		// 找不到记录也不报错，可能会话过期或其他原因
		return nil
	}

	seconds := int(logoutAt.Sub(latest.CreatedAt).Seconds())
	if seconds < 0 {
		seconds = 0
	}

	_, err = c.client.LoginAudit.UpdateOne(latest).
		SetLogoutAt(logoutAt).
		SetSessionSeconds(seconds).
		Save(ctx)
	return err
}

func (c *loginAuditClient) List(ctx context.Context, args *LoginAuditListArgs) (*LoginAuditListResult, error) {
	page := args.Page
	if page < 1 {
		page = 1
	}
	pageSize := args.PageSize
	if pageSize < 1 {
		pageSize = 20
	}

	q := c.client.LoginAudit.Query()

	if args.UserID != nil {
		q = q.Where(loginaudit.UserID(*args.UserID))
	}
	if args.IP != "" {
		q = q.Where(loginaudit.IPContains(args.IP))
	}
	if args.Status != nil {
		switch *args.Status {
		case LoginAuditStatusSuccess:
			q = q.Where(loginaudit.StatusEQ(loginaudit.StatusSuccess))
		case LoginAuditStatusFail:
			q = q.Where(loginaudit.StatusEQ(loginaudit.StatusFail))
		}
	}
	if args.Username != "" {
		q = q.Where(loginaudit.UsernameTriedContains(args.Username))
	}
	if args.StartTime != nil {
		q = q.Where(loginaudit.CreatedAtGTE(*args.StartTime))
	}
	if args.EndTime != nil {
		q = q.Where(loginaudit.CreatedAtLT(*args.EndTime))
	}

	// 软删除：LoginAudit 使用 CommonMixin，默认需要走 SkipSoftDelete 排除软删除过滤？
	// 这里审计记录不应该被软删除，但仍遵循默认查询规则，不主动跳过
	ctx = schema.SkipSoftDelete(ctx)

	total, err := q.Clone().Count(ctx)
	if err != nil {
		return nil, err
	}

	items, err := q.
		Order(ent.Desc(loginaudit.FieldCreatedAt)).
		WithUser().
		Limit(pageSize).
		Offset((page - 1) * pageSize).
		All(ctx)
	if err != nil {
		return nil, err
	}

	return &LoginAuditListResult{Total: total, Items: items}, nil
}

func (c *loginAuditClient) DailyStats(ctx context.Context, start, end time.Time) ([]LoginAuditDailyItem, error) {
	ctx = schema.SkipSoftDelete(ctx)

	// 先取得所有符合时间范围的记录，在内存中按日期聚合
	items, err := c.client.LoginAudit.Query().
		Where(
			loginaudit.CreatedAtGTE(start),
			loginaudit.CreatedAtLT(end.AddDate(0, 0, 1)),
		).
		Select(loginaudit.FieldCreatedAt, loginaudit.FieldStatus).
		All(ctx)
	if err != nil {
		return nil, err
	}

	agg := make(map[string]*LoginAuditDailyItem)
	for _, it := range items {
		day := it.CreatedAt.Format("2006-01-02")
		v, ok := agg[day]
		if !ok {
			v = &LoginAuditDailyItem{Date: day}
			agg[day] = v
		}
		switch it.Status {
		case loginaudit.StatusSuccess:
			v.Success++
		default:
			v.Fail++
		}
		v.Total++
	}

	// 补齐 start 到 end 之间每一天，按日期排序
	out := make([]LoginAuditDailyItem, 0)
	for d := start; !d.After(end); d = d.AddDate(0, 0, 1) {
		day := d.Format("2006-01-02")
		if v, ok := agg[day]; ok {
			out = append(out, *v)
		} else {
			out = append(out, LoginAuditDailyItem{Date: day})
		}
	}

	return out, nil
}