update

WebScan/lib/client.go

@@ -3,9 +3,13 @@ package lib
 import (
 	"context"
 	"crypto/tls"
+	"embed"
 	"errors"
+	"fmt"
 	"github.com/shadow1ng/fscan/common"
 	"golang.org/x/net/proxy"
+	"gopkg.in/yaml.v2"
+	"io/ioutil"
 	"log"
 	"net"
 	"net/http"
@@ -86,3 +90,166 @@ func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) err
 	}
 	return nil
 }
+
+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"`
+}
+
+type MapSlice = yaml.MapSlice
+
+type StrMap []StrItem
+type ListMap []ListItem
+type RuleMap []RuleItem
+
+type StrItem struct {
+	Key, Value string
+}
+
+type ListItem struct {
+	Key   string
+	Value []string
+}
+
+type RuleItem struct {
+	Key   string
+	Value []Rules
+}
+
+func (r *StrMap) 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
+}
+
+//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
+//}
+
+func (r *RuleMap) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	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
+	}
+
+	for _, one := range tmp1 {
+		key := one.Key.(string)
+		value := tmp[key]
+		*r = append(*r, RuleItem{key, value})
+	}
+	return nil
+}
+
+func (r *ListMap) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var tmp yaml.MapSlice
+	if err := unmarshal(&tmp); err != nil {
+		return err
+	}
+	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)
+		}
+		*r = append(*r, ListItem{key, value})
+	}
+	return nil
+}
+
+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"`
+}
+
+type Detail struct {
+	Author      string   `yaml:"author"`
+	Links       []string `yaml:"links"`
+	Description string   `yaml:"description"`
+	Version     string   `yaml:"version"`
+}
+
+func LoadMultiPoc(Pocs embed.FS, pocname string) []*Poc {
+	var pocs []*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)
+		}
+	}
+	return pocs
+}
+
+func LoadPoc(fileName string, Pocs embed.FS) (*Poc, error) {
+	p := &Poc{}
+	yamlFile, err := Pocs.ReadFile("pocs/" + fileName)
+
+	if err != nil {
+		fmt.Printf("[-] load poc %s error1: %v\n", fileName, err)
+		return nil, err
+	}
+	err = yaml.Unmarshal(yamlFile, p)
+	if err != nil {
+		fmt.Printf("[-] load poc %s error2: %v\n", fileName, err)
+		return nil, err
+	}
+	return p, err
+}
+
+func SelectPoc(Pocs embed.FS, pocname string) []string {
+	entries, err := Pocs.ReadDir("pocs")
+	if err != nil {
+		fmt.Println(err)
+	}
+	var foundFiles []string
+	for _, entry := range entries {
+		if strings.Contains(entry.Name(), pocname) {
+			foundFiles = append(foundFiles, entry.Name())
+		}
+	}
+	return foundFiles
+}
+
+func LoadPocbyPath(fileName string) (*Poc, error) {
+	p := &Poc{}
+	data, err := ioutil.ReadFile(fileName)
+	if err != nil {
+		fmt.Printf("[-] load poc %s error3: %v\n", fileName, err)
+		return nil, err
+	}
+	err = yaml.Unmarshal(data, p)
+	if err != nil {
+		fmt.Printf("[-] load poc %s error4: %v\n", fileName, err)
+		return nil, err
+	}
+	return p, err
+}