Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
projectdiscovery
GitHub Repository: projectdiscovery/nuclei
Path: blob/dev/pkg/catalog/aws/catalog.go
2070 views
1
package aws
2
3
import (
4
"bytes"
5
"context"
6
"errors"
7
"fmt"
8
"io"
9
"path"
10
"slices"
11
"strings"
12
13
"github.com/aws/aws-sdk-go-v2/aws"
14
"github.com/aws/aws-sdk-go-v2/config"
15
"github.com/aws/aws-sdk-go-v2/credentials"
16
"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
17
"github.com/aws/aws-sdk-go-v2/service/s3"
18
)
19
20
// Catalog manages the AWS S3 template catalog
21
type Catalog struct {
22
svc client
23
}
24
25
// client interface abstracts S3 connections
26
type client interface {
27
getAllKeys() ([]string, error)
28
downloadKey(name string) (io.ReadCloser, error)
29
setBucket(bucket string)
30
}
31
32
type s3svc struct {
33
client *s3.Client
34
bucket string
35
}
36
37
// NewCatalog creates a new AWS Catalog object given a required S3 bucket name and optional configurations. If
38
// no configurations to set AWS keys are provided then environment variables will be used to obtain AWS credentials.
39
func NewCatalog(bucket string, configurations ...func(*Catalog) error) (Catalog, error) {
40
var c Catalog
41
42
for _, configuration := range configurations {
43
err := configuration(&c)
44
if err != nil {
45
return c, err
46
}
47
}
48
49
if c.svc == nil {
50
cfg, err := config.LoadDefaultConfig(context.TODO())
51
if err != nil {
52
return c, err
53
}
54
55
c.svc = &s3svc{
56
client: s3.NewFromConfig(cfg),
57
}
58
}
59
c.svc.setBucket(bucket)
60
61
return c, nil
62
}
63
64
// WithAWSKeys enables explicitly setting the AWS access key, secret key and region
65
func WithAWSKeys(accessKey, secretKey, region string) func(*Catalog) error {
66
return func(c *Catalog) error {
67
cfg, err := config.LoadDefaultConfig(context.TODO(),
68
config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(accessKey, secretKey, "")),
69
config.WithRegion(region))
70
if err != nil {
71
return err
72
}
73
74
c.svc = &s3svc{
75
client: s3.NewFromConfig(cfg),
76
bucket: "",
77
}
78
79
return nil
80
}
81
}
82
83
// OpenFile downloads a file from S3 and returns the contents as an io.ReadCloser
84
func (c Catalog) OpenFile(filename string) (io.ReadCloser, error) {
85
if filename == "" {
86
return nil, errors.New("empty filename")
87
}
88
89
return c.svc.downloadKey(filename)
90
}
91
92
// GetTemplatePath looks for a target string performing a simple substring check
93
// against all S3 keys. If the input includes a wildcard (*) it is removed.
94
func (c Catalog) GetTemplatePath(target string) ([]string, error) {
95
target = strings.ReplaceAll(target, "*", "")
96
97
keys, err := c.svc.getAllKeys()
98
if err != nil {
99
return nil, err
100
}
101
102
var matches []string
103
for _, key := range keys {
104
if strings.Contains(key, target) {
105
matches = append(matches, key)
106
}
107
}
108
109
return matches, nil
110
}
111
112
// GetTemplatesPath returns all templates from S3
113
func (c Catalog) GetTemplatesPath(definitions []string) ([]string, map[string]error) {
114
keys, err := c.svc.getAllKeys()
115
if err != nil {
116
// necessary to implement the Catalog interface
117
return nil, map[string]error{"aws": err}
118
}
119
120
return keys, nil
121
}
122
123
// ResolvePath gets a full S3 key given the first param. If the second parameter is
124
// provided it tries to find paths relative to the second path.
125
func (c Catalog) ResolvePath(templateName, second string) (string, error) {
126
keys, err := c.svc.getAllKeys()
127
if err != nil {
128
return "", err
129
}
130
131
// if c second path is given, it's c folder and we join the two and check against keys
132
if second != "" {
133
// Note: Do not replace `path` with `filepath` since filepath is aware of Os path separator
134
// and we only see `/` in s3 paths changing it to filepath cause build fail and other errors
135
target := path.Join(path.Dir(second), templateName)
136
for _, key := range keys {
137
if key == target {
138
return key, nil
139
}
140
}
141
}
142
143
// check if templateName is already an absolute path to c key
144
if slices.Contains(keys, templateName) {
145
return templateName, nil
146
}
147
148
return "", fmt.Errorf("no such path found: %s%s for keys: %v", second, templateName, keys)
149
}
150
151
func (s *s3svc) getAllKeys() ([]string, error) {
152
paginator := s3.NewListObjectsV2Paginator(s.client, &s3.ListObjectsV2Input{
153
Bucket: &s.bucket,
154
})
155
156
var keys []string
157
158
for paginator.HasMorePages() {
159
page, err := paginator.NextPage(context.TODO())
160
if err != nil {
161
return nil, err
162
}
163
for _, obj := range page.Contents {
164
key := aws.ToString(obj.Key)
165
keys = append(keys, key)
166
}
167
}
168
169
return keys, nil
170
}
171
172
func (s *s3svc) downloadKey(name string) (io.ReadCloser, error) {
173
downloader := manager.NewDownloader(s.client)
174
buf := manager.NewWriteAtBuffer([]byte{})
175
_, err := downloader.Download(context.TODO(), buf, &s3.GetObjectInput{
176
Bucket: aws.String(s.bucket),
177
Key: aws.String(name),
178
})
179
if err != nil {
180
return nil, err
181
}
182
183
return io.NopCloser(bytes.NewReader(buf.Bytes())), nil
184
}
185
186
func (s *s3svc) setBucket(bucket string) {
187
s.bucket = bucket
188
}
189
190