From 93c9a353ca8b63406818bc96c1fb424909d07764 Mon Sep 17 00:00:00 2001 From: Evan Hosinski Date: Thu, 15 May 2025 15:08:32 -0400 Subject: [PATCH] Added useLocalDatabase global flag to allow for local db files for separate engagements. --- cmd/db.go | 44 ++++---- cmd/export.go | 223 ++++++++++++++++++++++++++++++++++++-- cmd/root.go | 36 +++++- internal/badger/badger.go | 46 ++++++-- internal/export/export.go | 56 ++++++++++ internal/sqlite/db.go | 54 +++++++++ 6 files changed, 409 insertions(+), 50 deletions(-) diff --git a/cmd/db.go b/cmd/db.go index 0193d01..ee22ad5 100644 --- a/cmd/db.go +++ b/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 diff --git a/cmd/export.go b/cmd/export.go index bef66e5..dab4bb9 100644 --- a/cmd/export.go +++ b/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()) } diff --git a/cmd/root.go b/cmd/root.go index a7ba62c..8b57bfa 100644 --- a/cmd/root.go +++ b/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) diff --git a/internal/badger/badger.go b/internal/badger/badger.go index 8452fef..3950a55 100644 --- a/internal/badger/badger.go +++ b/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), ) } diff --git a/internal/export/export.go b/internal/export/export.go index 17fb675..18781ad 100644 --- a/internal/export/export.go +++ b/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) +} diff --git a/internal/sqlite/db.go b/internal/sqlite/db.go index bcc6d71..b5e3db9 100644 --- a/internal/sqlite/db.go +++ b/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 +}