1  // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
 2  //
 3  // Licensed under the Apache License, Version 2.0 (the "License");
 4  // you may not use this file except in compliance with the License.
 5  // You may obtain a copy of the License at
 6  //
 7  //     http://www.apache.org/licenses/LICENSE-2.0
 8  //
 9  // Unless required by applicable law or agreed to in writing, software
10  // distributed under the License is distributed on an "AS IS" BASIS,
11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12  // See the License for the specific language governing permissions and
13  // limitations under the License.
14  //
15  // Requirement: Any integration or derivative work must explicitly attribute
16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
17  // documentation or user interface, as detailed in the NOTICE file.
18  
19  // Package preload mlflow漏洞go语言写法
20  package preload
21  
22  import (
23  	"errors"
24  	"github.com/Tencent/AI-Infra-Guard/common/utils"
25  	"github.com/Tencent/AI-Infra-Guard/pkg/httpx"
26  	"net/url"
27  	"regexp"
28  	"strings"
29  )
30  
31  // Mlflow struct implements fingerprint detection for MLflow services
32  type Mlflow struct {
33  }
34  
35  // Match checks if the given URI points to an MLflow service
36  // 通过检查页面标题来判断是否为 MLflow 服务
37  func (m Mlflow) Match(httpx *httpx.HTTPX, uri string) bool {
38  	headers := map[string]string{
39  		"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.88 Safari/537.36 Edg/87.0.664.66",
40  	}
41  	resp, err := httpx.Get(uri+"/", headers)
42  	if err != nil {
43  		return false
44  	}
45  	if resp.StatusCode == 200 && strings.Contains(resp.DataStr, "<title>MLflow</title>") {
46  		return true
47  	}
48  	return false
49  }
50  
51  // GetVersion attempts to extract the MLflow version from the service
52  // 通过以下步骤获取 MLflow 版本：
53  // 1. 获取主页面
54  // 2. 提取 JS 文件路径
55  // 3. 请求 JS 文件
56  // 4. 通过正则表达式匹配版本号
57  func (m Mlflow) GetVersion(httpx *httpx.HTTPX, uri string) (string, error) {
58  	flag := `{INTERNAL_ERROR:"INTERNAL_ERROR",INVALID_PARAMETER_VALUE:"INVALID_PARAMETER_VALUE",RESOURCE_DOES_NOT_EXIST:"RESOURCE_DOES_NOT_EXIST",PERMISSION_DENIED:"PERMISSION_DENIED",RESOURCE_CONFLICT:"RESOURCE_CONFLICT"}`
59  	headers := map[string]string{
60  		"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.88 Safari/537.36 Edg/87.0.664.66",
61  	}
62  	resp, err := httpx.Get(uri+"/", headers)
63  	if err != nil {
64  		return "", err
65  	}
66  	if resp.StatusCode != 200 {
67  		return "", errors.New("request code != 200")
68  	}
69  	jsPath := utils.GetMiddleText(`script defer="defer" src="`, `">`, resp.DataStr)
70  	if jsPath == "" {
71  		return "", errors.New("not found js path")
72  	}
73  	newURL, err := url.JoinPath(uri, jsPath)
74  	if err != nil {
75  		return "", err
76  	}
77  	resp2, err := httpx.Get(newURL, headers)
78  	if err != nil {
79  		return "", err
80  	}
81  	index := strings.Index(resp2.DataStr, flag)
82  	if index == -1 {
83  		return "", errors.New("not found flag")
84  	}
85  	parrern := `\d+\.\d+\.\d+`
86  	regex := regexp.MustCompile(parrern)
87  	match := regex.FindString(resp2.DataStr[index+len(flag):])
88  	return match, nil
89  }
90  
91  // Name returns the identifier for this fingerprint detector
92  func (m Mlflow) Name() string {
93  	return "mlflow"
94  }