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 runner
 20  
 21  import (
 22  	"context"
 23  	"errors"
 24  	"fmt"
 25  	"strings"
 26  
 27  	"github.com/Tencent/AI-Infra-Guard/common/utils/chromium"
 28  	"github.com/Tencent/AI-Infra-Guard/common/utils/models"
 29  	"github.com/Tencent/AI-Infra-Guard/internal/gologger"
 30  	"github.com/Tencent/AI-Infra-Guard/pkg/vulstruct"
 31  )
 32  
 33  func LoadSensitivePrompt(language string) string {
 34  	prompt := `网站首页源码：
 35  	{{html}}
 36  网站页面渲染后内容：
 37  	{{screenshot_summary}}
 38  上面是目标站点 {{url}} 的网页源码和页面内容描述，请按照如下要求对目标站点 {{url}} 进行分析：
 39  
 40  [角色设定]
 41  你是一个专业的安全工程师，非常擅长渗透测试和红蓝对抗，对安全攻防实战有非常深入的理解。你正在进行企业内网的安全巡检，内网开放了大量的Web网站。你需要评估这些站点在内网的安全风险，筛选出存在敏感功能未鉴权风险的危险站点。
 42  
 43  [概念定义]
 44  网络环境：当前所有操作都是在内网环境中进行的
 45  有鉴权：存在访问控制、登录措施等权限控制措施。例如：重定向到登录页面、存在登录表单等
 46  未鉴权：没有登录措施或者访问控制，可以直接访问到业务数据或者业务功能接口
 47  
 48  [评分依据]
 49  评分说明:
 50  high：该站点存在严重未鉴权风险，内网攻击者可以通过这个站点获取业务敏感数据、或者进一步篡改危害业务的机密性、完整性、可用性。
 51  medium：该站点存在未鉴权风险，但是功能不是很敏感，风险比较低，但还是建议增加鉴权
 52  low：该站点属于内网公共站点，可以开放，不存在未鉴权问题
 53  
 54  low表示低风险，medium表示中风险，high表示高风险。
 55  
 56  [高风险案例]
 57  1. 对于常见的通用组件未鉴权暴露，并且实际可利用时，应该给予较高风险评分，例如：K8S控制台、Hadoop等大数据控制台等等。
 58  2. 特定的业务运维、运营、数据分析系统，这些站点上可能包含大量敏感的业务功能或者业务数据，应该只开放给少量的特定业务员工，如果这些站点暴露了，则风险较高
 59  
 60  [低风险案例]
 61  1. 403 页面、登录页面等，此类页面表示拒绝访问、已经存在访问控制措施，则不存在未鉴权漏洞，风险较低
 62  2. 状态信息页面、无业务信息的中间件默认页面，此类页面如果没有具体业务功能或者数据，则风险较低
 63  3. 页面信息不足以证明其包含敏感功能或者敏感数据，则风险较低
 64  4. 在内网中，某些站点是开放给所有员工使用的公共站点，例如：HR站点、OA门户站点等等。此类站点应给与较低风险。
 65  
 66  [系统类型判定]
 67  分为两种类型：自研业务系统、第三方系统
 68  1. 自研业务系统，公司内部自己研发的网站系统，例如：内部的管理平台、运营系统、运维门户等等
 69  2. 第三方系统：外部开源或者商业的系统，例如：Hadoop、Grafana、cAdvisor、Apache Flink、Alluxio等等
 70  
 71  
 72  [任务目标]
 73  根据目标站点 {{url}} 的首页返回内容，判断该站点是否存在敏感功能或数据风险。
 74  重要：现在是在内网环境下，你现在是以内网普通员工的身份来访问这些站点，请你根据页面上的信息，结合数据敏感度、功能敏感度、漏洞情况等维度，判断该站点的内网未鉴权安全风险，根据评分依据给予评分。
 75  
 76  在内网中，某些站点是开放给所有员工使用的公共站点，例如：HR站点、OA门户站点、代码仓库站点等等。
 77  但是还有某些站点不应该开放给所有员工，例如：某个业务产品的运营系统、某个业务线的运维DevOps平台、某些业务的数据分析平台等等，这些站点上可能包含大量敏感的业务功能或者业务数据，应该只开放给少量的特定业务员工，如果这些站点暴露了，那就属于未鉴权风险。
 78  
 79  [输出格式]
 80  <result>
 81  <title>漏洞总结的精简标题</title>
 82  <details>完成任务目标，包含具体分析信息,自然段落回复</details>
 83  <summary>根据details的精简回复，回复以基于页面的真实浏览器截图判断开头.</summary>
 84  <severity>依据评分依据评分:low，medium，high</severity>
 85  </result>
 86  	`
 87  	if language == "en" {
 88  		prompt += "## Return in English"
 89  	}
 90  	return prompt
 91  }
 92  
 93  func LoadWebPageScreenShotSummary(language string) string {
 94  	prompt := `请作为网页安全分析专家，仔细观察这张网页截图，并按以下结构化格式详细描述：
 95  
 96  1. 页面类型，这是一个什么样的页面
 97  2. 页面布局于功能描述，导航菜单、页面分区、功能模块、数据内容等等
 98  3. 交互元素，页面上有哪些交互元素
 99  4. 敏感信息/敏感功能识别，例如：密钥/密码/token、个人隐私数据信息、内部运维运营数据、敏感的运维运营功能等等
100  6. 是否包含认证与授权信息，登录相关元素（用户名/密码框、验证码、第三方登录）、用户权限提示（管理员、普通用户等）、权限控制相关的UI元素
101  7. 进行整体总结`
102  	if language == "en" {
103  		prompt += "## Return in English"
104  	}
105  	return prompt
106  }
107  
108  func ScreenShot(url string) ([]byte, error) {
109  	instance, err := chromium.NewWebScreenShotWithOptions()
110  	if err != nil {
111  		gologger.WithError(err).Errorf("new screenshot instance error: %s", err)
112  		return nil, err
113  	}
114  	shotData, err := instance.Screen(url)
115  	if err != nil {
116  		gologger.WithError(err).Errorf("get screenshot error: %s", err)
117  		return nil, err
118  	}
119  	if len(shotData) == 0 {
120  		return nil, errors.New("截图失败")
121  	}
122  	return shotData, nil
123  }
124  
125  func Analysis(url string, resp string, language string, model *models.OpenAI) ([]byte, *vulstruct.Info, string, error) {
126  	var shotData []byte
127  	shotData, err := ScreenShot(url)
128  	if err != nil {
129  		gologger.WithError(err).Errorf("screenshot error: %s", err)
130  		return nil, nil, "", err
131  	}
132  	summary, err := model.ChatWithImageByte(context.Background(), LoadWebPageScreenShotSummary(language), shotData)
133  	if err != nil {
134  		gologger.WithError(err).Errorf("chat with image byte error: %s", err)
135  		return shotData, nil, "", err
136  	}
137  	// 敏感信息分析
138  	sensitive := LoadSensitivePrompt(language)
139  	sensitive = strings.ReplaceAll(sensitive, "{{url}}", url)
140  	sensitive = strings.ReplaceAll(sensitive, "{{html}}", resp)
141  	sensitive = strings.ReplaceAll(sensitive, "{{screenshot_summary}}", summary)
142  	ret := ""
143  	for word := range model.ChatStream(context.Background(), []map[string]string{{"role": "user", "content": sensitive}}) {
144  		ret += word
145  	}
146  	gologger.Infof("敏感信息分析: %s", ret)
147  	info := vulstruct.Info{}
148  	info.Details = extractTag(ret, "details")
149  	info.Severity = extractTag(ret, "severity")
150  	info.Summary = extractTag(ret, "title")
151  	x := extractTag(ret, "summary")
152  	return shotData, &info, x, nil
153  }
154  
155  func extractTag(text, tag string) string {
156  	startText := fmt.Sprintf("<%s>", tag)
157  	endText := fmt.Sprintf("</%s>", tag)
158  	startIndex := strings.Index(text, startText)
159  	if startIndex == -1 {
160  		return ""
161  	}
162  	tmp := text[startIndex+len(startText):]
163  	if strings.Index(tmp, endText) == -1 {
164  		return ""
165  	}
166  	endIndex := strings.Index(tmp, endText) + startIndex + len(startText)
167  	if endIndex == -1 || endIndex <= startIndex {
168  		return ""
169  	}
170  	return strings.TrimSpace(text[startIndex+len(startText) : endIndex])
171  }