Added useLocalDatabase global flag to allow for local db files for separate engagements.

internal/badger/badger.go

@@ -2,6 +2,7 @@ package badger
 
 import (
 	"crypto/sha256"
+	"errors"
 	"github.com/dgraph-io/badger/v4"
 	"go.uber.org/zap"
 	"log"
@@ -123,28 +124,42 @@ func GetKey() string {
 	return apiKey
 }
 
-func GetEmail() string {
-	var email string
+func GetUseLocalDB() bool {
+	var useLocal bool
 
 	err := db.View(func(txn *badger.Txn) error {
-		item, err := txn.Get([]byte("cfg:email"))
+		item, err := txn.Get([]byte("cfg:use_local_db"))
 		if err != nil {
-			return err // could be ErrKeyNotFound
+			// If key not found, set default to false and return nil
+			if errors.Is(err, badger.ErrKeyNotFound) {
+				// Store the default value for future use
+				err = StoreUseLocalDB(false)
+				if err != nil {
+					zap.L().Error("store_use_local_db",
+						zap.String("message", "failed to store use_local_db"),
+						zap.Error(err),
+					)
+					return err
+				}
+				return nil
+			}
+			// Return other errors
+			return err
 		}
 		return item.Value(func(val []byte) error {
-			email = string(val)
+			useLocal = val[0] == 1
 			return nil
 		})
 	})
 
 	if err != nil {
-		zap.L().Error("get_email",
-			zap.String("message", "failed to get email"),
+		zap.L().Error("get_use_local_db",
+			zap.String("message", "failed to get use_local_db"),
 			zap.Error(err),
 		)
 	}
 
-	return email
+	return useLocal
 }
 
 func StoreKey(apiKey string) error {
@@ -160,13 +175,20 @@ func StoreKey(apiKey string) error {
 	return err
 }
 
-func StoreEmail(email string) error {
+func StoreUseLocalDB(useLocal bool) error {
+	var local byte
+	if useLocal {
+		local = 1
+	} else {
+		local = 0
+	}
+
 	err := db.Update(func(txn *badger.Txn) error {
-		return txn.Set([]byte("cfg:email"), []byte(email))
+		return txn.Set([]byte("cfg:use_local_db"), []byte{local})
 	})
 	if err != nil {
-		zap.L().Error("set_email",
-			zap.String("message", "failed to set email"),
+		zap.L().Error("set_use_local_db",
+			zap.String("message", "failed to set use_local_db"),
 			zap.Error(err),
 		)
 	}

internal/export/export.go

@@ -75,3 +75,59 @@ func WriteToFile(results sqlite.DehashedResults, outputFile string, fileType fil
 	filePath := fmt.Sprintf("%s.%s", outputFile, fileType)
 	return ioutil.WriteFile(filePath, data, 0644)
 }
+
+// WriteQueryResultsToFile writes query results to a file in the specified format
+func WriteQueryResultsToFile(results []map[string]interface{}, outputFile string, fileType files.FileType) error {
+	var data []byte
+	var err error
+
+	switch fileType {
+	case files.JSON:
+		data, err = json.MarshalIndent(results, "", "  ")
+	case files.XML:
+		data, err = xml.MarshalIndent(results, "", "  ")
+	case files.YAML:
+		data, err = yaml.Marshal(results)
+	case files.TEXT:
+		var outStrings []string
+		for _, r := range results {
+			var rowStrings []string
+			for k, v := range r {
+				// Format the value to avoid array notation
+				var valueStr string
+				switch val := v.(type) {
+				case []string:
+					valueStr = strings.Join(val, ", ")
+				case []interface{}:
+					strSlice := make([]string, len(val))
+					for i, item := range val {
+						if item == nil {
+							strSlice[i] = ""
+						} else {
+							strSlice[i] = fmt.Sprintf("%v", item)
+						}
+					}
+					valueStr = strings.Join(strSlice, ", ")
+				default:
+					if v == nil {
+						valueStr = ""
+					} else {
+						valueStr = fmt.Sprintf("%v", v)
+					}
+				}
+				rowStrings = append(rowStrings, fmt.Sprintf("%s: %s", k, valueStr))
+			}
+			outStrings = append(outStrings, strings.Join(rowStrings, "\n")+"\n\n")
+		}
+		data = []byte(strings.Join(outStrings, ""))
+	default:
+		return errors.New("unsupported file type")
+	}
+
+	if err != nil {
+		return err
+	}
+
+	filePath := fmt.Sprintf("%s.%s", outputFile, fileType.String())
+	return os.WriteFile(filePath, data, 0644)
+}

internal/sqlite/db.go

@@ -363,3 +363,57 @@ func GetCredsCount(options *DBOptions) (int64, error) {
 
 	return count, nil
 }
+
+// ExecuteRawQuery executes a raw SQL query and returns the results as a slice of maps
+func ExecuteRawQuery(query string) ([]map[string]interface{}, error) {
+	db := GetDB()
+	rows, err := db.Raw(query).Rows()
+	if err != nil {
+		zap.L().Error("raw_query",
+			zap.String("message", "failed to execute raw query"),
+			zap.Error(err),
+		)
+		return nil, fmt.Errorf("failed to execute raw query: %w", err)
+	}
+	defer rows.Close()
+
+	columns, err := rows.Columns()
+	if err != nil {
+		zap.L().Error("raw_query",
+			zap.String("message", "failed to get columns from raw query"),
+			zap.Error(err),
+		)
+		return nil, fmt.Errorf("failed to get columns from raw query: %w", err)
+	}
+
+	var results []map[string]interface{}
+
+	for rows.Next() {
+		// Create a slice of interface{} to hold the values
+		values := make([]interface{}, len(columns))
+		pointers := make([]interface{}, len(columns))
+		for i := range values {
+			pointers[i] = &values[i]
+		}
+
+		// Scan the result into the pointers
+		if err := rows.Scan(pointers...); err != nil {
+			zap.L().Error("raw_query",
+				zap.String("message", "failed to scan row from raw query"),
+				zap.Error(err),
+			)
+			return nil, fmt.Errorf("failed to scan row from raw query: %w", err)
+		}
+
+		// Create a map for this row
+		rowMap := make(map[string]interface{})
+		for i, col := range columns {
+			val := values[i]
+			rowMap[col] = val
+		}
+
+		results = append(results, rowMap)
+	}
+
+	return results, nil
+}