package templates

import (
	"fmt"
	"io"
	"reflect"
	"sync"
	"sync/atomic"

	"github.com/logrusorgru/aurora/v4"
	"github.com/pkg/errors"
	"github.com/projectdiscovery/nuclei/v3/pkg/utils/yaml"

	"github.com/projectdiscovery/gologger"
	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
	"github.com/projectdiscovery/nuclei/v3/pkg/js/compiler"
	"github.com/projectdiscovery/nuclei/v3/pkg/model/types/severity"
	"github.com/projectdiscovery/nuclei/v3/pkg/operators"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/globalmatchers"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/offlinehttp"
	"github.com/projectdiscovery/nuclei/v3/pkg/templates/signer"
	"github.com/projectdiscovery/nuclei/v3/pkg/tmplexec"
	"github.com/projectdiscovery/nuclei/v3/pkg/utils"
	"github.com/projectdiscovery/nuclei/v3/pkg/utils/json"
	"github.com/projectdiscovery/utils/errkit"
	stringsutil "github.com/projectdiscovery/utils/strings"
)

var (
	ErrCreateTemplateExecutor          = errors.New("cannot create template executer")
	ErrIncompatibleWithOfflineMatching = errors.New("template can't be used for offline matching")
	// track how many templates are verified and by which signer
	SignatureStats = map[string]*atomic.Uint64{}
)

const (
	Unsigned = "unsigned"
)

func init() {
	for _, verifier := range signer.DefaultTemplateVerifiers {
		SignatureStats[verifier.Identifier()] = &atomic.Uint64{}
	}
	SignatureStats[Unsigned] = &atomic.Uint64{}
}

// updateRequestOptions updates options for all request types in a template
func updateRequestOptions(template *Template) {
	for i, r := range template.RequestsDNS {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsDNS[i] = &rCopy
	}
	for i, r := range template.RequestsHTTP {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsHTTP[i] = &rCopy
	}
	for i, r := range template.RequestsCode {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsCode[i] = &rCopy
	}
	for i, r := range template.RequestsFile {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsFile[i] = &rCopy
	}
	for i, r := range template.RequestsHeadless {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsHeadless[i] = &rCopy
	}
	for i, r := range template.RequestsNetwork {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsNetwork[i] = &rCopy
	}
	for i, r := range template.RequestsJavascript {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsJavascript[i] = &rCopy
	}
	for i, r := range template.RequestsSSL {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsSSL[i] = &rCopy
	}
	for i, r := range template.RequestsWHOIS {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsWHOIS[i] = &rCopy
	}
	for i, r := range template.RequestsWebsocket {
		rCopy := *r
		rCopy.UpdateOptions(template.Options)
		template.RequestsWebsocket[i] = &rCopy
	}
}

// parseFromSource parses a template from source with caching support
func parseFromSource(filePath string, preprocessor Preprocessor, options *protocols.ExecutorOptions, parser *Parser) (*Template, error) {
	reader, err := utils.ReaderFromPathOrURL(filePath, options.Catalog)
	if err != nil {
		return nil, err
	}

	defer func() {
		_ = reader.Close()
	}()

	options = options.Copy()
	options.TemplatePath = filePath

	template, err := ParseTemplateFromReader(reader, preprocessor, options)
	if err != nil {
		return nil, err
	}

	if template.isGlobalMatchersEnabled() {
		item := &globalmatchers.Item{
			TemplateID:   template.ID,
			TemplatePath: filePath,
			TemplateInfo: template.Info,
		}

		for _, request := range template.RequestsHTTP {
			item.Operators = append(item.Operators, request.CompiledOperators)
		}

		options.GlobalMatchers.AddOperator(item)

		return nil, nil
	}

	// Compile the workflow request
	if len(template.Workflows) > 0 {
		compiled := &template.Workflow

		compileWorkflow(filePath, preprocessor, options, compiled, options.WorkflowLoader)
		template.CompiledWorkflow = compiled
		template.CompiledWorkflow.Options = options
	}

	template.Path = filePath
	if !options.DoNotCache {
		parser.compiledTemplatesCache.StoreWithoutRaw(filePath, template, err)
	}

	return template, nil
}

// getParser returns a cached parser instance
func getParser(options *protocols.ExecutorOptions) *Parser {
	parser, ok := options.Parser.(*Parser)
	if !ok || parser == nil {
		panic("invalid parser")
	}

	return parser
}

// Parse parses a yaml request template file
// TODO make sure reading from the disk the template parsing happens once: see parsers.ParseTemplate vs templates.Parse
func Parse(filePath string, preprocessor Preprocessor, options *protocols.ExecutorOptions) (*Template, error) {
	parser := getParser(options)

	if !options.DoNotCache {
		if value, _, _ := parser.compiledTemplatesCache.Has(filePath); value != nil {
			// Copy the template, apply new options, and recompile requests
			tplCopy := *value
			newBase := options.Copy()
			newBase.TemplateID = tplCopy.Options.TemplateID
			newBase.TemplatePath = tplCopy.Options.TemplatePath
			newBase.TemplateInfo = tplCopy.Options.TemplateInfo
			newBase.TemplateVerifier = tplCopy.Options.TemplateVerifier
			newBase.RawTemplate = tplCopy.Options.RawTemplate

			if tplCopy.Options.Variables.Len() > 0 {
				newBase.Variables = tplCopy.Options.Variables
			}

			if len(tplCopy.Options.Constants) > 0 {
				newBase.Constants = tplCopy.Options.Constants
			}

			tplCopy.Options = newBase
			tplCopy.Options.ApplyNewEngineOptions(options)

			if tplCopy.CompiledWorkflow != nil {
				tplCopy.CompiledWorkflow.Options.ApplyNewEngineOptions(options)
				for _, w := range tplCopy.CompiledWorkflow.Workflows {
					for _, ex := range w.Executers {
						ex.Options.ApplyNewEngineOptions(options)
					}
				}
			}

			// Update options for all request types
			updateRequestOptions(&tplCopy)
			template := &tplCopy

			if template.isGlobalMatchersEnabled() {
				item := &globalmatchers.Item{
					TemplateID:   template.ID,
					TemplatePath: filePath,
					TemplateInfo: template.Info,
				}

				for _, request := range template.RequestsHTTP {
					item.Operators = append(item.Operators, request.CompiledOperators)
				}

				options.GlobalMatchers.AddOperator(item)

				return nil, nil
			}

			// Compile the workflow request
			if len(template.Workflows) > 0 {
				compiled := &template.Workflow
				compileWorkflow(filePath, preprocessor, tplCopy.Options, compiled, tplCopy.Options.WorkflowLoader)
				template.CompiledWorkflow = compiled
				template.CompiledWorkflow.Options = tplCopy.Options
			}

			if isCachedTemplateValid(template) {
				// options.Logger.Error().Msgf("returning cached template %s after recompiling %d requests", tplCopy.Options.TemplateID, tplCopy.Requests())
				return template, nil
			}

			// else: fallthrough to re-parse template from scratch
		}
	}

	return parseFromSource(filePath, preprocessor, options, parser)
}

// isGlobalMatchersEnabled checks if any of requests in the template
// have global matchers enabled. It iterates through all requests and
// returns true if at least one request has global matchers enabled;
// otherwise, it returns false. If global matchers templates are not
// enabled in the options, the method will immediately return false.
//
// Note: This method only checks the `RequestsHTTP`
// field of the template, which is specific to http-protocol-based
// templates.
//
// TODO: support all protocols.
func (template *Template) isGlobalMatchersEnabled() bool {
	caps := CapabilitiesFromOptions(template.Options.Options)
	if !caps.Has(CapabilityGlobalMatchers) {
		return false
	}

	return template.requiresGlobalMatchers()
}

// parseSelfContainedRequests parses the self contained template requests.
func (template *Template) parseSelfContainedRequests() {
	if template.Signature.Value.String() != "" {
		for _, request := range template.RequestsHTTP {
			request.Signature = template.Signature
		}
	}
	if !template.SelfContained {
		return
	}
	for _, request := range template.RequestsHTTP {
		request.SelfContained = true
	}
	for _, request := range template.RequestsNetwork {
		request.SelfContained = true
	}
	for _, request := range template.RequestsHeadless {
		request.SelfContained = true
	}
}

// Requests returns the total request count for the template
func (template *Template) Requests() int {
	return len(template.RequestsDNS) +
		len(template.RequestsHTTP) +
		len(template.RequestsFile) +
		len(template.RequestsNetwork) +
		len(template.RequestsHeadless) +
		len(template.Workflows) +
		len(template.RequestsSSL) +
		len(template.RequestsWebsocket) +
		len(template.RequestsWHOIS) +
		len(template.RequestsCode) +
		len(template.RequestsJavascript)
}

// compileProtocolRequests compiles all the protocol requests for the template
func (template *Template) compileProtocolRequests(options *protocols.ExecutorOptions) error {
	templateRequests := template.Requests()

	if templateRequests == 0 {
		return fmt.Errorf("no requests defined for %s", template.ID)
	}

	if options.Options.OfflineHTTP {
		return template.compileOfflineHTTPRequest(options)
	}

	var requests []protocols.Request
	caps := CapabilitiesFromOptions(options.Options)

	if template.hasMultipleRequests() {
		// when multiple requests are present preserve the order of requests and protocols
		// which is already done during unmarshalling
		requests = template.RequestsQueue
		if options.Flow == "" {
			options.IsMultiProtocol = true
		}
	} else {
		if template.HasDNSRequest() {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsDNS)...)
		}
		if template.HasFileRequest() {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsFile)...)
		}
		if template.HasNetworkRequest() {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsNetwork)...)
		}
		if template.HasHTTPRequest() {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsHTTP)...)
		}
		if template.HasHeadlessRequest() && caps.Has(CapabilityHeadless) {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsHeadless)...)
		}
		if template.HasSSLRequest() {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsSSL)...)
		}
		if template.HasWebsocketRequest() {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsWebsocket)...)
		}
		if template.HasWHOISRequest() {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsWHOIS)...)
		}
		if template.HasCodeRequest() && caps.Has(CapabilityCode) {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsCode)...)
		}
		if template.HasJavascriptRequest() {
			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsJavascript)...)
		}
	}
	var err error
	template.Executer, err = tmplexec.NewTemplateExecuter(requests, options)
	return err
}

// convertRequestToProtocolsRequest is a convenience wrapper to convert
// arbitrary interfaces which are slices of requests from the template to a
// slice of protocols.Request interface items.
func (template *Template) convertRequestToProtocolsRequest(requests interface{}) []protocols.Request {
	switch reflect.TypeOf(requests).Kind() {
	case reflect.Slice:
		s := reflect.ValueOf(requests)

		requestSlice := make([]protocols.Request, s.Len())
		for i := 0; i < s.Len(); i++ {
			value := s.Index(i)
			valueInterface := value.Interface()
			requestSlice[i] = valueInterface.(protocols.Request)
		}
		return requestSlice
	}
	return nil
}

// compileOfflineHTTPRequest iterates all requests if offline http mode is
// specified and collects all matchers for all the base request templates
// (those with URL {{BaseURL}} and it's slash variation.)
func (template *Template) compileOfflineHTTPRequest(options *protocols.ExecutorOptions) error {
	operatorsList := []*operators.Operators{}

mainLoop:
	for _, req := range template.RequestsHTTP {
		hasPaths := len(req.Path) > 0
		if !hasPaths {
			break mainLoop
		}
		for _, path := range req.Path {
			pathIsBaseURL := stringsutil.EqualFoldAny(path, "{{BaseURL}}", "{{BaseURL}}/", "/")
			if !pathIsBaseURL {
				break mainLoop
			}
		}
		operatorsList = append(operatorsList, &req.Operators)
	}
	if len(operatorsList) > 0 {
		options.Operators = operatorsList
		var err error
		template.Executer, err = tmplexec.NewTemplateExecuter([]protocols.Request{&offlinehttp.Request{}}, options)
		if err != nil {
			// it seems like flow executor cannot be used for offline http matching (ex:http(1) && http(2))
			return ErrIncompatibleWithOfflineMatching
		}
		return err
	}

	return ErrIncompatibleWithOfflineMatching
}

// ParseTemplateFromReader parses a template from an [io.Reader] with optional
// preprocessing.
func ParseTemplateFromReader(reader io.Reader, preprocessor Preprocessor, options *protocols.ExecutorOptions) (*Template, error) {
	data, err := io.ReadAll(reader)
	if err != nil {
		return nil, err
	}

	// a preprocessor is a variable like
	// {{randstr}} which is replaced before unmarshalling
	// as it is known to be a random static value per template
	hasPreprocessor := false
	allPreprocessors := getPreprocessors(preprocessor)
	for _, preprocessor := range allPreprocessors {
		if preprocessor.Exists(data) {
			hasPreprocessor = true
			break
		}
	}

	if !hasPreprocessor {
		// if no preprocessors exists parse template and exit
		template, err := parseTemplate(data, options)
		if err != nil {
			return nil, err
		}
		if !template.Verified && len(template.Workflows) == 0 {
			if config.DefaultConfig.LogAllEvents {
				gologger.DefaultLogger.Print().Msgf("[%v] Template %s is not signed or tampered\n", aurora.Yellow("WRN").String(), template.ID)
			}
		}
		return template, nil
	}

	// if preprocessor is required / exists in this template
	// expand all preprocessors, parse once, then verify against original data
	generatedConstants := map[string]interface{}{}

	// ==== execute preprocessors ======
	processedData := data
	for _, v := range allPreprocessors {
		var replaced map[string]interface{}
		processedData, replaced = v.ProcessNReturnData(processedData)
		// preprocess kind of act like a constant and are generated while loading
		// and stay constant for the template lifecycle
		generatedConstants = generators.MergeMaps(generatedConstants, replaced)
	}

	template, err := parseTemplateNoVerify(processedData, options)
	if err != nil {
		return nil, err
	}

	// add generated constants to constants map and executer options
	template.Constants = generators.MergeMaps(template.Constants, generatedConstants)
	template.Options.Constants = template.Constants
	applyTemplateVerification(template, data)

	if !template.Verified && len(template.Workflows) == 0 {
		// workflows are not signed by default
		if config.DefaultConfig.LogAllEvents {
			gologger.DefaultLogger.Print().Msgf("[%v] Template %s is not signed or tampered\n", aurora.Yellow("WRN").String(), template.ID)
		}
	}

	return template, nil
}

// parseTemplate parses the template and applies verification.
func parseTemplate(data []byte, srcOptions *protocols.ExecutorOptions) (*Template, error) {
	template, err := parseTemplateNoVerify(data, srcOptions)
	if err != nil {
		return nil, err
	}
	applyTemplateVerification(template, data)

	return template, nil
}

// parseTemplateNoVerify parses the template without applying any verification.
func parseTemplateNoVerify(data []byte, srcOptions *protocols.ExecutorOptions) (*Template, error) {
	// Create a copy of the options specifically for this template
	options := srcOptions.Copy()

	template := &Template{}
	var err error
	switch config.GetTemplateFormatFromExt(template.Path) {
	case config.JSON:
		err = json.Unmarshal(data, template)
	case config.YAML:
		err = yaml.Unmarshal(data, template)
	default:
		// assume its yaml
		if err = yaml.Unmarshal(data, template); err != nil {
			return nil, err
		}
	}
	if err != nil {
		return nil, errkit.Wrapf(err, "failed to parse %s", template.Path)
	}

	if utils.IsBlank(template.Info.Name) {
		return nil, errors.New("no template name field provided")
	}
	if template.Info.Authors.IsEmpty() {
		return nil, errors.New("no template author field provided")
	}

	numberOfWorkflows := len(template.Workflows)
	if numberOfWorkflows > 0 && numberOfWorkflows != template.Requests() {
		return nil, errors.New("workflows cannot have other protocols")
	}

	// use default unknown severity
	if len(template.Workflows) == 0 {
		if template.Info.SeverityHolder.Severity == severity.Undefined {
			// set unknown severity with counter and forced warning
			template.Info.SeverityHolder.Severity = severity.Unknown
			if options.Options.Validate {
				// when validating return error
				return nil, errors.New("no template severity field provided")
			}
		}
	}

	// Setting up variables regarding template metadata
	options.TemplateID = template.ID
	options.TemplateInfo = template.Info
	options.StopAtFirstMatch = template.StopAtFirstMatch

	if template.Variables.Len() > 0 {
		options.Variables = template.Variables
	}

	// if more than 1 request per protocol exist we add request id to protocol request
	// since in template context we have proto_prefix for each protocol it is overwritten
	// if request id is not present
	template.validateAllRequestIDs()

	// create empty context args for template scope
	options.CreateTemplateCtxStore()
	options.ProtocolType = template.Type()
	options.Constants = template.Constants

	// initialize the js compiler if missing
	if options.JsCompiler == nil {
		options.JsCompiler = GetJsCompiler() // this is a singleton
	}

	template.Options = options
	// If no requests, and it is also not a workflow, return error.
	if template.Requests() == 0 {
		return nil, fmt.Errorf("no requests defined for %s", template.ID)
	}

	// load `flow` and `source` in code protocol from file
	// if file is referenced instead of actual source code
	if err := template.ImportFileRefs(template.Options); err != nil {
		return nil, errkit.Wrapf(err, "failed to load file refs for %s", template.ID)
	}

	if err := template.compileProtocolRequests(template.Options); err != nil {
		return nil, err
	}

	if template.Executer != nil {
		if err := template.Executer.Compile(); err != nil {
			return nil, errors.Wrap(err, "could not compile request")
		}
		template.TotalRequests = template.Executer.Requests()
	}
	if template.Executer == nil && template.CompiledWorkflow == nil {
		return nil, ErrCreateTemplateExecutor
	}
	template.parseSelfContainedRequests()

	return template, nil
}

// applyTemplateVerification verifies a parsed template against the provided data.
func applyTemplateVerification(template *Template, data []byte) {
	if template == nil || template.Options == nil {
		return
	}

	options := template.Options
	// check if the template is verified
	// only valid templates can be verified or signed
	if options.TemplateVerificationCallback != nil && options.TemplatePath != "" {
		if cached := options.TemplateVerificationCallback(options.TemplatePath); cached != nil {
			template.Verified = cached.Verified
			template.TemplateVerifier = cached.Verifier
			options.TemplateVerifier = cached.Verifier
			//nolint
			if !(template.Verified && template.TemplateVerifier == "projectdiscovery/nuclei-templates") {
				template.Options.RawTemplate = data
			}
			return
		}
	}

	var verifier *signer.TemplateSigner
	for _, verifier = range signer.DefaultTemplateVerifiers {
		template.Verified, _ = verifier.Verify(data, template)
		if config.DefaultConfig.LogAllEvents {
			gologger.Verbose().Msgf("template %v verified by %s : %v", template.ID, verifier.Identifier(), template.Verified)
		}
		if template.Verified {
			template.TemplateVerifier = verifier.Identifier()
			break
		}
	}
	options.TemplateVerifier = template.TemplateVerifier

	//nolint
	if !(template.Verified && verifier.Identifier() == "projectdiscovery/nuclei-templates") {
		template.Options.RawTemplate = data
	}
}

// isCachedTemplateValid validates that a cached template is still usable after
// option updates
func isCachedTemplateValid(template *Template) bool {
	// no requests or workflows
	if template.Requests() == 0 && len(template.Workflows) == 0 {
		return false
	}

	// options not initialized
	if template.Options == nil {
		return false
	}

	// executer not available for non-workflow template
	if len(template.Workflows) == 0 && template.Executer == nil {
		return false
	}

	// compiled workflow not available
	if len(template.Workflows) > 0 && template.CompiledWorkflow == nil {
		return false
	}

	// template ID mismatch
	if template.Options.TemplateID != template.ID {
		return false
	}

	// executer exists but no requests or flow available
	if template.Executer != nil {
		// NOTE(dwisiswant0): This is a basic sanity check since we can't access
		// private fields, but we can check requests tho
		if template.Requests() == 0 && template.Options.Flow == "" {
			return false
		}
	}

	if template.Options.Options == nil {
		return false
	}

	return true
}

var (
	jsCompiler     *compiler.Compiler
	jsCompilerOnce = sync.OnceFunc(func() {
		jsCompiler = compiler.New()
	})
)

func GetJsCompiler() *compiler.Compiler {
	jsCompilerOnce()
	return jsCompiler
}
