mirror of
https://github.com/projectdiscovery/nuclei.git
synced 2025-12-21 20:15:24 +00:00
d5a09e733a
* Add s3 bucket template provider - Refactor the custom github template code - add interface for template provider * Validate if aws creds are passed if bucket flag - refactor s3 provider struct to take client - add function which returns the aws s3 client - update error messages * Add aws s3 bucket flags documentation in README.md - Rename the github_test.go to customTemplate_test.go * go mod update * Move template provider code to pkg/external/customtemplates dir * Remove github and aws update variables from flag * Rename CustomTemplateProvider to Provider * Update integration and function command in makefile * Update github test case, accept token * readme update * go mod tidy * Update build-test.yml * handle empty dir in s3 * Add requested changes - download/update s3 and github only when `-ut` is passed - only print the missing env variable for s3 - add the custom templates path in ~/.config/nuclei/.template-config.json * print custom paths only if exists in config file * misc update * tag update Co-authored-by: sandeep <8293321+ehsandeep@users.noreply.github.com> Co-authored-by: Sandeep Singh <sandeep@projectdiscovery.io>
86 lines
2.5 KiB
Go
86 lines
2.5 KiB
Go
package customtemplates
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"path/filepath"
|
|
|
|
"github.com/aws/aws-sdk-go-v2/aws"
|
|
"github.com/aws/aws-sdk-go-v2/config"
|
|
"github.com/aws/aws-sdk-go-v2/credentials"
|
|
"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
|
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
|
"github.com/projectdiscovery/gologger"
|
|
"github.com/projectdiscovery/stringsutil"
|
|
)
|
|
|
|
type customTemplateS3Bucket struct {
|
|
s3Client *s3.Client
|
|
bucketName string
|
|
prefix string
|
|
Location string
|
|
}
|
|
|
|
// download custom templates from s3 bucket
|
|
func (bk *customTemplateS3Bucket) Download(location string, ctx context.Context) {
|
|
downloadPath := filepath.Join(location, CustomS3TemplateDirectory, bk.bucketName)
|
|
|
|
manager := manager.NewDownloader(bk.s3Client)
|
|
paginator := s3.NewListObjectsV2Paginator(bk.s3Client, &s3.ListObjectsV2Input{
|
|
Bucket: &bk.bucketName,
|
|
Prefix: &bk.prefix,
|
|
})
|
|
|
|
for paginator.HasMorePages() {
|
|
page, err := paginator.NextPage(context.TODO())
|
|
if err != nil {
|
|
gologger.Error().Msgf("error downloading s3 bucket %s %s", bk.bucketName, err)
|
|
return
|
|
}
|
|
for _, obj := range page.Contents {
|
|
if err := downloadToFile(manager, downloadPath, bk.bucketName, aws.ToString(obj.Key)); err != nil {
|
|
gologger.Error().Msgf("error downloading s3 bucket %s %s", bk.bucketName, err)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
gologger.Info().Msgf("AWS bucket %s successfully cloned successfully at %s", bk.bucketName, downloadPath)
|
|
}
|
|
|
|
// download custom templates from s3 bucket
|
|
func (bk *customTemplateS3Bucket) Update(location string, ctx context.Context) {
|
|
bk.Download(location, ctx)
|
|
}
|
|
|
|
func downloadToFile(downloader *manager.Downloader, targetDirectory, bucket, key string) error {
|
|
// Create the directories in the path
|
|
file := filepath.Join(targetDirectory, key)
|
|
// If empty dir in s3
|
|
if stringsutil.HasSuffixI(key, "/") {
|
|
return os.MkdirAll(file, 0775)
|
|
}
|
|
if err := os.MkdirAll(filepath.Dir(file), 0775); err != nil {
|
|
return err
|
|
}
|
|
|
|
// Set up the local file
|
|
fd, err := os.Create(file)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer fd.Close()
|
|
|
|
// Download the file using the AWS SDK for Go
|
|
_, err = downloader.Download(context.TODO(), fd, &s3.GetObjectInput{Bucket: &bucket, Key: &key})
|
|
|
|
return err
|
|
}
|
|
|
|
func getS3Client(ctx context.Context, acccessKey, secretKey, region string) (*s3.Client, error) {
|
|
cfg, err := config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(acccessKey, secretKey, "")), config.WithRegion(region))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return s3.NewFromConfig(cfg), nil
|
|
}
|