HY2XS_flamy/apps/service/peer.go

package service

import (
	"errors"
	"fmt"
	"strings"

	"github.com/skip2/go-qrcode"
	"hy2xs-admin/dao"
	"hy2xs-admin/model/bo"
	"hy2xs-admin/model/constant"
	"hy2xs-admin/model/dto"
	"hy2xs-admin/model/entity"
	"hy2xs-admin/model/vo"
	"hy2xs-admin/util"
)

func PagePeer(peerPageDto dto.PeerPageDto) ([]vo.PeerVo, int64, error) {
	peers, total, err := dao.PagePeer(peerPageDto)
	if err != nil {
		return nil, 0, err
	}
	onlineUsers, _ := Hysteria2Online()
	result := make([]vo.PeerVo, 0, len(peers))
	for _, p := range peers {
		item := vo.PeerVo{
			BaseVo:           vo.BaseVo{Id: *p.Id, CreateTime: *p.CreateTime},
			Name:             strVal(p.Name),
			Remark:           strVal(p.Remark),
			AuthId:           strVal(p.AuthId),
			QuotaBytes:       int64Val(p.QuotaBytes),
			DownloadBytes:    int64Val(p.DownloadBytes),
			UploadBytes:      int64Val(p.UploadBytes),
			ExpiresAt:        int64Val(p.ExpiresAt),
			MaxDevices:       int64Val(p.MaxDevices),
			Disabled:         int64Val(p.Disabled),
			BannedUntil:      int64Val(p.BannedUntil),
			LastConnectionAt: int64Val(p.LastConnectionAt),
		}
		authID := strVal(p.AuthId)
		if v, ok := onlineUsers[authID]; ok {
			item.Online = true
			item.OnlineDevices = v
		}
		result = append(result, item)
	}
	return result, total, nil
}

func CreatePeer(peerDto dto.PeerSaveDto) (vo.PeerVo, error) {
	if peerDto.Name == nil || *peerDto.Name == "" {
		return vo.PeerVo{}, errors.New(constant.InvalidError)
	}
	if ExistPeerName(*peerDto.Name, 0) {
		return vo.PeerVo{}, errors.New(fmt.Sprintf("name %s already exists", *peerDto.Name))
	}
	secret := ""
	if peerDto.Secret != nil && *peerDto.Secret != "" {
		secret = *peerDto.Secret
	} else {
		generated, err := util.RandomString(24)
		if err != nil {
			return vo.PeerVo{}, err
		}
		secret = fmt.Sprintf("%s.%s", *peerDto.Name, generated)
	}
	authId, err := util.RandomString(18)
	if err != nil {
		return vo.PeerVo{}, err
	}
	secretDigest, err := PeerSecretDigest(secret)
	if err != nil {
		return vo.PeerVo{}, err
	}
	secretEncrypted, err := EncryptPeerSecret(secret)
	if err != nil {
		return vo.PeerVo{}, err
	}
	peer := entity.Peer{
		Name:            peerDto.Name,
		Remark:          peerDto.Remark,
		AuthId:          &authId,
		SecretDigest:    &secretDigest,
		SecretEncrypted: &secretEncrypted,
		QuotaBytes:      peerDto.QuotaBytes,
		ExpiresAt:       peerDto.ExpiresAt,
		MaxDevices:      peerDto.MaxDevices,
		Disabled:        peerDto.Disabled,
	}
	id, saveErr := dao.SavePeer(peer)
	if saveErr != nil {
		return vo.PeerVo{}, saveErr
	}
	return GetPeerVo(id)
}

func UpdatePeer(id int64, peerDto dto.PeerUpdateDto) error {
	updates := map[string]interface{}{}
	if peerDto.Name != nil && *peerDto.Name != "" {
		updates["name"] = *peerDto.Name
	}
	if peerDto.Secret != nil && *peerDto.Secret != "" {
		digest, err := PeerSecretDigest(*peerDto.Secret)
		if err != nil {
			return err
		}
		enc, err := EncryptPeerSecret(*peerDto.Secret)
		if err != nil {
			return err
		}
		updates["secret_digest"] = digest
		updates["secret_ciphertext"] = enc
	}
	if peerDto.QuotaBytes != nil {
		updates["quota_bytes"] = *peerDto.QuotaBytes
	}
	if peerDto.ExpiresAt != nil {
		updates["expires_at"] = *peerDto.ExpiresAt
	}
	if peerDto.MaxDevices != nil {
		updates["max_devices"] = *peerDto.MaxDevices
	}
	if peerDto.Disabled != nil {
		updates["disabled"] = *peerDto.Disabled
	}
	if peerDto.Remark != nil {
		updates["remark"] = *peerDto.Remark
	}
	return dao.UpdatePeer([]int64{id}, updates)
}

func DeletePeer(id int64) error { return dao.DeletePeer([]int64{id}) }

func GetPeerVo(id int64) (vo.PeerVo, error) {
	p, err := dao.GetPeer("id = ?", id)
	if err != nil {
		return vo.PeerVo{}, err
	}
	return vo.PeerVo{
		BaseVo:           vo.BaseVo{Id: *p.Id, CreateTime: *p.CreateTime},
		Name:             strVal(p.Name),
		Remark:           strVal(p.Remark),
		AuthId:           strVal(p.AuthId),
		QuotaBytes:       int64Val(p.QuotaBytes),
		DownloadBytes:    int64Val(p.DownloadBytes),
		UploadBytes:      int64Val(p.UploadBytes),
		ExpiresAt:        int64Val(p.ExpiresAt),
		MaxDevices:       int64Val(p.MaxDevices),
		Disabled:         int64Val(p.Disabled),
		BannedUntil:      int64Val(p.BannedUntil),
		LastConnectionAt: int64Val(p.LastConnectionAt),
	}, nil
}

func ResetPeerTraffic(id int64) error {
	return dao.UpdatePeer([]int64{id}, map[string]interface{}{"download_bytes": 0, "upload_bytes": 0})
}
func ReleaseKickPeer(id int64) error {
	return dao.UpdatePeer([]int64{id}, map[string]interface{}{"banned_until": 0})
}

func KickPeer(id int64, bannedUntil int64) error {
	if err := dao.UpdatePeer([]int64{id}, map[string]interface{}{"banned_until": bannedUntil}); err != nil {
		return err
	}
	return Hysteria2Kick([]int64{id}, bannedUntil)
}

func BuildPeerClientConfig(id int64) (vo.PeerClientConfigVo, error) {
	url, err := Hysteria2Url(id)
	if err != nil {
		return vo.PeerClientConfigVo{}, err
	}
	qrCode, err := qrcode.Encode(url, qrcode.Medium, 300)
	if err != nil {
		return vo.PeerClientConfigVo{}, err
	}
	return vo.PeerClientConfigVo{Url: url, QrCode: qrCode}, nil
}

func ListExportPeer(includeSecrets bool) ([]bo.PeerExport, error) {
	peers, err := dao.ListPeer("1=1")
	if err != nil {
		return nil, errors.New(constant.SysError)
	}
	out := make([]bo.PeerExport, 0, len(peers))
	for _, item := range peers {
		ex := bo.PeerExport{
			Id:               int64Val(item.Id),
			AuthId:           strVal(item.AuthId),
			Name:             strVal(item.Name),
			Remark:           strVal(item.Remark),
			QuotaBytes:       int64Val(item.QuotaBytes),
			DownloadBytes:    int64Val(item.DownloadBytes),
			UploadBytes:      int64Val(item.UploadBytes),
			ExpiresAt:        int64Val(item.ExpiresAt),
			MaxDevices:       int64Val(item.MaxDevices),
			Disabled:         int64Val(item.Disabled),
			BannedUntil:      int64Val(item.BannedUntil),
			LastConnectionAt: int64Val(item.LastConnectionAt),
		}
		if includeSecrets && item.SecretEncrypted != nil {
			if dec, derr := DecryptPeerSecret(*item.SecretEncrypted); derr == nil {
				ex.Secret = dec
			}
		}
		out = append(out, ex)
	}
	return out, nil
}

func UpsertPeerExport(items []bo.PeerExport) error {
	for _, item := range items {
		name := strings.TrimSpace(item.Name)
		if name == "" {
			return errors.New(constant.InvalidError)
		}

		var existing entity.Peer
		var err error
		authID := strings.TrimSpace(item.AuthId)
		if authID != "" {
			existing, err = dao.GetPeer("auth_id = ?", authID)
		}
		if (err != nil || existing.Id == nil) && name != "" {
			existing, err = dao.GetPeer("name = ?", name)
		}

		quota := item.QuotaBytes
		expires := item.ExpiresAt
		maxDevices := item.MaxDevices
		if maxDevices <= 0 {
			maxDevices = 3
		}
		disabled := item.Disabled
		remark := item.Remark

		if err == nil && existing.Id != nil {
			updates := map[string]interface{}{
				"name":               name,
				"remark":             remark,
				"quota_bytes":        quota,
				"download_bytes":     item.DownloadBytes,
				"upload_bytes":       item.UploadBytes,
				"expires_at":         expires,
				"max_devices":        maxDevices,
				"disabled":           disabled,
				"banned_until":       item.BannedUntil,
				"last_connection_at": item.LastConnectionAt,
			}
			if authID != "" {
				updates["auth_id"] = authID
			}
			if strings.TrimSpace(item.Secret) != "" {
				digest, derr := PeerSecretDigest(item.Secret)
				if derr != nil {
					return derr
				}
				enc, eerr := EncryptPeerSecret(item.Secret)
				if eerr != nil {
					return eerr
				}
				updates["secret_digest"] = digest
				updates["secret_ciphertext"] = enc
			}
			if uerr := dao.UpdatePeer([]int64{*existing.Id}, updates); uerr != nil {
				return uerr
			}
			continue
		}

		createSecret := strings.TrimSpace(item.Secret)
		if createSecret == "" {
			generated, gerr := util.RandomString(24)
			if gerr != nil {
				return gerr
			}
			createSecret = fmt.Sprintf("%s.%s", name, generated)
		}
		secretDigest, derr := PeerSecretDigest(createSecret)
		if derr != nil {
			return derr
		}
		secretEncrypted, eerr := EncryptPeerSecret(createSecret)
		if eerr != nil {
			return eerr
		}
		if authID == "" {
			authID, err = util.RandomString(18)
			if err != nil {
				return err
			}
		}

		peer := entity.Peer{
			Name:             &name,
			Remark:           &remark,
			AuthId:           &authID,
			SecretDigest:     &secretDigest,
			SecretEncrypted:  &secretEncrypted,
			QuotaBytes:       &quota,
			DownloadBytes:    &item.DownloadBytes,
			UploadBytes:      &item.UploadBytes,
			ExpiresAt:        &expires,
			MaxDevices:       &maxDevices,
			Disabled:         &disabled,
			BannedUntil:      &item.BannedUntil,
			LastConnectionAt: &item.LastConnectionAt,
		}
		if _, serr := dao.SavePeer(peer); serr != nil {
			return serr
		}
	}
	return nil
}

func ExistPeerName(name string, id int64) bool {
	var err error
	if id != 0 {
		_, err = dao.GetPeer("name = ? and id != ?", name, id)
	} else {
		_, err = dao.GetPeer("name = ?", name)
	}
	return err == nil
}

func UpdatePeerLastConnectionAt(id int64, conAt int64) error {
	return dao.UpdatePeer([]int64{id}, map[string]interface{}{"last_connection_at": conAt})
}

func strVal(v *string) string {
	if v == nil {
		return ""
	}
	return *v
}
func int64Val(v *int64) int64 {
	if v == nil {
		return 0
	}
	return *v
}