pcm-ac/internal/common/auth.go

package common

import (
	"flag"
	"github.com/zeromicro/go-zero/core/conf"
	"gitlink.org.cn/jcce-pcm/pcm-ac/internal/config"
	"strconv"
)

var (
	as = NewAuthService()
)

type AuthService struct {
	C                  *config.Config
	Token              string //区域用户认证token
	ClusterId          string
	AiCenterUrlPrefix  string
	HpcCenterUrlPrefix string
	EFileUrlPrefix     string
	JobManagerId       int
	GroupId            string
	QueueId            string
}

func NewAuthService() *AuthService {
	cfg := config.Config{}
	configFile := flag.String("c", "etc/hpcac.yaml", "the config file")
	conf.MustLoad(*configFile, &cfg)
	flag.Parse()
	token, clusterId := getTokenAndClusterId(&cfg)
	return &AuthService{
		C:                  &cfg,
		Token:              token,
		ClusterId:          clusterId,
		AiCenterUrlPrefix:  getAiCenterUrl(&cfg, token),
		HpcCenterUrlPrefix: getHpcCenterUrl(&cfg, token),
		EFileUrlPrefix:     getEFileUrlPrefix(&cfg, token),
		JobManagerId:       getJobManagerId(&cfg, token),
		GroupId:            getGroupId(&cfg, token),
		QueueId:            getQueueId(&cfg, token),
	}
}

func (a *AuthService) reGet() {
	a.Token, a.ClusterId = getTokenAndClusterId(a.C)
}

func (a *AuthService) valid() bool {
	tokenStateUrl := a.C.AiConf.BaseUrlAi + a.C.TokenStateUrl
	var ts TokenState
	req := GetRestyRequest(3)
	_, err := req.
		SetHeader("token", a.Token).
		SetResult(&ts).
		Get(tokenStateUrl)

	if err != nil || ts.Code != "0" {
		return false
	}

	return true
}

func GetToken() string {
	if !as.valid() {
		as.reGet()
	}
	return as.Token
}

func GetQueueId() string {
	return as.QueueId
}

func GetGroupId() string {
	return as.GroupId
}

func GetClusterId() string {
	return as.ClusterId
}

func GetJobManagerId() int {
	return as.JobManagerId
}

func AiCenterUrlPrefix() string {
	return as.AiCenterUrlPrefix
}

func HpcCenterUrlPrefix() string {
	return as.HpcCenterUrlPrefix
}

func EFileUrlPrefix() string {
	return as.EFileUrlPrefix
}

func getTokenAndClusterId(cfg *config.Config) (string, string) {
	authUrl := cfg.AiConf.BaseUrlAi + cfg.AuthUrl
	var tr TokenResp
	req := GetRestyRequest(3)
	_, err := req.
		SetHeader("user", cfg.User).
		SetHeader("password", cfg.Password).
		SetHeader("orgId", cfg.OrgId).
		SetResult(&tr).
		Post(authUrl)

	if err != nil || tr.Code != "0" {
		return "", ""
	}

	for _, datum := range tr.Data {
		if datum.ClusterId != "0" {
			return datum.Token, datum.ClusterId
		}
	}

	return "", ""
}

func getAiCenterUrl(cfg *config.Config, token string) string {
	centerUrl := cfg.AiConf.BaseUrlAi + cfg.CenterUrl
	var cr CenterResp
	req := GetRestyRequest(3)
	_, err := req.
		SetHeader("token", token).
		SetResult(&cr).
		Get(centerUrl)

	if err != nil || cr.Code != "0" {
		return ""
	}

	for _, url := range cr.Data.AiUrls {
		if url.Enable == "true" {
			return url.Url
		}
	}
	return ""
}

func getHpcCenterUrl(cfg *config.Config, token string) string {
	centerUrl := cfg.AcBaseUrl + cfg.CenterUrl
	var cr CenterResp
	req := GetRestyRequest(3)
	_, err := req.
		SetHeader("token", token).
		SetResult(&cr).
		Get(centerUrl)

	if err != nil || cr.Code != "0" {
		return ""
	}

	for _, url := range cr.Data.HpcUrls {
		if url.Enable == "true" {
			return url.Url
		}
	}
	return ""
}

func getEFileUrlPrefix(cfg *config.Config, token string) string {
	centerUrl := cfg.AcBaseUrl + cfg.CenterUrl
	var cr CenterResp
	req := GetRestyRequest(3)
	_, err := req.
		SetHeader("token", token).
		SetResult(&cr).
		Get(centerUrl)

	if err != nil || cr.Code != "0" {
		return ""
	}

	for _, url := range cr.Data.EfileUrls {
		if url.Enable == "true" {
			return url.Url
		}
	}
	return ""
}

func getJobManagerId(cfg *config.Config, token string) int {
	url := getHpcCenterUrl(cfg, token) + cfg.CPConf.ClusUrl
	var cr ClusterResp
	req := GetRestyRequest(3)
	_, err := req.
		SetHeader("token", token).
		SetResult(&cr).
		Get(url)

	if err != nil || cr.Code != "0" {
		return 0
	}

	for _, datum := range cr.Data {
		return datum.Id
	}
	return 0
}

func getGroupId(cfg *config.Config, token string) string {
	url := cfg.AcBaseUrl + cfg.UserConf.GetGroupMembers
	var gm GroupMembers
	req := GetRestyRequest(3)
	_, err := req.
		SetHeader("token", token).
		SetResult(&gm).
		Get(url)

	if err != nil || gm.Code != "0" {
		return ""
	}

	for _, datum := range gm.Data {
		// Todo multiple groups filtering
		return datum.GroupId
	}

	return ""
}

func getQueueId(cfg *config.Config, token string) string {
	url := getHpcCenterUrl(cfg, token) + cfg.JobConf.GetQueueNames

	var qn QueueNamesResp
	req := GetRestyRequest(3)
	_, err := req.
		SetHeader("token", token).
		SetPathParam("username", cfg.User).
		SetQueryString("strJobManagerID" + "=" + strconv.Itoa(getJobManagerId(cfg, token))).
		SetResult(&qn).
		Get(url)

	if err != nil || qn.Code != "0" {
		return ""
	}

	for _, datum := range qn.Data {
		// Todo multiple groups filtering
		return datum.Id
	}

	return ""
}