添加templates以及templates generator

This commit is contained in:
M09Ic 2022-10-28 00:01:05 +08:00
parent 6cc0a71dae
commit e88d780e81
4 changed files with 142 additions and 2 deletions

View File

@ -1,4 +1,4 @@
package main
package cmd
import (
"fmt"
@ -7,7 +7,7 @@ import (
"github.com/jessevdk/go-flags"
)
func main() {
func Spray() {
var option internal.Option
parser := flags.NewParser(&option, flags.Default)
_, err := parser.Parse()

13
internal/templates.go Normal file
View File

@ -0,0 +1,13 @@
package internal
import (
"github.com/chainreactors/files"
"github.com/chainreactors/parsers"
)
func LoadConfig(typ string) []byte {
if typ == "http" {
return files.UnFlate(parsers.Base64Decode(""))
}
return []byte{}
}

116
internal/templates_gen.go Normal file
View File

@ -0,0 +1,116 @@
//go:build ignore
// +build ignore
package main
import (
"encoding/json"
"fmt"
"github.com/chainreactors/files"
"github.com/chainreactors/parsers"
"io"
"os"
"path/filepath"
"sigs.k8s.io/yaml"
)
func Encode(input []byte) string {
return parsers.Base64Encode(files.Flate(input))
}
func loadYamlFile2JsonString(filename string) string {
var err error
file, err := os.Open("templates/" + filename)
if err != nil {
panic(err.Error())
}
bs, _ := io.ReadAll(file)
jsonstr, err := yaml.YAMLToJSON(bs)
if err != nil {
panic(filename + err.Error())
}
return Encode(jsonstr)
}
func visit(files *[]string) filepath.WalkFunc {
return func(path string, info os.FileInfo, err error) error {
if err != nil {
panic(err)
}
if !info.IsDir() {
*files = append(*files, path)
}
return nil
}
}
func recuLoadYamlFiles2JsonString(dir string, single bool) string {
var files []string
err := filepath.Walk("templates/"+dir, visit(&files))
if err != nil {
panic(err)
}
var pocs []interface{}
for _, file := range files {
var tmp interface{}
bs, err := os.ReadFile(file)
if err != nil {
panic(err)
}
err = yaml.Unmarshal(bs, &tmp)
if err != nil {
print(file)
panic(err)
}
if tmp == nil {
continue
}
if single {
pocs = append(pocs, tmp)
} else {
pocs = append(pocs, tmp.([]interface{})...)
}
}
jsonstr, err := json.Marshal(pocs)
if err != nil {
panic(err)
}
return Encode(jsonstr)
}
func main() {
template := `package internal
import (
"github.com/chainreactors/files"
"github.com/chainreactors/parsers"
)
func LoadConfig(typ string) []byte {
if typ == "http" {
return files.UnFlate(parsers.Base64Decode("%s"))
}
return []byte{}
}
`
template = fmt.Sprintf(template,
recuLoadYamlFiles2JsonString("fingers/http", false),
)
f, err := os.OpenFile("internal/templates.go", os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
if err != nil {
panic(err)
}
f.WriteString(template)
f.Sync()
f.Close()
println("generate templates.go successfully")
}

11
spray.go Normal file
View File

@ -0,0 +1,11 @@
//go:generate go run internal/templates_gen.go
package main
import "github.com/chainreactors/spray/cmd"
func main() {
//f, _ := os.Create("cpu.txt")
//pprof.StartCPUProfile(f)
//defer pprof.StopCPUProfile()
cmd.Spray()
}