-
Notifications
You must be signed in to change notification settings - Fork 12
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
create base enricher and make it follow conventions
- Loading branch information
1 parent
a3df9eb
commit 1d3908b
Showing
8 changed files
with
218 additions
and
274 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,108 +1,59 @@ | ||
| package main | ||
|
|
||
| import ( | ||
| "flag" | ||
| "fmt" | ||
| "log" | ||
| "path/filepath" | ||
| "time" | ||
| "log/slog" | ||
|
|
||
| v1 "github.com/ocurity/dracon/api/proto/v1" | ||
|
|
||
| "github.com/spf13/cobra" | ||
| "github.com/spf13/viper" | ||
| "github.com/ocurity/dracon/components/enrichers" | ||
|
|
||
| "github.com/ocurity/dracon/pkg/db" | ||
| "github.com/ocurity/dracon/pkg/enrichment" | ||
| "github.com/ocurity/dracon/pkg/putil" | ||
| ) | ||
|
|
||
| var ( | ||
| connStr string | ||
| readPath string | ||
| writePath string | ||
| connStr string | ||
| ) | ||
|
|
||
| var rootCmd = &cobra.Command{ | ||
| Use: "enricher", | ||
| Short: "enricher", | ||
| Long: "tool to enrich issues against a database", | ||
| RunE: func(cmd *cobra.Command, args []string) error { | ||
| connStr = viper.GetString("db_connection") | ||
| dbUrl, err := db.ParseConnectionStr(connStr) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| conn, err := dbUrl.Connect() | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| readPath = viper.GetString("read_path") | ||
| res, err := putil.LoadTaggedToolResponse(readPath) | ||
| if err != nil { | ||
| log.Fatal(err) | ||
| } | ||
|
|
||
| log.Printf("Loaded %d tagged tool responses\n", len(res)) | ||
| writePath = viper.GetString("write_path") | ||
| for _, r := range res { | ||
| enrichedIssues := []*v1.EnrichedIssue{} | ||
| log.Printf("enriching %d issues", len(r.GetIssues())) | ||
| for _, i := range r.GetIssues() { | ||
| eI, err := enrichment.EnrichIssue(conn, i) | ||
| if err != nil { | ||
| log.Printf("error enriching issue %s, err: %#v\n", i.Uuid, err) | ||
| continue | ||
| } | ||
| enrichedIssues = append(enrichedIssues, eI) | ||
| log.Printf("enriched issue '%s'", eI.GetRawIssue().GetUuid()) | ||
| } | ||
| if len(enrichedIssues) > 0 { | ||
| if err := putil.WriteEnrichedResults(r, enrichedIssues, | ||
| filepath.Join(writePath, fmt.Sprintf("%s.enriched.pb", r.GetToolName())), | ||
| ); err != nil { | ||
| return err | ||
| } | ||
| } | ||
| if len(r.GetIssues()) > 0 { | ||
| scanStartTime := r.GetScanInfo().GetScanStartTime().AsTime() | ||
| if err := putil.WriteResults( | ||
| r.GetToolName(), | ||
| r.GetIssues(), | ||
| filepath.Join(writePath, fmt.Sprintf("%s.raw.pb", r.GetToolName())), | ||
| r.GetScanInfo().GetScanUuid(), | ||
| scanStartTime.Format(time.RFC3339), | ||
| r.GetScanInfo().GetScanTags(), | ||
| ); err != nil { | ||
| log.Fatalf("could not write results: %s", err) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| return nil | ||
| }, | ||
| func main() { | ||
| flag.StringVar(&connStr, "db_connection", enrichers.LookupEnvOrString("ENRICHER_DB_CONNECTION", ""), "the database connection DSN") | ||
| if err := enrichers.ParseFlags(); err != nil { | ||
| log.Fatal(err) | ||
| } | ||
| if err := run(); err != nil { | ||
| log.Fatal(err) | ||
| } | ||
| } | ||
|
|
||
| func init() { | ||
| rootCmd.Flags().StringVar(&connStr, "db_connection", "", "the database connection DSN") | ||
| rootCmd.Flags().StringVar(&readPath, "read_path", "", "the path to read LaunchToolResponses from") | ||
| rootCmd.Flags().StringVar(&writePath, "write_path", "", "the path to write enriched results to") | ||
| if err := viper.BindPFlag("db_connection", rootCmd.Flags().Lookup("db_connection")); err != nil { | ||
| log.Fatalf("could not bind db_connection flag: %s", err) | ||
| func run() error { | ||
| dbURL, err := db.ParseConnectionStr(connStr) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| if err := viper.BindPFlag("read_path", rootCmd.Flags().Lookup("read_path")); err != nil { | ||
| log.Fatalf("could not bind read_path flag: %s", err) | ||
| conn, err := dbURL.Connect() | ||
| if err != nil { | ||
| return err | ||
| } | ||
| if err := viper.BindPFlag("write_path", rootCmd.Flags().Lookup("write_path")); err != nil { | ||
| log.Fatalf("could not bind write_path flag: %s", err) | ||
| res, err := enrichers.LoadData() | ||
| if err != nil { | ||
| return err | ||
| } | ||
| viper.SetEnvPrefix("enricher") | ||
| viper.AutomaticEnv() | ||
| } | ||
|
|
||
| func main() { | ||
| if err := rootCmd.Execute(); err != nil { | ||
| log.Fatal(err) | ||
| log.Printf("Loaded %d tagged tool responses\n", len(res)) | ||
| for _, r := range res { | ||
| enrichedIssues := []*v1.EnrichedIssue{} | ||
| log.Printf("enriching %d issues", len(r.GetIssues())) | ||
| for _, i := range r.GetIssues() { | ||
| eI, err := enrichment.EnrichIssue(conn, i) | ||
| if err != nil { | ||
| slog.Error(fmt.Sprintf("error enriching issue %s, err: %#v\n", i.Uuid, err)) | ||
| continue | ||
| } | ||
| enrichedIssues = append(enrichedIssues, eI) | ||
| log.Printf("enriched issue '%s'", eI.GetRawIssue().GetUuid()) | ||
| } | ||
| return enrichers.WriteData(enrichedIssues, r, "deduplication") | ||
| } | ||
| return nil | ||
| } |
Oops, something went wrong.