HY2XS_flamy/apps/service/cron.go

package service

import (
	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
	"hy2xs-admin/dao"
	"hy2xs-admin/model/constant"
	"hy2xs-admin/model/entity"
	"hy2xs-admin/proxy"
	"hy2xs-admin/util"
	"sync"
	"time"
)

var trafficMutex sync.Mutex
var kickMutex sync.Mutex

func CronHandleAccount() {
	go func() {
		if !Hysteria2IsRunning() {
			return
		}

		apiPort, err := GetHysteria2ApiPort()
		if err != nil {
			return
		}

		trafficSecretConfig, err := dao.GetConfig("key = ?", constant.Hysteria2TrafficStatsSecret)
		if err != nil {
			return
		}

		// Сохранение данных трафика
		go saveAccountTraffic(apiPort, *trafficSecretConfig.Value)

		// Принудительное отключение
		go kickAccount(apiPort, *trafficSecretConfig.Value)
	}()
}

func CronResetTraffic() {
	peers, err := dao.ListPeer("1=1")
	if err != nil {
		return
	}
	var ids []int64
	for _, item := range peers {
		ids = append(ids, *item.Id)
	}
	idsList := util.SplitArr(ids, 100)
	for _, item := range idsList {
		if err := dao.UpdatePeer(item, map[string]interface{}{"download_bytes": 0, "upload_bytes": 0}); err != nil {
			continue
		}
	}
}

func saveAccountTraffic(apiPort int64, trafficStatsSecret string) {
	if !trafficMutex.TryLock() {
		return
	}
	defer trafficMutex.Unlock()

	users, err := proxy.NewHysteria2Api(apiPort).ListUsers(true, trafficStatsSecret)
	if err != nil {
		return
	}
	if len(users) == 0 {
		return
	}

	nowMs := time.Now().UnixMilli()
	hourStart := nowMs - (nowMs % int64(time.Hour/time.Millisecond))
	dayStart := nowMs - (nowMs % int64(24*time.Hour/time.Millisecond))

	for key, traffic := range users {
		rxBytes := traffic.Rx
		txBytes := traffic.Tx
		if rxBytes == 0 && txBytes == 0 {
			continue
		}

		peer, peerErr := dao.GetPeer("auth_id = ?", key)
		if peerErr != nil {
			peer, peerErr = dao.GetPeer("name = ?", key)
			if peerErr != nil {
				continue
			}
		}
		if peer.Id == nil {
			continue
		}

		authId := key
		if peer.AuthId != nil && *peer.AuthId != "" {
			authId = *peer.AuthId
		}

		sample := entity.TrafficSample{
			PeerId:    peer.Id,
			AuthId:    &authId,
			RxBytes:   &rxBytes,
			TxBytes:   &txBytes,
			SampledAt: &nowMs,
		}
		if err = dao.SaveTrafficSample(sample); err != nil {
			logrus.Errorf("save traffic_sample failed: %v", err)
			continue
		}

		if err = dao.UpdatePeer([]int64{*peer.Id}, map[string]interface{}{
			"download_bytes": gorm.Expr("download_bytes + ?", rxBytes),
			"upload_bytes":   gorm.Expr("upload_bytes + ?", txBytes),
		}); err != nil {
			logrus.Errorf("update peer traffic failed: %v", err)
			continue
		}

		_ = dao.UpsertTrafficAggregateHourly(*peer.Id, hourStart, rxBytes, txBytes)
		_ = dao.UpsertTrafficAggregateDaily(*peer.Id, dayStart, rxBytes, txBytes)
	}
}
func kickAccount(apiPort int64, trafficStatsSecret string) {
	if !kickMutex.TryLock() {
		return
	}
	defer kickMutex.Unlock()

	users, err := proxy.NewHysteria2Api(apiPort).OnlineUsers(trafficStatsSecret)
	if err != nil {
		return
	}
	if len(users) > 0 {
		i := 0
		usernames := make([]string, len(users))
		for k := range users {
			usernames[i] = k
			i++
		}
		usernameLists := util.SplitArr(usernames, 10)
		var wg sync.WaitGroup
		for _, usernameList := range usernameLists {
			wg.Add(1)
			go func(usernameList []string) {
				defer wg.Done()
				now := time.Now().UnixMilli()
				peers, err := dao.ListPeer("name in ? and (disabled = 1 or (quota_bytes > 0 and quota_bytes < download_bytes + upload_bytes) or ? > expires_at or ? < banned_until)", usernameList, now, now)
				if err != nil {
					return
				}
				kickUsernames := make([]string, len(peers))
				j := 0
				for _, item := range peers {
					kickUsernames[j] = *item.Name
					j++
				}
				if err = proxy.NewHysteria2Api(apiPort).KickUsers(kickUsernames, trafficStatsSecret); err != nil {
					return
				}
			}(usernameList)
		}
		wg.Wait()
	}
}