Skip to content

Commit

Permalink
Gitlab Custom Templates (#3570)
Browse files Browse the repository at this point in the history 
* Configuration options for GitLab template pulls

* GitLab client creation

* GitLab hooks and property renames

* Fix filesystem writing and update environment variables

* Fix type error in formatted error message

* Migrate directory config to new nucleiconfig file

* refactor + add custom templates to tm

* typo fix + only show installed ct with -tv

* add default gitlab url if not given

* fix template valid failure

---------

Co-authored-by: Tarun Koyalwar <[email protected]>
  • Loading branch information
@kchason @tarunKoyalwar
kchason and tarunKoyalwar authored Apr 19, 2023
1 parent b211d6f commit dcb0032
Show file tree
Hide file tree
Showing 13 changed files with 442 additions and 110 deletions.
11 changes: 9 additions & 2 deletions v2/cmd/nuclei/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -417,12 +417,19 @@ func printVersion() {
func printTemplateVersion() {
cfg := config.DefaultConfig
gologger.Info().Msgf("Public nuclei-templates version: %s (%s)\n", cfg.TemplateVersion, cfg.TemplatesDirectory)
if cfg.CustomS3TemplatesDirectory != "" {

if fileutil.FolderExists(cfg.CustomS3TemplatesDirectory) {
gologger.Info().Msgf("Custom S3 templates location: %s\n", cfg.CustomS3TemplatesDirectory)
}
if cfg.CustomGithubTemplatesDirectory != "" {
if fileutil.FolderExists(cfg.CustomGithubTemplatesDirectory) {
gologger.Info().Msgf("Custom Github templates location: %s ", cfg.CustomGithubTemplatesDirectory)
}
if fileutil.FolderExists(cfg.CustomGitLabTemplatesDirectory) {
gologger.Info().Msgf("Custom Gitlab templates location: %s ", cfg.CustomGitLabTemplatesDirectory)
}
if fileutil.FolderExists(cfg.CustomAzureTemplatesDirectory) {
gologger.Info().Msgf("Custom Azure templates location: %s ", cfg.CustomAzureTemplatesDirectory)
}
os.Exit(0)
}

Expand Down
16 changes: 13 additions & 3 deletions v2/internal/installer/template.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package installer

import (
"bytes"
"context"
"crypto/md5"
"fmt"
"io"
Expand All @@ -14,6 +15,7 @@ import (
"github.com/olekukonko/tablewriter"
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/nuclei/v2/pkg/catalog/config"
"github.com/projectdiscovery/nuclei/v2/pkg/external/customtemplates"
errorutil "github.com/projectdiscovery/utils/errors"
fileutil "github.com/projectdiscovery/utils/file"
stringsutil "github.com/projectdiscovery/utils/strings"
Expand Down Expand Up @@ -54,7 +56,9 @@ func (t *templateUpdateResults) String() string {

// TemplateManager is a manager for templates.
// It downloads / updates / installs templates.
type TemplateManager struct{}
type TemplateManager struct {
CustomTemplates *customtemplates.CustomTemplatesManager // optional if given tries to download custom templates
}

// FreshInstallIfNotExists installs templates if they are not already installed
// if templates directory already exists, it does nothing
Expand All @@ -63,7 +67,13 @@ func (t *TemplateManager) FreshInstallIfNotExists() error {
return nil
}
gologger.Info().Msgf("nuclei-templates are not installed, installing...")
return t.installTemplatesAt(config.DefaultConfig.TemplatesDirectory)
if err := t.installTemplatesAt(config.DefaultConfig.TemplatesDirectory); err != nil {
return errorutil.NewWithErr(err).Msgf("failed to install templates at %s", config.DefaultConfig.TemplatesDirectory)
}
if t.CustomTemplates != nil {
t.CustomTemplates.Download(context.TODO())
}
return nil
}

// UpdateIfOutdated updates templates if they are outdated
Expand Down Expand Up @@ -310,7 +320,7 @@ func (t *TemplateManager) calculateChecksumMap(dir string) (map[string]string, e
return err
}
// skip checksums of custom templates i.e github and s3
if stringsutil.HasPrefixAny(path, config.DefaultConfig.CustomGithubTemplatesDirectory, config.DefaultConfig.CustomS3TemplatesDirectory) {
if stringsutil.HasPrefixAny(path, config.DefaultConfig.GetAllCustomTemplateDirs()...) {
return nil
}

Expand Down
48 changes: 47 additions & 1 deletion v2/internal/runner/options.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"fmt"
"os"
"path/filepath"
"strconv"
"strings"

"github.com/pkg/errors"
Expand Down Expand Up @@ -129,7 +130,7 @@ func validateOptions(options *types.Options) error {
}
validateCertificatePaths([]string{options.ClientCertFile, options.ClientKeyFile, options.ClientCAFile})
}
// Verify aws secrets are passed if s3 template bucket passed
// Verify AWS secrets are passed if a S3 template bucket is passed
if options.AwsBucketName != "" && options.UpdateTemplates {
missing := validateMissingS3Options(options)
if missing != nil {
Expand All @@ -145,6 +146,14 @@ func validateOptions(options *types.Options) error {
}
}

// Verify that all GitLab options are provided if the GitLab server or token is provided
if options.GitLabToken != "" && options.UpdateTemplates {
missing := validateMissingGitLabOptions(options)
if missing != nil {
return fmt.Errorf("gitlab server details are missing. Please provide %s", strings.Join(missing, ","))
}
}

// verify that a valid ip version type was selected (4, 6)
if len(options.IPVersion) == 0 {
// add ipv4 as default
Expand Down Expand Up @@ -186,6 +195,8 @@ func validateCloudOptions(options *types.Options) error {
missing = validateMissingS3Options(options)
case "github":
missing = validateMissingGithubOptions(options)
case "gitlab":
missing = validateMissingGitLabOptions(options)
case "azure":
missing = validateMissingAzureOptions(options)
}
Expand Down Expand Up @@ -244,6 +255,18 @@ func validateMissingGithubOptions(options *types.Options) []string {
return missing
}

func validateMissingGitLabOptions(options *types.Options) []string {
var missing []string
if options.GitLabToken == "" {
missing = append(missing, "GITLAB_TOKEN")
}
if len(options.GitLabTemplateRepositoryIDs) == 0 {
missing = append(missing, "GITLAB_REPOSITORY_IDS")
}

return missing
}

// configureOutput configures the output logging levels to be displayed on the screen
func configureOutput(options *types.Options) {
// If the user desires verbose output, show verbose output
Expand Down Expand Up @@ -333,6 +356,29 @@ func readEnvInputVars(options *types.Options) {
if repolist != "" {
options.GithubTemplateRepo = append(options.GithubTemplateRepo, stringsutil.SplitAny(repolist, ",")...)
}

// GitLab options for downloading templates from a repository
options.GitLabServerURL = os.Getenv("GITLAB_SERVER_URL")
if options.GitLabServerURL == "" {
options.GitLabServerURL = "https://gitlab.com"
}
options.GitLabToken = os.Getenv("GITLAB_TOKEN")
repolist = os.Getenv("GITLAB_REPOSITORY_IDS")
// Convert the comma separated list of repository IDs to a list of integers
if repolist != "" {
for _, repoID := range stringsutil.SplitAny(repolist, ",") {
// Attempt to convert the repo ID to an integer
repoIDInt, err := strconv.Atoi(repoID)
if err != nil {
gologger.Warning().Msgf("Invalid GitLab template repository ID: %s", repoID)
continue
}

// Add the int repository ID to the list
options.GitLabTemplateRepositoryIDs = append(options.GitLabTemplateRepositoryIDs, repoIDInt)
}
}

// AWS options for downloading templates from an S3 bucket
options.AwsAccessKey = os.Getenv("AWS_ACCESS_KEY")
options.AwsSecretKey = os.Getenv("AWS_SECRET_KEY")
Expand Down
25 changes: 21 additions & 4 deletions v2/internal/runner/runner.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,6 @@ type Runner struct {
hostErrors hosterrorscache.CacheInterface
resumeCfg *types.ResumeCfg
pprofServer *http.Server
customTemplates []customtemplates.Provider
cloudClient *nucleicloud.Client
cloudTargets []string
}
Expand Down Expand Up @@ -102,8 +101,16 @@ func New(options *types.Options) (*Runner, error) {
gologger.Error().Msgf("nuclei version check failed got: %s\n", err)
}
}

// check for custom template updates and update if available
ctm, err := customtemplates.NewCustomTemplatesManager(options)
if err != nil {
gologger.Error().Label("custom-templates").Msgf("Failed to create custom templates manager: %s\n", err)
}

// Check for template updates and update if available
tm := &installer.TemplateManager{}
// if custom templates manager is not nil, we will install custom templates if there is fresh installation
tm := &installer.TemplateManager{CustomTemplates: ctm}
if err := tm.FreshInstallIfNotExists(); err != nil {
gologger.Warning().Msgf("failed to install nuclei templates: %s\n", err)
}
Expand All @@ -116,6 +123,18 @@ func New(options *types.Options) (*Runner, error) {
gologger.Warning().Msgf("failed to update nuclei ignore file: %s\n", err)
}
}

if options.UpdateTemplates {
// we automatically check for updates unless explicitly disabled
// this print statement is only to inform the user that there are no updates
if !config.DefaultConfig.NeedsTemplateUpdate() {
gologger.Info().Msgf("No new updates found for nuclei templates")
}
// manually trigger update of custom templates
if ctm != nil {
ctm.Update(context.TODO())
}
}
}

if options.Validate {
Expand All @@ -125,8 +144,6 @@ func New(options *types.Options) (*Runner, error) {
// TODO: refactor to pass options reference globally without cycles
parsers.NoStrictSyntax = options.NoStrictSyntax
yaml.StrictSyntax = !options.NoStrictSyntax
// parse the runner.options.GithubTemplateRepo and store the valid repos in runner.customTemplateRepos
runner.customTemplates = customtemplates.ParseCustomTemplates(runner.options)

if options.Headless {
if engine.MustDisableSandbox() {
Expand Down
8 changes: 6 additions & 2 deletions v2/pkg/catalog/config/constants.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,6 @@ import (
const (
TemplateConfigFileName = ".templates-config.json"
NucleiTemplatesDirName = "nuclei-templates"
CustomS3TemplatesDirName = "s3"
CustomGithubTemplatesDirName = "github"
OfficialNucleiTeamplatesRepoName = "nuclei-templates"
NucleiIgnoreFileName = ".nuclei-ignore"
NucleiTemplatesCheckSumFileName = ".checksum"
Expand All @@ -19,6 +17,12 @@ const (
ReportingConfigFilename = "reporting-config.yaml"
// Version is the current version of nuclei
Version = `v2.9.2-dev`

// Directory Names of custom templates
CustomS3TemplatesDirName = "s3"
CustomGithubTemplatesDirName = "github"
CustomAzureTemplatesDirName = "azure"
CustomGitLabTemplatesDirName = "gitlab"
)

// IsOutdatedVersion compares two versions and returns true
Expand Down
24 changes: 19 additions & 5 deletions v2/pkg/catalog/config/nucleiconfig.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,12 @@ var DefaultConfig *Config
type Config struct {
TemplatesDirectory string `json:"nuclei-templates-directory,omitempty"`

// customtemplates exists in templates directory with the name of custom-templates provider
// below custom paths are absolute paths to respecitive custom-templates directories
CustomS3TemplatesDirectory string `json:"custom-s3-templates-directory"`
CustomGithubTemplatesDirectory string `json:"custom-github-templates-directory"`
CustomGitLabTemplatesDirectory string `json:"custom-gitlab-templates-directory"`
CustomAzureTemplatesDirectory string `json:"custom-azure-templates-directory"`

TemplateVersion string `json:"nuclei-templates-version,omitempty"`
NucleiIgnoreHash string `json:"nuclei-ignore-hash,omitempty"`
Expand Down Expand Up @@ -104,6 +108,11 @@ func (c *Config) GetConfigDir() string {
return c.configDir
}

// GetAllCustomTemplateDirs returns all custom template directories
func (c *Config) GetAllCustomTemplateDirs() []string {
return []string{c.CustomS3TemplatesDirectory, c.CustomGithubTemplatesDirectory, c.CustomGitLabTemplatesDirectory, c.CustomAzureTemplatesDirectory}
}

// GetReportingConfigFilePath returns the nuclei reporting config file path
func (c *Config) GetReportingConfigFilePath() string {
return filepath.Join(c.configDir, ReportingConfigFilename)
Expand Down Expand Up @@ -175,7 +184,9 @@ func (c *Config) SetTemplatesDir(dirPath string) {
c.TemplatesDirectory = dirPath
// Update the custom templates directory
c.CustomGithubTemplatesDirectory = filepath.Join(dirPath, CustomGithubTemplatesDirName)
c.CustomS3TemplatesDirectory = filepath.Join(dirPath, CustomGithubTemplatesDirName)
c.CustomS3TemplatesDirectory = filepath.Join(dirPath, CustomS3TemplatesDirName)
c.CustomGitLabTemplatesDirectory = filepath.Join(dirPath, CustomGitLabTemplatesDirName)
c.CustomAzureTemplatesDirectory = filepath.Join(dirPath, CustomAzureTemplatesDirName)
}

// SetTemplatesVersion sets the new nuclei templates version
Expand All @@ -202,8 +213,6 @@ func (c *Config) ReadTemplatesConfig() error {
return errorutil.NewWithErr(err).Msgf("could not unmarshal nuclei config file at %s", c.getTemplatesConfigFilePath())
}
// apply config
c.CustomGithubTemplatesDirectory = cfg.CustomGithubTemplatesDirectory
c.CustomS3TemplatesDirectory = cfg.CustomS3TemplatesDirectory
c.TemplatesDirectory = cfg.TemplatesDirectory
c.TemplateVersion = cfg.TemplateVersion
c.NucleiIgnoreHash = cfg.NucleiIgnoreHash
Expand Down Expand Up @@ -279,6 +288,11 @@ func init() {
gologger.Error().Msgf("failed to write config file at %s got: %s", DefaultConfig.getTemplatesConfigFilePath(), err)
}
}
// Loads/updates paths of custom templates
// Note: custom templates paths should not be updated in config file
// and even if it is changed we don't follow it since it is not expected behavior
// If custom templates are in default locations only then they are loaded while running nuclei
DefaultConfig.SetTemplatesDir(DefaultConfig.TemplatesDirectory)
}

func getDefaultConfigDir() string {
Expand All @@ -297,6 +311,6 @@ func getDefaultConfigDir() string {
// Add Default Config adds default when .templates-config.json file is not present
func applyDefaultConfig() {
DefaultConfig.TemplatesDirectory = filepath.Join(DefaultConfig.homeDir, NucleiTemplatesDirName)
DefaultConfig.CustomGithubTemplatesDirectory = filepath.Join(DefaultConfig.TemplatesDirectory, CustomGithubTemplatesDirName)
DefaultConfig.CustomS3TemplatesDirectory = filepath.Join(DefaultConfig.TemplatesDirectory, CustomS3TemplatesDirName)
// updates all necessary paths
DefaultConfig.SetTemplatesDir(DefaultConfig.TemplatesDirectory)
}
42 changes: 35 additions & 7 deletions v2/pkg/external/customtemplates/azure_blob.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,19 +3,47 @@ package customtemplates
import (
"bytes"
"context"
"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
"github.com/projectdiscovery/gologger"
"os"
"path/filepath"
"strings"

"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/nuclei/v2/pkg/catalog/config"
"github.com/projectdiscovery/nuclei/v2/pkg/types"
errorutil "github.com/projectdiscovery/utils/errors"
)

var _ Provider = &customTemplateAzureBlob{}

type customTemplateAzureBlob struct {
azureBlobClient *azblob.Client
containerName string
}

// NewAzureProviders creates a new Azure Blob Storage provider for downloading custom templates
func NewAzureProviders(options *types.Options) ([]*customTemplateAzureBlob, error) {
providers := []*customTemplateAzureBlob{}
if options.AzureContainerName != "" {
// Establish a connection to Azure and build a client object with which to download templates from Azure Blob Storage
azClient, err := getAzureBlobClient(options.AzureTenantID, options.AzureClientID, options.AzureClientSecret, options.AzureServiceURL)
if err != nil {
return nil, errorutil.NewWithErr(err).Msgf("Error establishing Azure Blob client for %s", options.AzureContainerName)
}

// Create a new Azure Blob Storage container object
azTemplateContainer := &customTemplateAzureBlob{
azureBlobClient: azClient,
containerName: options.AzureContainerName,
}

// Add the Azure Blob Storage container object to the list of custom templates
providers = append(providers, azTemplateContainer)
}
return providers, nil
}

func getAzureBlobClient(tenantID string, clientID string, clientSecret string, serviceURL string) (*azblob.Client, error) {
// Create an Azure credential using the provided credentials
credentials, err := azidentity.NewClientSecretCredential(tenantID, clientID, clientSecret, nil)
Expand All @@ -34,12 +62,12 @@ func getAzureBlobClient(tenantID string, clientID string, clientSecret string, s
return client, nil
}

func (bk *customTemplateAzureBlob) Download(location string, ctx context.Context) {
func (bk *customTemplateAzureBlob) Download(ctx context.Context) {
// Set an incrementer for the number of templates downloaded
var templatesDownloaded = 0

// Define the local path to which the templates will be downloaded
downloadPath := filepath.Join(location, CustomAzureTemplateDirectory, bk.containerName)
downloadPath := filepath.Join(config.DefaultConfig.CustomAzureTemplatesDirectory, bk.containerName)

// Get the list of all templates from the container
pager := bk.azureBlobClient.NewListBlobsFlatPager(bk.containerName, &azblob.ListBlobsFlatOptions{
Expand Down Expand Up @@ -78,9 +106,9 @@ func (bk *customTemplateAzureBlob) Download(location string, ctx context.Context
// Update updates the templates from the Azure Blob Storage container to the local filesystem. This is effectively a
// wrapper of the Download function which downloads of all templates from the container and doesn't manage a
// differential update.
func (bk *customTemplateAzureBlob) Update(location string, ctx context.Context) {
func (bk *customTemplateAzureBlob) Update(ctx context.Context) {
// Treat the update as a download of all templates from the container
bk.Download(location, ctx)
bk.Download(ctx)
}

// downloadTemplate downloads a template from the Azure Blob Storage container to the local filesystem with the provided
Expand Down
Loading

0 comments on commit dcb0032

Please sign in to comment.