Nuclei Template Signature Verification Bypass

@@ -11,7 +11,6 @@ import (
 	"errors"
 	"fmt"
 	"os"
-	"regexp"
 	"strings"
 	"sync"
 
@@ -21,18 +20,21 @@ import (
 )
 
 var (
-	ReDigest            = regexp.MustCompile(`(?m)^#\sdigest:\s.+$`)
 	ErrUnknownAlgorithm = errors.New("unknown algorithm")
 	SignaturePattern    = "# digest: "
 	SignatureFmt        = SignaturePattern + "%x" + ":%v" // `#digest: <signature>:<fragment>`
 )
 
-func RemoveSignatureFromData(data []byte) []byte {
-	return bytes.Trim(ReDigest.ReplaceAll(data, []byte("")), "\n")
-}
-
-func GetSignatureFromData(data []byte) []byte {
-	return ReDigest.Find(data)
+// 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 = []byte(strings.TrimSpace(dataStr[:idx]))
+	} else {
+		content = data
+	}
+	return
 }
 
 // SignableTemplate is a template that can be signed
@@ -69,26 +71,29 @@ func (t *TemplateSigner) GetUserFragment() string {
 
 // 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)
+
 	// while re-signing template check if it has a code protocol
 	// if it does then verify that it is signed by current signer
 	// if not then return error
 	if tmpl.HasCodeProtocol() {
-		sig := GetSignatureFromData(data)
-		arr := strings.SplitN(string(sig), ":", 3)
-		if len(arr) == 2 {
-			// signature has no fragment
-			return "", errorutil.NewWithTag("signer", "re-signing code templates are not allowed for security reasons.")
-		}
-		if len(arr) == 3 {
-			// signature has fragment verify if it is equal to current fragment
-			fragment := t.GetUserFragment()
-			if fragment != arr[2] {
+		if len(existingSignature) > 0 {
+			arr := strings.SplitN(string(existingSignature), ":", 3)
+			if len(arr) == 2 {
+				// signature has no fragment
 				return "", errorutil.NewWithTag("signer", "re-signing code templates are not allowed for security reasons.")
 			}
+			if len(arr) == 3 {
+				// signature has fragment verify if it is equal to current fragment
+				fragment := t.GetUserFragment()
+				if fragment != arr[2] {
+					return "", errorutil.NewWithTag("signer", "re-signing code templates are not allowed for security reasons.")
+				}
+			}
 		}
 	}
 
-	buff := bytes.NewBuffer(RemoveSignatureFromData(data))
+	buff := bytes.NewBuffer(content)
 	// if file has any imports process them
 	for _, file := range tmpl.GetFileImports() {
 		bin, err := os.ReadFile(file)
@@ -123,20 +128,24 @@ func (t *TemplateSigner) sign(data []byte) (string, error) {
 
 // Verify verifies the given template with the template signer
 func (t *TemplateSigner) Verify(data []byte, tmpl SignableTemplate) (bool, error) {
-	digestData := ReDigest.Find(data)
-	if len(digestData) == 0 {
-		return false, errors.New("digest not found")
+	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(digestData, []byte(SignaturePattern)))
+	digestData := bytes.TrimSpace(bytes.TrimPrefix(signature, []byte(SignaturePattern)))
 	// remove fragment from digest as it is used for re-signing purposes only
 	digestString := strings.TrimSuffix(string(digestData), ":"+t.GetUserFragment())
 	digest, err := hex.DecodeString(digestString)
 	if err != nil {
 		return false, err
 	}
 
-	buff := bytes.NewBuffer(RemoveSignatureFromData(data))
+	buff := bytes.NewBuffer(content)
 	// if file has any imports process them
 	for _, file := range tmpl.GetFileImports() {
 		bin, err := os.ReadFile(file)

@@ -0,0 +1,126 @@
+package signer
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	testCertFile = "../../../integration_tests/protocols/keys/ci.crt"
+	testKeyFile  = "../../../integration_tests/protocols/keys/ci-private-key.pem"
+)
+
+type mockSignableTemplate struct {
+	imports []string
+	hasCode bool
+}
+
+func (m *mockSignableTemplate) GetFileImports() []string {
+	return m.imports
+}
+
+func (m *mockSignableTemplate) HasCodeProtocol() bool {
+	return m.hasCode
+}
+
+var signer, _ = NewTemplateSignerFromFiles(testCertFile, testKeyFile)
+
+func TestTemplateSignerSignAndVerify(t *testing.T) {
+	tempDir := t.TempDir()
+
+	tests := []struct {
+		name            string
+		data            []byte
+		tmpl            SignableTemplate
+		wantSignErr     bool
+		wantVerifyErr   bool
+		wantVerified    bool
+		modifyAfterSign func([]byte) []byte
+	}{
+		{
+			name:         "Simple template",
+			data:         []byte("id: test-template\ninfo:\n  name: Test Template"),
+			tmpl:         &mockSignableTemplate{},
+			wantVerified: true,
+		},
+		{
+			name: "Template with imports",
+			data: []byte("id: test-template\ninfo:\n  name: Test Template"),
+			tmpl: &mockSignableTemplate{imports: []string{
+				filepath.Join(tempDir, "import1.yaml"),
+				filepath.Join(tempDir, "import2.yaml"),
+			}},
+			wantVerified: true,
+		},
+		{
+			name:         "Template with code protocol",
+			data:         []byte("id: test-template\ninfo:\n  name: Test Template\n\ncode:\n  - engine: bash\n    source: echo 'Hello, World!'"),
+			tmpl:         &mockSignableTemplate{hasCode: true},
+			wantSignErr:  false,
+			wantVerified: true,
+		},
+		{
+			name: "Tampered template",
+			data: []byte("id: test-template\ninfo:\n  name: Test Template"),
+			tmpl: &mockSignableTemplate{},
+			modifyAfterSign: func(data []byte) []byte {
+				signatureIndex := bytes.LastIndex(data, []byte(SignaturePattern))
+				if signatureIndex == -1 {
+					return data
+				}
+				return append(data[:signatureIndex], append([]byte("# Tampered content\n"), data[signatureIndex:]...)...)
+			},
+			wantVerified: false,
+		},
+		{
+			name: "Invalid signature",
+			data: []byte("id: test-template\ninfo:\n  name: Test Template"),
+			tmpl: &mockSignableTemplate{},
+			modifyAfterSign: func(data []byte) []byte {
+				return append(bytes.TrimSuffix(data, []byte("\n")), []byte("\n# digest: invalid_signature:fragment")...)
+			},
+			wantVerifyErr: true,
+			wantVerified:  false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create import files if needed
+			for _, imp := range tt.tmpl.GetFileImports() {
+				err := os.WriteFile(imp, []byte("imported content"), 0644)
+				require.NoError(t, err, "Failed to create import file")
+			}
+
+			// Sign the template
+			signature, err := signer.Sign(tt.data, tt.tmpl)
+			if tt.wantSignErr {
+				assert.Error(t, err, "Expected an error during signing")
+				return
+			}
+			require.NoError(t, err, "Failed to sign template")
+
+			// Append signature to the template data
+			signedData := append(tt.data, []byte("\n"+signature)...)
+
+			// Apply any modifications after signing if specified
+			if tt.modifyAfterSign != nil {
+				signedData = tt.modifyAfterSign(signedData)
+			}
+
+			// Verify the signature
+			verified, err := signer.Verify(signedData, tt.tmpl)
+			if tt.wantVerifyErr {
+				assert.Error(t, err, "Expected an error during verification")
+			} else {
+				assert.NoError(t, err, "Unexpected error during verification")
+			}
+			assert.Equal(t, tt.wantVerified, verified, "Unexpected verification result")
+		})
+	}
+}

@@ -75,11 +75,12 @@ func SignTemplate(templateSigner *signer.TemplateSigner, templatePath string) er
 		return ErrNotATemplate
 	}
 	if !template.Verified {
+		_, content := signer.ExtractSignatureAndContent(bin)
 		signatureData, err := templateSigner.Sign(bin, template)
 		if err != nil {
 			return err
 		}
-		buff := bytes.NewBuffer(signer.RemoveSignatureFromData(bin))
+		buff := bytes.NewBuffer(content)
 		buff.WriteString("\n" + signatureData)
 		return os.WriteFile(templatePath, buff.Bytes(), 0644)
 	}