package signer

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/md5"
	"crypto/rand"
	"crypto/sha256"
	"encoding/gob"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"

	"github.com/projectdiscovery/gologger"
	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
	"github.com/projectdiscovery/utils/errkit"
)

var (
	ErrUnknownAlgorithm = errors.New("unknown algorithm")
	SignaturePattern    = "# digest: "
	SignatureFmt        = SignaturePattern + "%x" + ":%v" // `#digest: <signature>:<fragment>`
)

// ExtractSignatureAndContent extracts the signature (if present) and returns the content without the signature
func ExtractSignatureAndContent(data []byte) (signature, content []byte) {
	dataStr := string(data)
	if idx := strings.LastIndex(dataStr, SignaturePattern); idx != -1 {
		signature = []byte(strings.TrimSpace(dataStr[idx:]))
		content = bytes.TrimSpace(data[:idx])
	} else {
		content = data
	}
	content = bytes.TrimSpace(content)
	return signature, content
}

// SignableTemplate is a template that can be signed
type SignableTemplate interface {
	// GetFileImports returns a list of files that are imported by the template
	GetFileImports() []string
	// HasCodeProtocol returns true if the template has a code protocol section
	HasCodeProtocol() bool
}

type fileImportContentProvider interface {
	GetFileImportContents() ([][]byte, bool)
}

type javascriptSignableTemplate interface {
	HasJavascriptRequest(...int) bool
}

type TemplateSigner struct {
	sync.Once
	handler     *KeyHandler
	fragment    string
	fragmentErr error
}

// Identifier returns the identifier for the template signer
func (t *TemplateSigner) Identifier() string {
	return t.handler.cert.Subject.CommonName
}

// fragment is optional part of signature that is used to identify the user
// who signed the template via md5 hash of the public key x-coordinate
func (t *TemplateSigner) GetUserFragment() string {
	fragment, _ := t.userFragment()
	return fragment
}

func (t *TemplateSigner) userFragment() (string, error) {
	// wrap with sync.Once to reduce unnecessary md5 hashing
	t.Do(func() {
		t.fragment, t.fragmentErr = publicKeyFragment(t.handler.ecdsaPubKey)
	})
	return t.fragment, t.fragmentErr
}

func publicKeyFragment(publicKey *ecdsa.PublicKey) (string, error) {
	if publicKey == nil {
		return "", nil
	}
	publicKeyBytes, err := publicKey.Bytes()
	if err != nil {
		return "", fmt.Errorf("encode ecdsa public key: %w", err)
	}
	if len(publicKeyBytes) < 3 || publicKeyBytes[0] != 4 || (len(publicKeyBytes)-1)%2 != 0 {
		return "", fmt.Errorf("invalid uncompressed ecdsa public key")
	}
	xCoordinateLength := (len(publicKeyBytes) - 1) / 2
	// Keep the old fragment stable: big.Int.Bytes omitted leading zero bytes.
	xCoordinateBytes := bytes.TrimLeft(publicKeyBytes[1:1+xCoordinateLength], "\x00")
	hashed := md5.Sum(xCoordinateBytes)
	return fmt.Sprintf("%x", hashed), nil
}

// Sign signs the given template with the template signer and returns the signature
func (t *TemplateSigner) Sign(data []byte, tmpl SignableTemplate) (string, error) {
	existingSignature, content := ExtractSignatureAndContent(data)
	content = normalizeTemplateContentForSignature(content)

	// Executable templates can only be re-signed by the current signer.
	hasJavascript := false
	if javascriptTemplate, ok := tmpl.(javascriptSignableTemplate); ok {
		hasJavascript = javascriptTemplate.HasJavascriptRequest()
	}

	if tmpl.HasCodeProtocol() || hasJavascript {
		if len(existingSignature) > 0 {
			arr := strings.SplitN(string(existingSignature), ":", 3)
			if len(arr) == 2 {
				// signature has no fragment
				return "", errkit.New("re-signing executable templates are not allowed for security reasons.")
			}

			if len(arr) == 3 {
				// signature has fragment verify if it is equal to current fragment
				fragment, err := t.userFragment()
				if err != nil {
					return "", err
				}

				if fragment != arr[2] {
					return "", errkit.New("re-signing executable templates are not allowed for security reasons.")
				}
			}
		}
	}

	buff, err := templateContentWithImports(content, tmpl)
	if err != nil {
		return "", err
	}

	signatureData, err := t.sign(buff.Bytes())
	if err != nil {
		return "", err
	}

	return signatureData, nil
}

// Signs given data with the template signer
// Note: this should not be used for signing templates as file references
// in templates are not processed use template.SignTemplate() instead
func (t *TemplateSigner) sign(data []byte) (string, error) {
	dataHash := sha256.Sum256(data)

	ecdsaSignature, err := ecdsa.SignASN1(rand.Reader, t.handler.ecdsaKey, dataHash[:])
	if err != nil {
		return "", err
	}

	var signatureData bytes.Buffer
	if err := gob.NewEncoder(&signatureData).Encode(ecdsaSignature); err != nil {
		return "", err
	}

	fragment, err := t.userFragment()
	if err != nil {
		return "", err
	}

	return fmt.Sprintf(SignatureFmt, signatureData.Bytes(), fragment), nil
}

// Verify verifies the given template with the template signer
func (t *TemplateSigner) Verify(data []byte, tmpl SignableTemplate) (bool, error) {
	signature, content := ExtractSignatureAndContent(data)
	if len(signature) == 0 {
		return false, errors.New("no signature found")
	}

	if !bytes.HasPrefix(signature, []byte(SignaturePattern)) {
		return false, errors.New("signature must be at the end of the template")
	}

	digestData := bytes.TrimSpace(bytes.TrimPrefix(signature, []byte(SignaturePattern)))

	fragment, err := t.userFragment()
	if err != nil {
		return false, err
	}

	// remove fragment from digest as it is used for re-signing purposes only
	digestString := strings.TrimSuffix(string(digestData), ":"+fragment)

	digest, err := hex.DecodeString(digestString)
	if err != nil {
		return false, err
	}

	content = normalizeTemplateContentForSignature(content)

	buff, err := templateContentWithImports(content, tmpl)
	if err != nil {
		return false, err
	}

	return t.verify(buff.Bytes(), digest)
}

func templateContentWithImports(content []byte, tmpl SignableTemplate) (*bytes.Buffer, error) {
	buff := bytes.NewBuffer(content)
	imports := tmpl.GetFileImports()

	if provider, ok := tmpl.(fileImportContentProvider); ok {
		if importedContents, captured := provider.GetFileImportContents(); captured {
			if len(importedContents) != len(imports) {
				return nil, fmt.Errorf("imported-file content count does not match imported-file path count")
			}

			for _, importedContent := range importedContents {
				buff.WriteRune('\n')
				buff.Write(importedContent)
			}

			return buff, nil
		}
	}

	for _, file := range imports {
		bin, err := os.ReadFile(file)
		if err != nil {
			return nil, err
		}

		buff.WriteRune('\n')
		buff.Write(bin)
	}

	return buff, nil
}

// Fingerprint returns the SHA-256 fingerprint of the signer's public key.
func (t *TemplateSigner) Fingerprint() [sha256.Size]byte {
	if t == nil || t.handler == nil || t.handler.cert == nil {
		return [sha256.Size]byte{}
	}

	return sha256.Sum256(t.handler.cert.RawSubjectPublicKeyInfo)
}

func normalizeTemplateContentForSignature(content []byte) []byte {
	return bytes.ReplaceAll(content, []byte("\r\n"), []byte("\n"))
}

// Verify verifies the given data with the template signer
// Note: this should not be used for verifying templates as file references
// in templates are not processed
func (t *TemplateSigner) verify(data, signatureData []byte) (bool, error) {
	dataHash := sha256.Sum256(data)

	var signature []byte
	if err := gob.NewDecoder(bytes.NewReader(signatureData)).Decode(&signature); err != nil {
		return false, err
	}

	return ecdsa.VerifyASN1(t.handler.ecdsaPubKey, dataHash[:], signature), nil
}

// NewTemplateSigner creates a new signer for signing templates
func NewTemplateSigner(cert, privateKey []byte) (*TemplateSigner, error) {
	handler := &KeyHandler{}

	var err error

	if cert != nil || privateKey != nil {
		handler.UserCert = cert
		handler.PrivateKey = privateKey
	} else {
		err = handler.ReadCert(CertEnvVarName, config.DefaultConfig.GetKeysDir())
		if err == nil {
			err = handler.ReadPrivateKey(PrivateKeyEnvName, config.DefaultConfig.GetKeysDir())
		}
	}

	if err != nil && !SkipGeneratingKeys {
		if err != ErrNoCertificate && err != ErrNoPrivateKey {
			gologger.Info().Msgf("Invalid user cert found : %s\n", err)
		}

		// generating new keys
		handler.GenerateKeyPair()
		if err := handler.SaveToDisk(config.DefaultConfig.GetKeysDir()); err != nil {
			gologger.Fatal().Msgf("could not save generated keys to disk: %s\n", err)
		}

		// do not continue further let user re-run the command
		os.Exit(0)
	} else if err != nil && SkipGeneratingKeys {
		return nil, err
	}

	if err := handler.ParseUserCert(); err != nil {
		return nil, err
	}

	if err := handler.ParsePrivateKey(); err != nil {
		return nil, err
	}

	return &TemplateSigner{
		handler: handler,
	}, nil
}

// NewTemplateSignerFromFiles creates a new signer for signing templates
func NewTemplateSignerFromFiles(cert, privKey string) (*TemplateSigner, error) {
	certData, err := os.ReadFile(cert)
	if err != nil {
		return nil, err
	}

	privKeyData, err := os.ReadFile(privKey)
	if err != nil {
		return nil, err
	}

	return NewTemplateSigner(certData, privKeyData)
}

// NewTemplateSigVerifier creates a new signer for verifying templates
func NewTemplateSigVerifier(cert []byte) (*TemplateSigner, error) {
	handler := &KeyHandler{}
	if cert != nil {
		handler.UserCert = cert
	} else {
		if err := handler.ReadCert(CertEnvVarName, config.DefaultConfig.GetKeysDir()); err != nil {
			return nil, err
		}
	}

	if err := handler.ParseUserCert(); err != nil {
		return nil, err
	}

	return &TemplateSigner{
		handler: handler,
	}, nil
}
