[update] add Aseet protocols support

pkg/model/assets.go

@@ -38,6 +38,7 @@ type assetSorter struct {
 func (s *assetSorter) Len() int {
 	return len(s.data)
 }
+
 func (s *assetSorter) Swap(i, j int) {
 	s.data[i], s.data[j] = s.data[j], s.data[i]
 }
@@ -75,20 +76,60 @@ func assetSortByHostName(asset1, asset2 *Asset) bool {
 type NodeList []Node
 
 type Asset struct {
-	ID              string       `json:"id"`
-	Hostname        string       `json:"hostname"`
-	IP              string       `json:"ip"`
-	Port            int          `json:"port"`
-	SystemUsers     []SystemUser `json:"system_users_granted"`
-	IsActive        bool         `json:"is_active"`
-	SystemUsersJoin string       `json:"system_users_join"`
-	Os              string       `json:"os"`
-	Domain          string       `json:"domain"`
-	Platform        string       `json:"platform"`
-	Comment         string       `json:"comment"`
-	Protocol        string       `json:"protocol"`
-	OrgID           string       `json:"org_id"`
-	OrgName         string       `json:"org_name"`
+	ID              string         `json:"id"`
+	Hostname        string         `json:"hostname"`
+	IP              string         `json:"ip"`
+	Port            int            `json:"port"`
+	SystemUsers     []SystemUser   `json:"system_users_granted"`
+	IsActive        bool           `json:"is_active"`
+	SystemUsersJoin string         `json:"system_users_join"`
+	Os              string         `json:"os"`
+	Domain          string         `json:"domain"`
+	Platform        string         `json:"platform"`
+	Comment         string         `json:"comment"`
+	Protocol        string         `json:"protocol"`
+	Protocols       []protocolItem `json:"protocols,omitempty"`
+	OrgID           string         `json:"org_id"`
+	OrgName         string         `json:"org_name"`
+}
+
+func (a *Asset) ProtocolPort(protocol string) int {
+	// 向下兼容
+	if a.Protocols == nil {
+		return a.Port
+	}
+	for _, item := range a.Protocols {
+		if strings.ToLower(item.Name) == strings.ToLower(protocol) {
+			return item.Port
+		}
+	}
+	switch strings.ToLower(protocol) {
+	case "telnet":
+		return 23
+	case "vnc":
+		return 5901
+	case "rdp":
+		return 3389
+	default:
+		return 22
+	}
+}
+
+func (a *Asset) IsSupportProtocol(protocol string) bool {
+	if a.Protocols == nil {
+		return a.Protocol == protocol
+	}
+	for _, item := range a.Protocols {
+		if strings.ToLower(item.Name) == strings.ToLower(protocol) {
+			return true
+		}
+	}
+	return false
+}
+
+type protocolItem struct {
+	Name string `json:"name"`
+	Port int    `json:"port"`
 }
 
 type Gateway struct {

pkg/proxy/proxy.go

@@ -66,12 +66,12 @@ func (p *ProxyServer) getSystemUserUsernameIfNeed() {
 
 // checkProtocolMatch 检查协议是否匹配
 func (p *ProxyServer) checkProtocolMatch() bool {
-	return p.SystemUser.Protocol == p.Asset.Protocol
+	return p.Asset.IsSupportProtocol(p.SystemUser.Protocol)
 }
 
 // checkProtocolIsGraph 检查协议是否是图形化的
 func (p *ProxyServer) checkProtocolIsGraph() bool {
-	switch p.Asset.Protocol {
+	switch p.SystemUser.Protocol {
 	case "ssh", "telnet":
 		return false
 	default:

pkg/srvconn/connmanager.go

@@ -141,7 +141,7 @@ func MakeConfig(asset *model.Asset, systemUser *model.SystemUser, timeout time.D
 	}
 	conf = &SSHClientConfig{
 		Host:       asset.IP,
-		Port:       strconv.Itoa(asset.Port),
+		Port:       strconv.Itoa(asset.ProtocolPort("ssh")),
 		User:       systemUser.Username,
 		Password:   systemUser.Password,
 		PrivateKey: systemUser.PrivateKey,

pkg/srvconn/telnetconn.go

@@ -140,7 +140,7 @@ func (tc *ServerTelnetConnection) login(data []byte) AuthStatus {
 
 func (tc *ServerTelnetConnection) Connect(h, w int, term string) (err error) {
 	var ip = tc.Asset.IP
-	var port = strconv.Itoa(tc.Asset.Port)
+	var port = strconv.Itoa(tc.Asset.ProtocolPort("telnet"))
 	var asset = tc.Asset
 	var proxyConn *gossh.Client