Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
projectdiscovery
GitHub Repository: projectdiscovery/nuclei
Path: blob/dev/pkg/templates/signer/tmpl_signer_test.go
2070 views
1
package signer
2
3
import (
4
"bytes"
5
"os"
6
"path/filepath"
7
"testing"
8
9
"github.com/stretchr/testify/assert"
10
"github.com/stretchr/testify/require"
11
)
12
13
const (
14
testCertFile = "../../../integration_tests/protocols/keys/ci.crt"
15
testKeyFile = "../../../integration_tests/protocols/keys/ci-private-key.pem"
16
)
17
18
type mockSignableTemplate struct {
19
imports []string
20
hasCode bool
21
}
22
23
func (m *mockSignableTemplate) GetFileImports() []string {
24
return m.imports
25
}
26
27
func (m *mockSignableTemplate) HasCodeProtocol() bool {
28
return m.hasCode
29
}
30
31
var signer, _ = NewTemplateSignerFromFiles(testCertFile, testKeyFile)
32
33
func TestTemplateSignerSignAndVerify(t *testing.T) {
34
tempDir := t.TempDir()
35
36
tests := []struct {
37
name string
38
data []byte
39
tmpl SignableTemplate
40
wantSignErr bool
41
wantVerifyErr bool
42
wantVerified bool
43
modifyAfterSign func([]byte) []byte
44
}{
45
{
46
name: "Simple template",
47
data: []byte("id: test-template\ninfo:\n name: Test Template"),
48
tmpl: &mockSignableTemplate{},
49
wantVerified: true,
50
},
51
{
52
name: "Template with imports",
53
data: []byte("id: test-template\ninfo:\n name: Test Template"),
54
tmpl: &mockSignableTemplate{imports: []string{
55
filepath.Join(tempDir, "import1.yaml"),
56
filepath.Join(tempDir, "import2.yaml"),
57
}},
58
wantVerified: true,
59
},
60
{
61
name: "Template with code protocol",
62
data: []byte("id: test-template\ninfo:\n name: Test Template\n\ncode:\n - engine: bash\n source: echo 'Hello, World!'"),
63
tmpl: &mockSignableTemplate{hasCode: true},
64
wantSignErr: false,
65
wantVerified: true,
66
},
67
{
68
name: "Tampered template",
69
data: []byte("id: test-template\ninfo:\n name: Test Template"),
70
tmpl: &mockSignableTemplate{},
71
modifyAfterSign: func(data []byte) []byte {
72
signatureIndex := bytes.LastIndex(data, []byte(SignaturePattern))
73
if signatureIndex == -1 {
74
return data
75
}
76
return append(data[:signatureIndex], append([]byte("# Tampered content\n"), data[signatureIndex:]...)...)
77
},
78
wantVerified: false,
79
},
80
{
81
name: "Invalid signature",
82
data: []byte("id: test-template\ninfo:\n name: Test Template"),
83
tmpl: &mockSignableTemplate{},
84
modifyAfterSign: func(data []byte) []byte {
85
return append(bytes.TrimSuffix(data, []byte("\n")), []byte("\n# digest: invalid_signature:fragment")...)
86
},
87
wantVerifyErr: true,
88
wantVerified: false,
89
},
90
}
91
92
for _, tt := range tests {
93
t.Run(tt.name, func(t *testing.T) {
94
// Create import files if needed
95
for _, imp := range tt.tmpl.GetFileImports() {
96
err := os.WriteFile(imp, []byte("imported content"), 0644)
97
require.NoError(t, err, "Failed to create import file")
98
}
99
100
// Sign the template
101
signature, err := signer.Sign(tt.data, tt.tmpl)
102
if tt.wantSignErr {
103
assert.Error(t, err, "Expected an error during signing")
104
return
105
}
106
require.NoError(t, err, "Failed to sign template")
107
108
// Append signature to the template data
109
signedData := append(tt.data, []byte("\n"+signature)...)
110
111
// Apply any modifications after signing if specified
112
if tt.modifyAfterSign != nil {
113
signedData = tt.modifyAfterSign(signedData)
114
}
115
116
// Verify the signature
117
verified, err := signer.Verify(signedData, tt.tmpl)
118
if tt.wantVerifyErr {
119
assert.Error(t, err, "Expected an error during verification")
120
} else {
121
assert.NoError(t, err, "Unexpected error during verification")
122
}
123
assert.Equal(t, tt.wantVerified, verified, "Unexpected verification result")
124
})
125
}
126
}
127
128