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

cmd/db.go

@@ -11,6 +11,28 @@ import (
 	"strings"
 )
 
+func init() {
+	// Add whois command to root command
+	rootCmd.AddCommand(databaseQueryCmd)
+
+	// Add flags specific to whois command
+	databaseQueryCmd.Flags().StringVarP(&dbQueryTableName, "table", "t", "", "Table to query (results, creds, whois, subdomains, history, query_options)")
+	databaseQueryCmd.Flags().IntVarP(&dbQueryLimitRows, "limit", "l", 100, "Limit number of results")
+	databaseQueryCmd.Flags().StringVarP(&dbQueryNotNull, "not-null", "n", "", "Filter for non-null values (comma-separated list, e.g., 'password,email')")
+	databaseQueryCmd.Flags().StringVarP(&dbQueryColumns, "columns", "c", "", "Columns to display in output (comma-separated list, e.g., 'username,email,password')")
+	databaseQueryCmd.Flags().StringVarP(&dbQueryUserQuery, "query", "q", "", "User query to execute")
+	databaseQueryCmd.Flags().StringVarP(&dbQueryRawQuery, "raw-query", "r", "", "Raw SQL query to execute")
+	databaseQueryCmd.Flags().BoolVarP(&dbQueryListAll, "list-all", "a", false, "List all columns")
+
+	// Add mutually exclusive flags to query and raw-query
+	// Cannot use query and raw-query at the same time
+	databaseQueryCmd.MarkFlagsMutuallyExclusive("query", "raw-query")
+	// Raw query does not require a table
+	databaseQueryCmd.MarkFlagsMutuallyExclusive("query", "table")
+	// List all columns does not require a query or raw-query
+	databaseQueryCmd.MarkFlagsMutuallyExclusive("raw-query", "list-all")
+}
+
 var (
 	dbQueryTableName string
 	dbQueryLimitRows int
@@ -45,28 +67,6 @@ var (
 	}
 )
 
-func init() {
-	// Add whois command to root command
-	rootCmd.AddCommand(databaseQueryCmd)
-
-	// Add flags specific to whois command
-	databaseQueryCmd.Flags().StringVarP(&dbQueryTableName, "table", "t", "", "Table to query (results, creds, whois, subdomains, history, query_options)")
-	databaseQueryCmd.Flags().IntVarP(&dbQueryLimitRows, "limit", "l", 100, "Limit number of results")
-	databaseQueryCmd.Flags().StringVarP(&dbQueryNotNull, "not-null", "n", "", "Filter for non-null values (comma-separated list, e.g., 'password,email')")
-	databaseQueryCmd.Flags().StringVarP(&dbQueryColumns, "columns", "c", "", "Columns to display in output (comma-separated list, e.g., 'username,email,password')")
-	databaseQueryCmd.Flags().StringVarP(&dbQueryUserQuery, "query", "q", "", "User query to execute")
-	databaseQueryCmd.Flags().StringVarP(&dbQueryRawQuery, "raw-query", "r", "", "Raw SQL query to execute")
-	databaseQueryCmd.Flags().BoolVarP(&dbQueryListAll, "list-all", "a", false, "List all columns")
-
-	// Add mutually exclusive flags to query and raw-query
-	// Cannot use query and raw-query at the same time
-	databaseQueryCmd.MarkFlagsMutuallyExclusive("query", "raw-query")
-	// Raw query does not require a table
-	databaseQueryCmd.MarkFlagsMutuallyExclusive("query", "table")
-	// List all columns does not require a query or raw-query
-	databaseQueryCmd.MarkFlagsMutuallyExclusive("raw-query", "list-all")
-}
-
 func tableQuery(table Table) {
 
 	// Get the columns to query

cmd/export.go

@@ -1,14 +1,37 @@
 package cmd
 
 import (
+	"dehasher/internal/export"
+	"dehasher/internal/files"
+	"dehasher/internal/sqlite"
 	"fmt"
 	"github.com/spf13/cobra"
+	"go.uber.org/zap"
+	"strings"
 )
 
 func init() {
 	// Add Subcommand to db command
 	rootCmd.AddCommand(exportCmd)
 
+	// Add flags specific to export command
+	exportCmd.Flags().IntVarP(&exportLimitRows, "limit", "l", 100, "Limit number of results")
+	exportCmd.Flags().BoolVarP(&exportListAll, "list-all", "a", false, "List all columns")
+	exportCmd.Flags().StringVarP(&exportTableName, "table", "t", "", "Table to export")
+	exportCmd.Flags().StringVarP(&exportNotNull, "not-null", "n", "", "Filter for non-null values (comma-separated list, e.g., 'password,email')")
+	exportCmd.Flags().StringVarP(&exportColumns, "columns", "c", "", "Columns to display in output (comma-separated list, e.g., 'username,email,password')")
+	exportCmd.Flags().StringVarP(&exportUserQuery, "query", "q", "", "User query to execute")
+	exportCmd.Flags().StringVarP(&exportRawQuery, "raw-query", "r", "", "Raw SQL query to execute")
+	exportCmd.Flags().StringVarP(&exportFormat, "format", "f", "json", "Output format (json, yaml, xml, txt)")
+	exportCmd.Flags().StringVarP(&exportFile, "file", "o", "export", "File to output results to including extension")
+
+	// Add mutually exclusive flags to query and raw-query
+	// Cannot use query and raw-query at the same time
+	databaseQueryCmd.MarkFlagsMutuallyExclusive("query", "raw-query")
+	// Raw query does not require a table
+	databaseQueryCmd.MarkFlagsMutuallyExclusive("query", "table")
+	// List all columns does not require a query or raw-query
+	databaseQueryCmd.MarkFlagsMutuallyExclusive("raw-query", "list-all")
 }
 
 // DB export command
@@ -27,20 +50,196 @@ var (
 		Use:   "export",
 		Short: "Export database to file",
 		Run: func(cmd *cobra.Command, args []string) {
-			fmt.Println("Exporting database...")
+			fmt.Println("[*] Exporting database...")
+
+			// If Raw Query is set, execute it and export
+			if exportRawQuery != "" {
+				fmt.Println("[*] Executing Raw Query...")
+				exportRawDBQuery()
+				return
+			}
+
+			// Determine which table to query based on the tableTypeDBQuery parameter
+			table := GetTable(exportTableName)
+			if table == UnknownTable {
+				fmt.Printf("Error: Unknown table type '%s'.\n", exportTableName)
+				cmd.Help()
+				return
+			}
+
+			fmt.Println("[*] Querying Database...")
+			exportTableQuery(table)
 		},
 	}
 )
 
-func init() {
-	// Add flags specific to export command
-	exportCmd.Flags().IntVarP(&exportLimitRows, "limit", "l", 100, "Limit number of results")
-	exportCmd.Flags().BoolVarP(&exportListAll, "list-all", "a", false, "List all columns")
-	exportCmd.Flags().StringVarP(&exportTableName, "table", "t", "", "Table to export")
-	exportCmd.Flags().StringVarP(&exportNotNull, "not-null", "n", "", "Filter for non-null values (comma-separated list, e.g., 'password,email')")
-	exportCmd.Flags().StringVarP(&exportColumns, "columns", "c", "", "Columns to display in output (comma-separated list, e.g., 'username,email,password')")
-	exportCmd.Flags().StringVarP(&exportUserQuery, "query", "q", "", "User query to execute")
-	exportCmd.Flags().StringVarP(&exportRawQuery, "raw-query", "r", "", "Raw SQL query to execute")
-	exportCmd.Flags().StringVarP(&exportFormat, "format", "f", "json", "Output format (json, yaml, xml, txt)")
-	exportCmd.Flags().StringVarP(&exportFile, "file", "o", "export", "File to output results to including extension")
+// exportTableQuery queries a table and exports the results
+func exportTableQuery(table Table) {
+	// Get the columns to query
+	columns := []string{"*"}
+	if exportColumns != "" {
+		columns = strings.Split(exportColumns, ",")
+	}
+
+	// Get the not null fields
+	notNullFields := []string{}
+	if exportNotNull != "" {
+		notNullFields = strings.Split(exportNotNull, ",")
+	}
+
+	// Get the user query
+	userQuery := ""
+	if exportUserQuery != "" {
+		userQuery = exportUserQuery
+	}
+
+	// Get the limit
+	limit := exportLimitRows
+
+	// Get the object for the table
+	object := table.Object()
+
+	// Query the database
+	db := sqlite.GetDB()
+	query := db.Model(object).Select(columns)
+	if len(notNullFields) > 0 {
+		for _, field := range notNullFields {
+			query = query.Where(fmt.Sprintf("%s IS NOT NULL", field))
+		}
+	}
+	if userQuery != "" {
+		query = query.Where(userQuery)
+	}
+	if limit > 0 {
+		query = query.Limit(limit)
+	}
+	rows, err := query.Rows()
+	if err != nil {
+		zap.L().Error("export_query",
+			zap.String("message", "failed to execute query"),
+			zap.Error(err),
+		)
+		fmt.Printf("[!] Error executing query: %v\n", err)
+		return
+	}
+	defer rows.Close()
+
+	// Get the columns
+	cols, err := rows.Columns()
+	if err != nil {
+		zap.L().Error("export_query",
+			zap.String("message", "failed to get columns from query"),
+			zap.Error(err),
+		)
+		fmt.Printf("[!] Error getting columns from query: %v\n", err)
+		return
+	}
+
+	// Prepare data for export
+	var results []map[string]interface{}
+
+	// Process the rows
+	for rows.Next() {
+		values := make([]interface{}, len(cols))
+		pointers := make([]interface{}, len(cols))
+		for i := range values {
+			pointers[i] = &values[i]
+		}
+		if err := rows.Scan(pointers...); err != nil {
+			zap.L().Error("export_query",
+				zap.String("message", "failed to scan row from query"),
+				zap.Error(err),
+			)
+			fmt.Printf("[!] Error scanning row from query: %v\n", err)
+			return
+		}
+
+		// Create a map for this row
+		rowMap := make(map[string]interface{})
+		for i, col := range cols {
+			val := values[i]
+			rowMap[col] = val
+		}
+
+		results = append(results, rowMap)
+	}
+
+	// Export the results
+	exportResults(results)
+}
+
+// exportRawDBQuery executes a raw query and exports the results
+func exportRawDBQuery() {
+	db := sqlite.GetDB()
+	rows, err := db.Raw(exportRawQuery).Rows()
+	if err != nil {
+		zap.L().Error("export_raw_query",
+			zap.String("message", "failed to execute raw query"),
+			zap.Error(err),
+		)
+		fmt.Printf("[!] Error executing raw query: %v\n", err)
+		return
+	}
+	defer rows.Close()
+
+	columns, err := rows.Columns()
+	if err != nil {
+		zap.L().Error("export_raw_query",
+			zap.String("message", "failed to get columns from raw query"),
+			zap.Error(err),
+		)
+		fmt.Printf("[!] Error getting columns from raw query: %v\n", err)
+		return
+	}
+
+	// Prepare data for export
+	var results []map[string]interface{}
+
+	// Process the rows
+	for rows.Next() {
+		values := make([]interface{}, len(columns))
+		pointers := make([]interface{}, len(columns))
+		for i := range values {
+			pointers[i] = &values[i]
+		}
+		if err := rows.Scan(pointers...); err != nil {
+			zap.L().Error("export_raw_query",
+				zap.String("message", "failed to scan row from raw query"),
+				zap.Error(err),
+			)
+			fmt.Printf("[!] Error scanning row from raw query: %v\n", err)
+			return
+		}
+
+		// 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)
+	}
+
+	// Export the results
+	exportResults(results)
+}
+
+// exportResults exports the results to a file
+func exportResults(results []map[string]interface{}) {
+	// Get file type
+	fileType := files.GetFileType(exportFormat)
+
+	// Export results
+	err := export.WriteQueryResultsToFile(results, exportFile, fileType)
+	if err != nil {
+		zap.L().Error("export_results",
+			zap.String("message", "failed to write to file"),
+			zap.Error(err),
+		)
+		fmt.Printf("[!] Error writing to file: %v\n", err)
+		return
+	}
+
+	fmt.Printf("[+] Exported %d records to file: %s%s\n", len(results), exportFile, fileType.Extension())
 }

cmd/root.go

@@ -6,11 +6,12 @@ import (
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
 	"os"
+	"strings"
 )
 
 var (
 	// Global Flags
-	alternativeDatabase string
+	useLocalDatabase bool
 
 	// rootCmd is the base command for the CLI.
 	rootCmd = &cobra.Command{
@@ -58,14 +59,15 @@ func Execute() {
 }
 
 func init() {
+	// Attempt to retreive the useLocalDB
+	useLocalDatabase = badger.GetUseLocalDB()
+
 	// Hide the default help command
 	rootCmd.CompletionOptions.HiddenDefaultCmd = true
 
-	// Add global flags
-	rootCmd.PersistentFlags().StringVar(&alternativeDatabase, "alt-db", "", "Alternative database path")
-
 	// Add subcommands
 	rootCmd.AddCommand(setKeyCmd)
+	rootCmd.AddCommand(setLocalDb)
 }
 
 // Command to set API key
@@ -85,6 +87,32 @@ var setKeyCmd = &cobra.Command{
 	},
 }
 
+var setLocalDb = &cobra.Command{
+	Use:   "set-local-db [true|false]",
+	Short: "Set dehasher to use a local database path instead of the default (must be unset to use default)",
+	Args:  cobra.ExactArgs(1),
+	Run: func(cmd *cobra.Command, args []string) {
+		useLocal := strings.ToLower(args[0])
+
+		if useLocal == "true" {
+			useLocalDatabase = true
+		} else if useLocal == "false" {
+			useLocalDatabase = false
+		} else {
+			fmt.Println("Invalid argument. Please use 'true' or 'false'.")
+			return
+		}
+
+		// Store useLocal in badger DB
+		err := badger.StoreUseLocalDB(useLocalDatabase)
+		if err != nil {
+			fmt.Printf("Error storing local database useLocal: %v\n", err)
+			return
+		}
+		fmt.Println("Local database useLocal stored successfully")
+	},
+}
+
 // Helper functions to store API credentials
 func storeApiKey(key string) error {
 	err := badger.StoreKey(key)

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
+}