(fix) dont create empty sqlite dbs

fixes #244
sensepost · Oct 10, 2024 · 83b61c3 · 83b61c3
1 parent 216c03d
commit 83b61c3
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 6 deletions.
diff --git a/cmd/report_convert.go b/cmd/report_convert.go
@@ -115,7 +115,7 @@ func init() {
 
 func convertFromDbTo(from string, writer writers.Writer) error {
 	var results = []*models.Result{}
-	conn, err := database.Connection(fmt.Sprintf("sqlite://%s", from), false)
+	conn, err := database.Connection(fmt.Sprintf("sqlite://%s", from), true, false)
 	if err != nil {
 		return err
 	}

diff --git a/cmd/report_generate.go b/cmd/report_generate.go
@@ -102,7 +102,7 @@ The output file is a zip archive with an index.html file containing the report.
 		}
 
 		// but, db-uri is the default
-		conn, err := database.Connection(generateCmdFlags.DbURI, false)
+		conn, err := database.Connection(generateCmdFlags.DbURI, true, false)
 		if err != nil {
 			log.Fatal("could not connect to database", "err", err)
 		}

diff --git a/cmd/report_list.go b/cmd/report_list.go
@@ -81,7 +81,7 @@ lines file.`)),
 		}
 
 		// db-uri is the default
-		conn, err := database.Connection(listCmdFlags.DbURI, false)
+		conn, err := database.Connection(listCmdFlags.DbURI, true, false)
 		if err != nil {
 			log.Error("could not connect to database", "err", err)
 			return

diff --git a/pkg/database/db.go b/pkg/database/db.go
@@ -2,7 +2,10 @@ package database
 
 import (
 	"errors"
+	"fmt"
 	"net/url"
+	"os"
+	"path/filepath"
 
 	"github.com/glebarez/sqlite"
 	"github.com/sensepost/gowitness/pkg/models"
@@ -13,7 +16,7 @@ import (
 )
 
 // Connection returns a Database connection based on a URI
-func Connection(uri string, debug bool) (*gorm.DB, error) {
+func Connection(uri string, shouldExist, debug bool) (*gorm.DB, error) {
 	var err error
 	var c *gorm.DB
 
@@ -31,6 +34,17 @@ func Connection(uri string, debug bool) (*gorm.DB, error) {
 
 	switch db.Scheme {
 	case "sqlite":
+		if shouldExist {
+			dbpath := filepath.Join(db.Host, db.Path)
+			dbpath = filepath.Clean(dbpath)
+
+			if _, err := os.Stat(dbpath); os.IsNotExist(err) {
+				return nil, fmt.Errorf("sqlite database file does not exist: %s", dbpath)
+			} else if err != nil {
+				return nil, fmt.Errorf("error checking sqlite database file: %w", err)
+			}
+		}
+
 		c, err = gorm.Open(sqlite.Open(db.Host+db.Path+"?cache=shared"), config)
 		if err != nil {
 			return nil, err

diff --git a/pkg/writers/db.go b/pkg/writers/db.go
@@ -22,7 +22,7 @@ type DbWriter struct {
 
 // NewDbWriter initialises a database writer
 func NewDbWriter(uri string, debug bool) (*DbWriter, error) {
-	c, err := database.Connection(uri, debug)
+	c, err := database.Connection(uri, false, debug)
 	if err != nil {
 		return nil, err
 	}

diff --git a/web/api/handler.go b/web/api/handler.go
@@ -18,7 +18,7 @@ type ApiHandler struct {
 func NewApiHandler(uri string, screenshotPath string) (*ApiHandler, error) {
 
 	// get a db handle
-	conn, err := database.Connection(uri, false)
+	conn, err := database.Connection(uri, true, false)
 	if err != nil {
 		return nil, err
 	}