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

2025-05-15 15:08:32 -04:00
parent bb43695822
commit 93c9a353ca
6 changed files with 409 additions and 50 deletions
@@ -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
@@ -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())
 }
@@ -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)