mirrors/fscan
1
0
Fork 0
mirror of https://github.com/shadow1ng/fscan.git synced 2025-07-13 21:02:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

perf: 优化Client.go的代码，添加注释，规范输出

Browse Source
This commit is contained in:
ZacharyZcR 2024-12-19 14:26:20 +08:00
parent 6a33a65c94
commit 02eb3d6f7a
1 changed files with 116 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

168
WebScan/lib/Client.go
View File

@ -17,34 +17,43 @@ import (
"time"
)
// 全局HTTP客户端变量
var (
Client *http.Client
ClientNoRedirect *http.Client
dialTimout = 5 * time.Second
keepAlive = 5 * time.Second
Client *http.Client // 标准HTTP客户端
ClientNoRedirect *http.Client // 不自动跟随重定向的HTTP客户端
dialTimout = 5 * time.Second // 连接超时时间
keepAlive = 5 * time.Second // 连接保持时间
)
// Inithttp 初始化HTTP客户端配置
func Inithttp() {
//Common.Proxy = "http://127.0.0.1:8080"
// 设置默认并发数
if Common.PocNum == 0 {
Common.PocNum = 20
}
// 设置默认超时时间
if Common.WebTimeout == 0 {
Common.WebTimeout = 5
}
// 初始化HTTP客户端
err := InitHttpClient(Common.PocNum, Common.Proxy, time.Duration(Common.WebTimeout)*time.Second)
if err != nil {
panic(err)
}
}
// InitHttpClient 创建HTTP客户端
func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) error {
type DialContext = func(ctx context.Context, network, addr string) (net.Conn, error)
// 配置基础连接参数
dialer := &net.Dialer{
Timeout: dialTimout,
KeepAlive: keepAlive,
}
// 配置Transport参数
tr := &http.Transport{
DialContext: dialer.DialContext,
MaxConnsPerHost: 5,
@ -56,6 +65,7 @@ func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) err
DisableKeepAlives: false,
}
// 配置Socks5代理
if Common.Socks5Proxy != "" {
dialSocksProxy, err := Common.Socks5Dialer(dialer)
if err != nil {
@ -64,9 +74,10 @@ func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) err
if contextDialer, ok := dialSocksProxy.(proxy.ContextDialer); ok {
tr.DialContext = contextDialer.DialContext
} else {
return errors.New("Failed type assertion to DialContext")
return errors.New("无法转换为DialContext类型")
}
} else if DownProxy != "" {
// 处理其他代理配置
if DownProxy == "1" {
DownProxy = "http://127.0.0.1:8080"
} else if DownProxy == "2" {
@ -74,9 +85,13 @@ func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) err
} else if !strings.Contains(DownProxy, "://") {
DownProxy = "http://127.0.0.1:" + DownProxy
}
// 验证代理类型
if !strings.HasPrefix(DownProxy, "socks") && !strings.HasPrefix(DownProxy, "http") {
return errors.New("no support this proxy")
return errors.New("不支持的代理类型")
}
// 解析代理URL
u, err := url.Parse(DownProxy)
if err != nil {
return err
@ -84,81 +99,100 @@ func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) err
tr.Proxy = http.ProxyURL(u)
}
// 创建标准HTTP客户端
Client = &http.Client{
Transport: tr,
Timeout: Timeout,
}
// 创建不跟随重定向的HTTP客户端
ClientNoRedirect = &http.Client{
Transport: tr,
Timeout: Timeout,
CheckRedirect: func(req *http.Request, via []*http.Request) error { return http.ErrUseLastResponse },
}
return nil
}
// Poc 定义漏洞检测配置结构
type Poc struct {
Name string `yaml:"name"`
Set StrMap `yaml:"set"`
Sets ListMap `yaml:"sets"`
Rules []Rules `yaml:"rules"`
Groups RuleMap `yaml:"groups"`
Detail Detail `yaml:"detail"`
Name string `yaml:"name"` // POC名称
Set StrMap `yaml:"set"` // 单值配置映射
Sets ListMap `yaml:"sets"` // 列表值配置映射
Rules []Rules `yaml:"rules"` // 检测规则列表
Groups RuleMap `yaml:"groups"` // 规则组映射
Detail Detail `yaml:"detail"` // 漏洞详情
}
// MapSlice 用于解析YAML的通用映射类型
type MapSlice = yaml.MapSlice
type StrMap []StrItem
type ListMap []ListItem
type RuleMap []RuleItem
// 自定义映射类型
type (
StrMap []StrItem // 字符串键值对映射
ListMap []ListItem // 字符串键列表值映射
RuleMap []RuleItem // 字符串键规则列表映射
)
type StrItem struct {
Key, Value string
// 映射项结构定义
type (
// StrItem 字符串键值对
StrItem struct {
Key string // 键名
Value string // 值
}
type ListItem struct {
Key string
Value []string
// ListItem 字符串键列表值对
ListItem struct {
Key string // 键名
Value []string // 值列表
}
type RuleItem struct {
Key string
Value []Rules
// RuleItem 字符串键规则列表对
RuleItem struct {
Key string // 键名
Value []Rules // 规则列表
}
)
// UnmarshalYAML 实现StrMap的YAML解析接口
func (r *StrMap) UnmarshalYAML(unmarshal func(interface{}) error) error {
// 临时使用MapSlice存储解析结果
var tmp yaml.MapSlice
if err := unmarshal(&tmp); err != nil {
return err
}
// 转换为StrMap结构
for _, one := range tmp {
key, value := one.Key.(string), one.Value.(string)
*r = append(*r, StrItem{key, value})
}
return nil
}
//func (r *RuleItem) UnmarshalYAML(unmarshal func(interface{}) error) error {
// var tmp yaml.MapSlice
// if err := unmarshal(&tmp); err != nil {
// return err
// }
// //for _,one := range tmp{
// // key,value := one.Key.(string),one.Value.(string)
// // *r = append(*r,StrItem{key,value})
// //}
// return nil
//}
// UnmarshalYAML 实现RuleMap的YAML解析接口
// 参数：
// - unmarshal: YAML解析函数
//
// 返回：
// - error: 解析错误
func (r *RuleMap) UnmarshalYAML(unmarshal func(interface{}) error) error {
// 使用MapSlice保持键的顺序
var tmp1 yaml.MapSlice
if err := unmarshal(&tmp1); err != nil {
return err
}
// 解析规则内容
var tmp = make(map[string][]Rules)
if err := unmarshal(&tmp); err != nil {
return err
}
// 按顺序转换为RuleMap结构
for _, one := range tmp1 {
key := one.Key.(string)
value := tmp[key]
@ -167,14 +201,24 @@ func (r *RuleMap) UnmarshalYAML(unmarshal func(interface{}) error) error {
return nil
}
// UnmarshalYAML 实现ListMap的YAML解析接口
// 参数：
// - unmarshal: YAML解析函数
//
// 返回：
// - error: 解析错误
func (r *ListMap) UnmarshalYAML(unmarshal func(interface{}) error) error {
// 解析YAML映射
var tmp yaml.MapSlice
if err := unmarshal(&tmp); err != nil {
return err
}
// 转换为ListMap结构
for _, one := range tmp {
key := one.Key.(string)
var value []string
// 将接口类型转换为字符串
for _, val := range one.Value.([]interface{}) {
v := fmt.Sprintf("%v", val)
value = append(value, v)
@ -184,58 +228,68 @@ func (r *ListMap) UnmarshalYAML(unmarshal func(interface{}) error) error {
return nil
}
// Rules 定义POC检测规则结构
type Rules struct {
Method string `yaml:"method"`
Path string `yaml:"path"`
Headers map[string]string `yaml:"headers"`
Body string `yaml:"body"`
Search string `yaml:"search"`
FollowRedirects bool `yaml:"follow_redirects"`
Expression string `yaml:"expression"`
Continue bool `yaml:"continue"`
Method string `yaml:"method"` // HTTP请求方法
Path string `yaml:"path"` // 请求路径
Headers map[string]string `yaml:"headers"` // 请求头
Body string `yaml:"body"` // 请求体
Search string `yaml:"search"` // 搜索模式
FollowRedirects bool `yaml:"follow_redirects"` // 是否跟随重定向
Expression string `yaml:"expression"` // 匹配表达式
Continue bool `yaml:"continue"` // 是否继续执行
}
// Detail 定义POC详情结构
type Detail struct {
Author string `yaml:"author"`
Links []string `yaml:"links"`
Description string `yaml:"description"`
Version string `yaml:"version"`
Author string `yaml:"author"` // POC作者
Links []string `yaml:"links"` // 相关链接
Description string `yaml:"description"` // POC描述
Version string `yaml:"version"` // POC版本
}
// LoadMultiPoc 加载多个POC文件
func LoadMultiPoc(Pocs embed.FS, pocname string) []*Poc {
var pocs []*Poc
// 遍历选中的POC文件
for _, f := range SelectPoc(Pocs, pocname) {
if p, err := LoadPoc(f, Pocs); err == nil {
pocs = append(pocs, p)
} else {
fmt.Println("[-] load poc ", f, " error:", err)
fmt.Printf("[-] POC加载失败 %s: %v\n", f, err)
}
}
return pocs
}
// LoadPoc 从内嵌文件系统加载单个POC
func LoadPoc(fileName string, Pocs embed.FS) (*Poc, error) {
p := &Poc{}
// 读取POC文件内容
yamlFile, err := Pocs.ReadFile("pocs/" + fileName)
if err != nil {
fmt.Printf("[-] load poc %s error1: %v\n", fileName, err)
fmt.Printf("[-] POC文件读取失败 %s: %v\n", fileName, err)
return nil, err
}
// 解析YAML内容
err = yaml.Unmarshal(yamlFile, p)
if err != nil {
fmt.Printf("[-] load poc %s error2: %v\n", fileName, err)
fmt.Printf("[-] POC解析失败 %s: %v\n", fileName, err)
return nil, err
}
return p, err
}
// SelectPoc 根据名称关键字选择POC文件
func SelectPoc(Pocs embed.FS, pocname string) []string {
entries, err := Pocs.ReadDir("pocs")
if err != nil {
fmt.Println(err)
fmt.Printf("[-] 读取POC目录失败: %v\n", err)
}
var foundFiles []string
// 查找匹配关键字的POC文件
for _, entry := range entries {
if strings.Contains(entry.Name(), pocname) {
foundFiles = append(foundFiles, entry.Name())
@ -244,16 +298,20 @@ func SelectPoc(Pocs embed.FS, pocname string) []string {
return foundFiles
}
// LoadPocbyPath 从文件系统路径加载POC
func LoadPocbyPath(fileName string) (*Poc, error) {
p := &Poc{}
// 读取POC文件内容
data, err := os.ReadFile(fileName)
if err != nil {
fmt.Printf("[-] load poc %s error3: %v\n", fileName, err)
fmt.Printf("[-] POC文件读取失败 %s: %v\n", fileName, err)
return nil, err
}
// 解析YAML内容
err = yaml.Unmarshal(data, p)
if err != nil {
fmt.Printf("[-] load poc %s error4: %v\n", fileName, err)
fmt.Printf("[-] POC解析失败 %s: %v\n", fileName, err)
return nil, err
}
return p, err