11package commands
22
33import (
4+ "errors"
45 "log/slog"
6+ "os"
57 "regexp"
68 "strings"
79 "time"
810
911 "github.com/l3montree-dev/devguard/internal/core"
1012 "github.com/l3montree-dev/devguard/internal/core/vulndb"
13+ "github.com/l3montree-dev/devguard/internal/database"
14+ "github.com/l3montree-dev/devguard/internal/database/models"
1115 "github.com/l3montree-dev/devguard/internal/database/repositories"
1216 "github.com/spf13/cobra"
1317)
@@ -51,19 +55,41 @@ func isValidCVE(cveID string) bool {
5155 return r .MatchString (cveID )
5256}
5357
58+ func migrateDB (db core.DB ) {
59+ // Run database migrations using the existing database connection
60+ disableAutoMigrate := os .Getenv ("DISABLE_AUTOMIGRATE" )
61+ if disableAutoMigrate != "true" {
62+ slog .Info ("running database migrations..." )
63+ if err := database .RunMigrationsWithDB (db ); err != nil {
64+ slog .Error ("failed to run database migrations" , "error" , err )
65+ panic (errors .New ("Failed to run database migrations" ))
66+ }
67+
68+ // Run hash migrations if needed (when algorithm version changes)
69+ if err := models .RunHashMigrationsIfNeeded (db ); err != nil {
70+ slog .Error ("failed to run hash migrations" , "error" , err )
71+ panic (errors .New ("Failed to run hash migrations" ))
72+ }
73+ } else {
74+ slog .Info ("automatic migrations disabled via DISABLE_AUTOMIGRATE=true" )
75+ }
76+ }
77+
5478func newImportCVECommand () * cobra.Command {
5579 importCmd := & cobra.Command {
5680 Use : "import-cve" ,
5781 Short : "Will import the vulnerability database" ,
5882 Args : cobra .ExactArgs (1 ),
5983 Run : func (cmd * cobra.Command , args []string ) {
6084 core .LoadConfig () // nolint
61- database , err := core .DatabaseFactory ()
85+ db , err := core .DatabaseFactory ()
6286 if err != nil {
6387 slog .Error ("could not connect to database" , "err" , err )
6488 return
6589 }
6690
91+ migrateDB (db )
92+
6793 cveID := args [0 ]
6894 cveID = strings .TrimSpace (strings .ToUpper (cveID ))
6995 // check if first argument is valid cve
@@ -72,9 +98,9 @@ func newImportCVECommand() *cobra.Command {
7298 return
7399 }
74100
75- cveRepository := repositories .NewCVERepository (database )
101+ cveRepository := repositories .NewCVERepository (db )
76102 nvdService := vulndb .NewNVDService (cveRepository )
77- osvService := vulndb .NewOSVService (repositories .NewAffectedComponentRepository (database ))
103+ osvService := vulndb .NewOSVService (repositories .NewAffectedComponentRepository (db ))
78104
79105 cve , err := nvdService .ImportCVE (cveID )
80106
@@ -155,27 +181,29 @@ func newSyncCommand() *cobra.Command {
155181
156182 core .LoadConfig () // nolint
157183
158- database , err := core .DatabaseFactory ()
184+ db , err := core .DatabaseFactory ()
159185 if err != nil {
160186 slog .Error ("could not connect to database" , "err" , err )
161187 return
162188 }
163189
190+ migrateDB (db )
191+
164192 databasesToSync , _ := cmd .Flags ().GetStringArray ("databases" )
165193
166- cveRepository := repositories .NewCVERepository (database )
167- cweRepository := repositories .NewCWERepository (database )
168- affectedCmpRepository := repositories .NewAffectedComponentRepository (database )
194+ cveRepository := repositories .NewCVERepository (db )
195+ cweRepository := repositories .NewCWERepository (db )
196+ affectedCmpRepository := repositories .NewAffectedComponentRepository (db )
169197 nvdService := vulndb .NewNVDService (cveRepository )
170198 mitreService := vulndb .NewMitreService (cweRepository )
171199 epssService := vulndb .NewEPSSService (nvdService , cveRepository )
172200 osvService := vulndb .NewOSVService (affectedCmpRepository )
173201 // cvelistService := vulndb.NewCVEListService(cveRepository)
174202 debianSecurityTracker := vulndb .NewDebianSecurityTracker (affectedCmpRepository )
175203
176- expoitDBService := vulndb .NewExploitDBService (nvdService , repositories .NewExploitRepository (database ))
204+ expoitDBService := vulndb .NewExploitDBService (nvdService , repositories .NewExploitRepository (db ))
177205
178- githubExploitDBService := vulndb .NewGithubExploitDBService (repositories .NewExploitRepository (database ))
206+ githubExploitDBService := vulndb .NewGithubExploitDBService (repositories .NewExploitRepository (db ))
179207
180208 if emptyOrContains (databasesToSync , "cwe" ) {
181209 now := time .Now ()
0 commit comments