-
Notifications
You must be signed in to change notification settings - Fork 2.3k
/
gitlab.go
148 lines (124 loc) · 5.22 KB
/
gitlab.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
package customtemplates
import (
"context"
"encoding/base64"
"os"
"path/filepath"
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
"github.com/projectdiscovery/nuclei/v3/pkg/types"
errorutil "github.com/projectdiscovery/utils/errors"
"github.com/xanzy/go-gitlab"
)
var _ Provider = &customTemplateGitLabRepo{}
type customTemplateGitLabRepo struct {
gitLabClient *gitlab.Client
serverURL string
projectIDs []int
}
// NewGitLabProviders returns a new list of GitLab providers for downloading custom templates
func NewGitLabProviders(options *types.Options) ([]*customTemplateGitLabRepo, error) {
providers := []*customTemplateGitLabRepo{}
if options.GitLabToken != "" && !options.GitLabTemplateDisableDownload {
// Establish a connection to GitLab and build a client object with which to download templates from GitLab
gitLabClient, err := getGitLabClient(options.GitLabServerURL, options.GitLabToken)
if err != nil {
return nil, errorutil.NewWithErr(err).Msgf("Error establishing GitLab client for %s %s", options.GitLabServerURL, err)
}
// Create a new GitLab service client
gitLabContainer := &customTemplateGitLabRepo{
gitLabClient: gitLabClient,
serverURL: options.GitLabServerURL,
projectIDs: options.GitLabTemplateRepositoryIDs,
}
// Add the GitLab service client to the list of custom templates
providers = append(providers, gitLabContainer)
}
return providers, nil
}
// Download downloads all .yaml files from a GitLab repository
func (bk *customTemplateGitLabRepo) Download(_ context.Context) {
// Define the project and template count
var projectCount = 0
var templateCount = 0
// Append the GitLab directory to the location
location := config.DefaultConfig.CustomGitLabTemplatesDirectory
// Ensure the CustomGitLabTemplateDirectory directory exists or create it if it doesn't yet exist
err := os.MkdirAll(filepath.Dir(location), 0755)
if err != nil {
gologger.Error().Msgf("Error creating directory: %v", err)
return
}
// Get the projects from the GitLab serverURL
for _, projectID := range bk.projectIDs {
// Get the project information from the GitLab serverURL to get the default branch and the project name
project, _, err := bk.gitLabClient.Projects.GetProject(projectID, nil)
if err != nil {
gologger.Error().Msgf("error retrieving GitLab project: %s %s", project, err)
return
}
// Add a subdirectory with the project ID as the subdirectory within the location
projectOutputPath := filepath.Join(location, project.Path)
// Ensure the subdirectory exists or create it if it doesn't yet exist
err = os.MkdirAll(projectOutputPath, 0755)
if err != nil {
gologger.Error().Msgf("Error creating subdirectory: %v", err)
return
}
// Get the directory listing for the files in the project
tree, _, err := bk.gitLabClient.Repositories.ListTree(projectID, &gitlab.ListTreeOptions{
Ref: gitlab.String(project.DefaultBranch),
Recursive: gitlab.Bool(true),
})
if err != nil {
gologger.Error().Msgf("error retrieving files from GitLab project: %s (%d) %s", project.Name, projectID, err)
}
// Loop through the tree and download the files
for _, file := range tree {
// If the object is not a file or file extension is not .yaml, skip it
if file.Type == "blob" && filepath.Ext(file.Path) == ".yaml" {
gf := &gitlab.GetFileOptions{
Ref: gitlab.String(project.DefaultBranch),
}
f, _, err := bk.gitLabClient.RepositoryFiles.GetFile(projectID, file.Path, gf)
if err != nil {
gologger.Error().Msgf("error retrieving GitLab project file: %d %s", projectID, err)
return
}
// Decode the file content from base64 into bytes so that it can be written to the local filesystem
contents, err := base64.StdEncoding.DecodeString(f.Content)
if err != nil {
gologger.Error().Msgf("error decoding GitLab project (%s) file: %s %s", project.Name, f.FileName, err)
return
}
// Write the downloaded template to the local filesystem at the location with the filename of the blob name
err = os.WriteFile(filepath.Join(projectOutputPath, f.FileName), contents, 0644)
if err != nil {
gologger.Error().Msgf("error writing GitLab project (%s) file: %s %s", project.Name, f.FileName, err)
return
}
// Increment the number of templates downloaded
templateCount++
}
}
// Increment the number of projects downloaded
projectCount++
gologger.Info().Msgf("GitLab project '%s' (%d) cloned successfully", project.Name, projectID)
}
// Print the number of projects and templates downloaded
gologger.Info().Msgf("%d templates downloaded from %d GitLab project(s) to: %s", templateCount, projectCount, location)
}
// Update is a wrapper around Download since it doesn't maintain a diff of the templates downloaded versus in the
// repository for simplicity.
func (bk *customTemplateGitLabRepo) Update(ctx context.Context) {
if len(bk.projectIDs) == 0 {
// No projects to download or update
return
}
bk.Download(ctx)
}
// getGitLabClient returns a GitLab client for the given serverURL and token
func getGitLabClient(server string, token string) (*gitlab.Client, error) {
client, err := gitlab.NewClient(token, gitlab.WithBaseURL(server))
return client, err
}