Skip to content

Commit

Permalink
Add --additional-packs argument
Browse files Browse the repository at this point in the history
  • Loading branch information
pwntester committed Oct 26, 2023
1 parent 2c13e0f commit cf4eb8a
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 35 deletions.
29 changes: 15 additions & 14 deletions cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,20 +31,21 @@ import (
)

var (
sessionNameFlag string
runIdFlag int
sessionPrefixFlag string
outputDirFlag string
downloadDBsFlag bool
nwoFlag string
jsonFlag bool
languageFlag string
listFileFlag string
listFlag string
codeqlPathFlag string
controllerFlag string
queryFileFlag string
querySuiteFileFlag string
sessionNameFlag string
runIdFlag int
sessionPrefixFlag string
outputDirFlag string
downloadDBsFlag bool
nwoFlag string
jsonFlag bool
languageFlag string
listFileFlag string
listFlag string
codeqlPathFlag string
controllerFlag string
queryFileFlag string
querySuiteFileFlag string
additionalPacksFlag string
)
var rootCmd = &cobra.Command{
Use: "gh-mrva",
Expand Down
33 changes: 23 additions & 10 deletions cmd/submit.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,15 @@ import (
)

var (
controller string
codeqlPath string
listFile string
listName string
language string
sessionName string
queryFile string
querySuiteFile string
additionalPacks string
controller string
codeqlPath string
listFile string
listName string
language string
sessionName string
queryFile string
querySuiteFile string
)
var submitCmd = &cobra.Command{
Use: "submit",
Expand All @@ -43,6 +44,7 @@ func init() {
submitCmd.Flags().StringVarP(&listFileFlag, "list-file", "f", "", "Path to repo list file (overrides config file)")
submitCmd.Flags().StringVarP(&listFlag, "list", "i", "", "Name of repo list")
submitCmd.Flags().StringVarP(&codeqlPathFlag, "codeql-path", "p", "", "Path to CodeQL distribution (overrides config file)")
submitCmd.Flags().StringVarP(&additionalPacksFlag, "additional-packs", "a", "", "Additional Packs")
submitCmd.MarkFlagRequired("session")
submitCmd.MarkFlagRequired("language")
submitCmd.MarkFlagsMutuallyExclusive("query", "query-suite")
Expand All @@ -69,6 +71,9 @@ func submitQuery() {
} else if configData.CodeQLPath != "" {
codeqlPath = configData.CodeQLPath
}
if additionalPacksFlag != "" {
additionalPacks = additionalPacksFlag
}
if languageFlag != "" {
language = languageFlag
}
Expand All @@ -85,6 +90,14 @@ func submitQuery() {
querySuiteFile = querySuiteFileFlag
}

if codeqlPath != "" {
if additionalPacks != "" {
additionalPacks = ":" + codeqlPath
} else {
additionalPacks = codeqlPath
}
}

if controller == "" {
fmt.Println("Please specify a controller.")
os.Exit(1)
Expand Down Expand Up @@ -118,13 +131,13 @@ func submitQuery() {
if queryFileFlag != "" {
queries = append(queries, queryFileFlag)
} else if querySuiteFileFlag != "" {
queries = utils.ResolveQueries(codeqlPath, querySuiteFile)
queries = utils.ResolveQueries(additionalPacks, querySuiteFile)
}

fmt.Printf("Submitting %d queries for %d repositories\n", len(queries), len(repositories))
var runs []models.Run
for _, query := range queries {
encodedBundle, queryId, err := utils.GenerateQueryPack(codeqlPath, query, language)
encodedBundle, queryId, err := utils.GenerateQueryPack(query, language, additionalPacks)
if err != nil {
log.Fatal(err)
}
Expand Down
25 changes: 14 additions & 11 deletions utils/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -259,9 +259,9 @@ func ResolveQueryId(queryFile string) (string, error) {
}
}

func ResolveQueries(codeqlPath string, querySuite string) []string {
func ResolveQueries(additionalPacks string, querySuite string) []string {
args := []string{"resolve", "queries", "--format=json", querySuite}
jsonBytes, err := RunCodeQLCommand(codeqlPath, false, args...)
jsonBytes, err := RunCodeQLCommand(additionalPacks, false, args...)
var queries []string
if strings.TrimSpace(string(jsonBytes)) == "" {
fmt.Println("No queries found in the specified query suite.")
Expand All @@ -275,9 +275,12 @@ func ResolveQueries(codeqlPath string, querySuite string) []string {
return queries
}

func RunCodeQLCommand(codeqlPath string, combined bool, args ...string) ([]byte, error) {
if codeqlPath != "" && !strings.Contains(strings.Join(args, " "), "packlist") {
args = append(args, fmt.Sprintf("--additional-packs=%s", codeqlPath))
func RunCodeQLCommand(additionalPacks string, combined bool, args ...string) ([]byte, error) {
if additionalPacks != "" {
args = append(args, "--additional-packs", additionalPacks)
}
if strings.Contains(strings.Join(args, " "), "pack install") {
args = append(args, "--no-strict-mode")
}
cmd := exec.Command("codeql", args...)
cmd.Env = os.Environ()
Expand All @@ -288,7 +291,7 @@ func RunCodeQLCommand(codeqlPath string, combined bool, args ...string) ([]byte,
}
}

func GenerateQueryPack(codeqlPath string, queryFile string, language string) (string, string, error) {
func GenerateQueryPack(queryFile string, language string, additionalPacks string) (string, string, error) {
fmt.Printf("Generating query pack for %s\n", queryFile)

// create a temporary directory to hold the query pack
Expand Down Expand Up @@ -355,7 +358,7 @@ defaultSuite:
} else {
// don't include all query files in the QLPacks. We only want the queryFile to be copied.
fmt.Printf("QLPack exists, stripping all other queries from %s\n", originalPackRoot)
toCopy := PackPacklist(codeqlPath, originalPackRoot, false)
toCopy := PackPacklist(originalPackRoot, false)
// also copy the lock file (either new name or old name) and the query file itself (these are not included in the packlist)
lockFileNew := filepath.Join(originalPackRoot, "qlpack.lock.yml")
lockFileOld := filepath.Join(originalPackRoot, "codeql-pack.lock.yml")
Expand Down Expand Up @@ -389,7 +392,7 @@ defaultSuite:
// install the pack dependencies
fmt.Print("Installing QLPack dependencies\n")
args := []string{"pack", "install", queryPackDir}
stdouterr, err := RunCodeQLCommand(codeqlPath, true, args...)
stdouterr, err := RunCodeQLCommand(additionalPacks, true, args...)
if err != nil {
fmt.Printf("`codeql pack bundle` failed with error: %v\n", string(stdouterr))
return "", "", fmt.Errorf("Failed to install query pack: %v", err)
Expand All @@ -398,7 +401,7 @@ defaultSuite:
fmt.Print("Compiling and bundling the QLPack (This may take a while)\n")
args = []string{"pack", "bundle", "-o", bundlePath, queryPackDir}
args = append(args, precompilationOpts...)
stdouterr, err = RunCodeQLCommand(codeqlPath, true, args...)
stdouterr, err = RunCodeQLCommand(additionalPacks, true, args...)
if err != nil {
fmt.Printf("`codeql pack bundle` failed with error: %v\n", string(stdouterr))
return "", "", fmt.Errorf("Failed to bundle query pack: %v\n", err)
Expand All @@ -419,14 +422,14 @@ defaultSuite:
return bundleBase64, queryId, nil
}

func PackPacklist(codeqlPath string, dir string, includeQueries bool) []string {
func PackPacklist(dir string, includeQueries bool) []string {
// since 2.7.1, packlist returns an object with a "paths" property that is a list of packs.
args := []string{"pack", "packlist", "--format=json"}
if !includeQueries {
args = append(args, "--no-include-queries")
}
args = append(args, dir)
jsonBytes, err := RunCodeQLCommand(codeqlPath, false, args...)
jsonBytes, err := RunCodeQLCommand("", false, args...)
var packlist map[string][]string
err = json.Unmarshal(jsonBytes, &packlist)
if err != nil {
Expand Down

0 comments on commit cf4eb8a

Please sign in to comment.