package nuclei

import (
	"bufio"
	"bytes"
	"context"
	"io"
	"os"
	"sync"

	"github.com/projectdiscovery/gologger"
	"github.com/projectdiscovery/nuclei/v3/pkg/authprovider"
	"github.com/projectdiscovery/nuclei/v3/pkg/catalog"
	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/loader"
	"github.com/projectdiscovery/nuclei/v3/pkg/core"
	"github.com/projectdiscovery/nuclei/v3/pkg/input/provider"
	providerTypes "github.com/projectdiscovery/nuclei/v3/pkg/input/types"
	"github.com/projectdiscovery/nuclei/v3/pkg/loader/workflow"
	"github.com/projectdiscovery/nuclei/v3/pkg/output"
	"github.com/projectdiscovery/nuclei/v3/pkg/progress"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/hosterrorscache"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/interactsh"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolinit"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/headless/engine"
	"github.com/projectdiscovery/nuclei/v3/pkg/reporting"
	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
	"github.com/projectdiscovery/nuclei/v3/pkg/templates/signer"
	"github.com/projectdiscovery/nuclei/v3/pkg/types"
	"github.com/projectdiscovery/ratelimit"
	"github.com/projectdiscovery/retryablehttp-go"
	"github.com/projectdiscovery/utils/errkit"
)

// NucleiSDKOptions contains options for nuclei SDK
type NucleiSDKOptions func(e *NucleiEngine) error

var (
	// ErrNotImplemented is returned when a feature is not implemented
	ErrNotImplemented = errkit.New("Not implemented")
	// ErrNoTemplatesAvailable is returned when no templates are available to execute
	ErrNoTemplatesAvailable = errkit.New("No templates available")
	// ErrNoTargetsAvailable is returned when no targets are available to scan
	ErrNoTargetsAvailable = errkit.New("No targets available")
	// ErrOptionsNotSupported is returned when an option is not supported in thread safe mode
	ErrOptionsNotSupported = errkit.New("Option not supported in thread safe mode")
)

type engineMode uint

const (
	singleInstance engineMode = iota
	threadSafe
)

// NucleiEngine is the Engine/Client for nuclei which
// runs scans using templates and returns results
type NucleiEngine struct {
	// user options
	resultCallbacks             []func(event *output.ResultEvent)
	onFailureCallback           func(event *output.InternalEvent)
	disableTemplatesAutoUpgrade bool
	enableStats                 bool
	onUpdateAvailableCallback   func(newVersion string)

	// ready-status fields
	templatesLoaded bool

	// unexported core fields
	ctx              context.Context
	interactshClient *interactsh.Client
	catalog          catalog.Catalog
	rateLimiter      *ratelimit.Limiter
	store            *loader.Store
	httpxClient      providerTypes.InputLivenessProbe
	inputProvider    provider.InputProvider
	engine           *core.Engine
	mode             engineMode
	browserInstance  *engine.Browser
	httpClient       *retryablehttp.Client
	parser           *templates.Parser
	ownsParser       bool // true when the engine created the parser and may purge it on Close
	authprovider     authprovider.AuthProvider

	// unexported meta options
	opts           *types.Options
	interactshOpts *interactsh.Options
	hostErrCache   *hosterrorscache.Cache
	customWriter   output.Writer
	customProgress progress.Progress
	rc             reporting.Client
	reportingOpts  *reporting.Options
	executerOpts   *protocols.ExecutorOptions

	// Logger instance for the engine
	Logger *gologger.Logger

	// Temporary directory for SDK-managed template files
	tmpDir string
}

// LoadAllTemplates loads all nuclei template based on given options
func (e *NucleiEngine) LoadAllTemplates() error {
	workflowLoader, err := workflow.NewLoader(e.executerOpts)
	if err != nil {
		return errkit.Wrapf(err, "Could not create workflow loader: %s", err)
	}
	e.executerOpts.WorkflowLoader = workflowLoader

	e.store, err = loader.New(loader.NewConfig(e.opts, e.catalog, e.executerOpts))
	if err != nil {
		return errkit.Wrapf(err, "Could not create loader client: %s", err)
	}
	if err := e.store.Load(); err != nil {
		return errkit.Wrapf(err, "Could not load templates: %s", err)
	}
	e.templatesLoaded = true
	return nil
}

// GetTemplates returns all nuclei templates that are loaded
func (e *NucleiEngine) GetTemplates() []*templates.Template {
	if !e.templatesLoaded {
		_ = e.LoadAllTemplates()
	}
	return e.store.Templates()
}

// GetWorkflows returns all nuclei workflows that are loaded
func (e *NucleiEngine) GetWorkflows() []*templates.Template {
	if !e.templatesLoaded {
		_ = e.LoadAllTemplates()
	}
	return e.store.Workflows()
}

// LoadTargets(urls/domains/ips only) adds targets to the nuclei engine
func (e *NucleiEngine) LoadTargets(targets []string, probeNonHttp bool) {
	for _, target := range targets {
		if probeNonHttp {
			_ = e.inputProvider.SetWithProbe(e.opts.ExecutionId, target, e.httpxClient)
		} else {
			e.inputProvider.Set(e.opts.ExecutionId, target)
		}
	}
}

// LoadTargetsFromReader adds targets(urls/domains/ips only) from reader to the nuclei engine
func (e *NucleiEngine) LoadTargetsFromReader(reader io.Reader, probeNonHttp bool) {
	buff := bufio.NewScanner(reader)
	for buff.Scan() {
		if probeNonHttp {
			_ = e.inputProvider.SetWithProbe(e.opts.ExecutionId, buff.Text(), e.httpxClient)
		} else {
			e.inputProvider.Set(e.opts.ExecutionId, buff.Text())
		}
	}
}

// LoadTargetsWithHttpData loads targets that contain http data from file it currently supports
// multiple formats like burp xml,openapi,swagger,proxify json
// Note: this is mutually exclusive with LoadTargets and LoadTargetsFromReader
func (e *NucleiEngine) LoadTargetsWithHttpData(filePath string, filemode string) error {
	e.opts.TargetsFilePath = filePath
	e.opts.InputFileMode = filemode
	httpProvider, err := provider.NewInputProvider(provider.InputOptions{Options: e.opts})
	if err != nil {
		e.opts.TargetsFilePath = ""
		e.opts.InputFileMode = ""
		return err
	}
	e.inputProvider = httpProvider
	return nil
}

// GetExecuterOptions returns the nuclei executor options
func (e *NucleiEngine) GetExecuterOptions() *protocols.ExecutorOptions {
	return e.executerOpts
}

// ParseTemplate parses a template from given data
// template verification status can be accessed from template.Verified
func (e *NucleiEngine) ParseTemplate(data []byte) (*templates.Template, error) {
	return templates.ParseTemplateFromReader(bytes.NewReader(data), nil, e.executerOpts)
}

// SignTemplate signs the tempalate using given signer
func (e *NucleiEngine) SignTemplate(tmplSigner *signer.TemplateSigner, data []byte) ([]byte, error) {
	tmpl, err := e.ParseTemplate(data)
	if err != nil {
		return data, err
	}
	if tmpl.Verified {
		// already signed
		return data, nil
	}
	if len(tmpl.Workflows) > 0 {
		return data, templates.ErrNotATemplate
	}
	signatureData, err := tmplSigner.Sign(data, tmpl)
	if err != nil {
		return data, err
	}
	_, content := signer.ExtractSignatureAndContent(data)
	buff := bytes.NewBuffer(content)
	buff.WriteString("\n" + signatureData)
	return buff.Bytes(), err
}

func (e *NucleiEngine) closeInternal() {
	if e.interactshClient != nil {
		e.interactshClient.Close()
	}
	if e.rc != nil {
		e.rc.Close()
	}
	if e.customWriter != nil {
		e.customWriter.Close()
	}
	if e.customProgress != nil {
		e.customProgress.Stop()
	}
	if e.hostErrCache != nil {
		e.hostErrCache.Close()
	}
	if e.executerOpts.RateLimiter != nil {
		e.executerOpts.RateLimiter.Stop()
	}
	if e.rateLimiter != nil {
		e.rateLimiter.Stop()
	}
	if e.inputProvider != nil {
		e.inputProvider.Close()
	}
	if e.browserInstance != nil {
		e.browserInstance.Close()
	}
	if e.httpxClient != nil {
		_ = e.httpxClient.Close()
	}
	if e.tmpDir != "" {
		_ = os.RemoveAll(e.tmpDir)
	}
	if e.opts != nil {
		generators.ClearOptionsPayloadMap(e.opts)
	}
	// Purge the template caches (compiled templates are heap-heavy) so a
	// long-running embedder does not retain them for the process lifetime.
	// Only do this when the engine created the parser; a caller-supplied
	// parser is an opt-in shared cache the caller owns.
	if e.ownsParser && e.parser != nil {
		e.parser.Purge()
	}
}

// Close all resources used by nuclei engine
func (e *NucleiEngine) Close() {
	e.closeInternal()
	protocolinit.Close(e.opts.ExecutionId)
}

// ExecuteCallbackWithCtx executes templates on targets and calls callback on each result(only if results are found)
// enable matcher-status option if you expect this callback to be called for all results regardless if it matched or not
func (e *NucleiEngine) ExecuteCallbackWithCtx(ctx context.Context, callback ...func(event *output.ResultEvent)) error {
	if !e.templatesLoaded {
		_ = e.LoadAllTemplates()
	}
	if len(e.store.Templates()) == 0 && len(e.store.Workflows()) == 0 {
		return ErrNoTemplatesAvailable
	}
	if e.inputProvider.Count() == 0 {
		return ErrNoTargetsAvailable
	}

	filtered := []func(event *output.ResultEvent){}
	for _, cb := range callback {
		if cb != nil {
			filtered = append(filtered, cb)
		}
	}
	e.resultCallbacks = append(e.resultCallbacks, filtered...)

	templatesAndWorkflows := append(e.store.Templates(), e.store.Workflows()...)
	if len(templatesAndWorkflows) == 0 {
		return ErrNoTemplatesAvailable
	}

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		_ = e.engine.ExecuteScanWithOpts(ctx, templatesAndWorkflows, e.inputProvider, false)
	}()

	// wait for context to be cancelled
	select {
	case <-ctx.Done():
		<-wait(&wg) // wait for scan to finish
		return ctx.Err()
	case <-wait(&wg):
		// scan finished
	}
	return nil
}

// ExecuteWithCallback is same as ExecuteCallbackWithCtx but with default context
// Note this is deprecated and will be removed in future major release
func (e *NucleiEngine) ExecuteWithCallback(callback ...func(event *output.ResultEvent)) error {
	ctx := context.Background()
	if e.ctx != nil {
		ctx = e.ctx
	}
	return e.ExecuteCallbackWithCtx(ctx, callback...)
}

// Options return nuclei Type Options
func (e *NucleiEngine) Options() *types.Options {
	return e.opts
}

// Engine returns core Executer of nuclei
func (e *NucleiEngine) Engine() *core.Engine {
	return e.engine
}

// Store returns store of nuclei
func (e *NucleiEngine) Store() *loader.Store {
	return e.store
}

// NewNucleiEngineCtx creates a new nuclei engine instance with given context
func NewNucleiEngineCtx(ctx context.Context, options ...NucleiSDKOptions) (*NucleiEngine, error) {
	// default options
	defaultOptions := types.DefaultOptions()
	e := &NucleiEngine{
		opts:   defaultOptions,
		mode:   singleInstance,
		ctx:    ctx,
		Logger: defaultOptions.Logger,
	}
	for _, option := range options {
		if err := option(e); err != nil {
			return nil, err
		}
	}
	if err := e.init(ctx); err != nil {
		return nil, err
	}
	return e, nil
}

// Deprecated: use NewNucleiEngineCtx instead
func NewNucleiEngine(options ...NucleiSDKOptions) (*NucleiEngine, error) {
	return NewNucleiEngineCtx(context.Background(), options...)
}

// GetParser returns the template parser with cache
func (e *NucleiEngine) GetParser() *templates.Parser {
	return e.parser
}

// wait for a waitgroup to finish
func wait(wg *sync.WaitGroup) <-chan struct{} {
	ch := make(chan struct{})
	go func() {
		defer close(ch)
		wg.Wait()
	}()
	return ch
}

// GetClusterTemplateIDs returns the template IDs for a given cluster ID
// Returns nil if the cluster ID doesn't exist or engine hasn't executed yet
func (e *NucleiEngine) GetClusterTemplateIDs(clusterID string) []string {
	if e.executerOpts == nil || e.executerOpts.ClusterMappings == nil {
		return nil
	}
	templateIDs, ok := e.executerOpts.ClusterMappings.Get(clusterID)
	if !ok {
		return nil
	}
	return templateIDs
}

// GetAllClusterMappings returns all cluster mappings
// Returns nil if engine hasn't executed yet
func (e *NucleiEngine) GetAllClusterMappings() map[string][]string {
	if e.executerOpts == nil || e.executerOpts.ClusterMappings == nil {
		return nil
	}

	return e.executerOpts.ClusterMappings.GetAll()
}