// nuclei-------------------------------------
// @file      : nuclei.go
// @author    : Autumn
// @contact   : rainy-autumn@outlook.com
// @time      : 2024/10/19 16:27
// -------------------------------------------

package nuclei

import (
	"context"
	"errors"
	"fmt"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/contextmanager"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/global"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/handler"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/interfaces"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/results"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/types"
	"github.com/Autumn-27/ScopeSentry-Scan/pkg/logger"
	"github.com/Autumn-27/ScopeSentry-Scan/pkg/utils"
	nuclei "github.com/projectdiscovery/nuclei/v3/lib"
	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
	"github.com/projectdiscovery/nuclei/v3/pkg/output"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

type Plugin struct {
	Name      string
	Module    string
	Parameter string
	PluginId  string
	Result    chan interface{}
	Custom    interface{}
	TaskId    string
	TaskName  string
}

func NewPlugin() *Plugin {
	return &Plugin{
		Name:     "nuclei",
		Module:   "VulnerabilityScan",
		PluginId: "ed93b8af6b72fe54a60efdb932cf6fbc",
	}
}

func (p *Plugin) SetTaskName(name string) {
	p.TaskName = name
}

func (p *Plugin) GetTaskName() string {
	return p.TaskName
}

func (p *Plugin) SetTaskId(id string) {
	p.TaskId = id
}
func (p *Plugin) UnInstall() error {
	return nil
}
func (p *Plugin) GetTaskId() string {
	return p.TaskId
}
func (p *Plugin) SetCustom(cu interface{}) {
	p.Custom = cu
}

func (p *Plugin) GetCustom() interface{} {
	return p.Custom
}
func (p *Plugin) SetPluginId(id string) {
	p.PluginId = id
}

func (p *Plugin) GetPluginId() string {
	return p.PluginId
}

func (p *Plugin) SetResult(ch chan interface{}) {
	p.Result = ch
}

func (p *Plugin) SetName(name string) {
	p.Name = name
}

func (p *Plugin) GetName() string {
	return p.Name
}

func (p *Plugin) SetModule(module string) {
	p.Module = module
}

func (p *Plugin) GetModule() string {
	return p.Module
}

func (p *Plugin) Install() error {
	return nil
}

func (p *Plugin) Check() error {
	return nil
}

func (p *Plugin) SetParameter(args string) {
	p.Parameter = args
}

func (p *Plugin) GetParameter() string {
	return p.Parameter
}

func (p *Plugin) Log(msg string, tp ...string) {
	var logTp string
	if len(tp) > 0 {
		logTp = tp[0] // 使用传入的参数
	} else {
		logTp = "i"
	}
	logger.PluginsLog(fmt.Sprintf("[Plugins %v] %v", p.GetName(), msg), logTp, p.GetModule(), p.GetPluginId())
}

func (p *Plugin) Execute(input interface{}) (interface{}, error) {
	var options []nuclei.NucleiSDKOptions
	var tmplateFilters nuclei.TemplateFilters
	var targets []string
	var tmpTags []string
	switch a := input.(type) {
	case []types.AssetOther:
		for _, assetOther := range a {
			targets = append(targets, assetOther.Host+":"+assetOther.Port)
		}
		tmplateFilters.ExcludeProtocolTypes = "http"
	case []types.AssetHttp:
		for _, assetHttp := range a {
			targets = append(targets, assetHttp.URL)
			tmpTags = append(tmpTags, assetHttp.Technologies...)
		}
		// 如果是http则限制为http的poc
		tmplateFilters.ProtocolTypes = "http"
	default:
		return nil, errors.New("input is not AssetHttp, AssetOther")
	}
	targets = utils.Tools.RemoveStringDuplicates(targets)
	p.Log(fmt.Sprintf("target %v start", targets))
	start := time.Now()
	parameter := p.GetParameter()
	var templates []string
	maxTokens := 100
	duration := 1 * time.Second
	var concurrency nuclei.Concurrency
	concurrency.HostConcurrency = 15
	concurrency.TemplateConcurrency = 15
	concurrency.HeadlessHostConcurrency = 10
	concurrency.HeadlessTemplateConcurrency = 10
	concurrency.JavascriptTemplateConcurrency = 80
	concurrency.TemplatePayloadConcurrency = 15
	concurrency.ProbeConcurrency = 5
	if parameter != "" {
		args, err := utils.Tools.ParseArgs(parameter, "t", "s", "es", "tags", "etags", "rl", "rld", "bs", "c", "hbs", "headc", "jsc", "pc", "prc", "as", "InteractshURL")
		if err != nil {
		} else {
			for key, value := range args {
				if value != "" {
					switch key {
					case "t":
						if value == "*" {
							templates = []string{global.PocDir}
						} else {
							templates = strings.Split(value, ",")
						}
					case "s":
						tmplateFilters.Severity = value
					case "es":
						tmplateFilters.ExcludeSeverities = value
					case "tags":
						tmplateFilters.Tags = append(tmplateFilters.Tags, strings.Split(value, ",")...)
					case "etags":
						tmplateFilters.ExcludeTags = strings.Split(value, ",")
					case "rl":
						maxTokens, _ = strconv.Atoi(value)
					case "rld":
						timeValue, _ := strconv.Atoi(value)
						duration = time.Duration(timeValue) * time.Second
					case "bs":
						bsValue, _ := strconv.Atoi(value)
						concurrency.HostConcurrency = bsValue
					case "c":
						cValue, _ := strconv.Atoi(value)
						concurrency.TemplateConcurrency = cValue
					case "hbs":
						hbsValue, _ := strconv.Atoi(value)
						concurrency.HeadlessHostConcurrency = hbsValue
					case "headc":
						headcValue, _ := strconv.Atoi(value)
						concurrency.HeadlessTemplateConcurrency = headcValue
					case "jsc":
						jscValue, _ := strconv.Atoi(value)
						concurrency.JavascriptTemplateConcurrency = jscValue
					case "pc":
						pcValue, _ := strconv.Atoi(value)
						concurrency.TemplatePayloadConcurrency = pcValue
					case "prc":
						prcValue, _ := strconv.Atoi(value)
						concurrency.ProbeConcurrency = prcValue
					case "as":
						tmplateFilters.Tags = append(tmplateFilters.Tags, tmpTags...)
					//case "InteractshURL":
					//	options = append(options, nuclei.WithInteractshOptions(nuclei.InteractshOpts(interactsh.Options{
					//		ServerURL: value,
					//	})))
					default:
						continue
					}
				}
			}
		}
	}
	// 全局速率限制
	options = append(options, nuclei.WithGlobalRateLimitCtx(context.Background(), maxTokens, duration))

	// 速率限制
	options = append(options, nuclei.WithConcurrency(concurrency))

	// 选择 poc
	if len(templates) != 0 {
		options = append(options, nuclei.WithTemplatesOrWorkflows(nuclei.TemplateSources{Templates: templates}))
	}

	// TmplateFilters 模板过滤
	options = append(options, nuclei.WithTemplateFilters(tmplateFilters))
	ctx := contextmanager.GlobalContextManagers.GetContext(p.GetTaskId())
	callBackFunc := func(event *output.ResultEvent) {
		vulName := event.Info.Name
		level := event.Info.SeverityHolder.Severity.String()
		if level == "" {
			level = "unknown"
		}
		url := event.URL
		if url == "" {
			url = event.Host + event.Port
		}
		tmpResult := types.VulnResult{
			Url:      url,
			VulnId:   strings.TrimSuffix(filepath.Base(event.TemplatePath), ".yaml"),
			VulName:  vulName,
			Matched:  event.Matched,
			Level:    level,
			Time:     utils.Tools.GetTimeNow(),
			Request:  event.Request,
			Response: event.Response,
		}
		tmpResult.TaskName = p.TaskName
		tmpResult.Status = 1
		go results.Handler.Vulnerability(&tmpResult)
	}

	config.DefaultConfig.TemplatesDirectory = filepath.Join(global.PocDir)
	ne := handler.NewNucleiEngine()
	handler.NucleiEngineWg.Add(1)

	ne.GlobalResultCallback(callBackFunc)
	err := ne.ExecuteNucleiWithOptsCtx(ctx, targets, options...)
	if err != nil {
		p.Log(fmt.Sprintf("Nuclei target %v to err: %s", targets, err), "e")
	}

	//ne, err := nuclei.NewNucleiEngineCtx(context.Background(), options...)
	//if err != nil {
	//	p.Log(fmt.Sprintf("Nuclei target %v to err: %s", target, err), "w")
	//	return nil, nil
	//}
	//ne.LoadTargets([]string{target}, false)
	//err = ne.ExecuteWithCallback(callBackFunc)
	//if err != nil {
	//	p.Log(fmt.Sprintf("Nuclei target %v to 2err: %s", target, err), "w")
	//	return nil, nil
	//}
	end := time.Now()
	runDuration := end.Sub(start)
	p.Log(fmt.Sprintf("target %v run time: %v", targets, runDuration))
	handler.NucleiEngineWg.Done()
	time.Sleep(3 * time.Second)
	return nil, nil
}

func (p *Plugin) Clone() interfaces.Plugin {
	return &Plugin{
		Name:     p.Name,
		Module:   p.Module,
		PluginId: p.PluginId,
		Custom:   p.Custom,
		TaskId:   p.TaskId,
	}
}
