diff --git a/cmd/export.go b/cmd/export.go index ec0409d..533d946 100644 --- a/cmd/export.go +++ b/cmd/export.go @@ -16,8 +16,8 @@ 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().BoolVarP(&exportListAll, "list-all", "a", false, "List all tables and their columns") + exportCmd.Flags().StringVarP(&exportTableName, "table", "t", "", "Table to export (results, creds, whois, subdomains, history, runs)") 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, "user-query", "q", "", "User query to execute") @@ -50,6 +50,12 @@ var ( Use: "export", Short: "Export database to file", Run: func(cmd *cobra.Command, args []string) { + // If list-all flag is set, list all tables and columns + if exportListAll { + listAvailableTables() + return + } + fmt.Println("[*] Exporting database...") // If Raw Query is set, execute it and export @@ -59,11 +65,65 @@ var ( return } + // Validate table name + if exportTableName == "" { + fmt.Println("[!] Error: Table name is required. Use -t or --table to specify a table.") + fmt.Println("[*] Available tables: results, creds, whois, subdomains, history, runs") + fmt.Println("[*] Use --list-all to see all tables and their columns.") + return + } + + if !isValidTable(exportTableName) { + fmt.Printf("[!] Error: Unknown table '%s'.\n", exportTableName) + fmt.Println("[*] Available tables: results, creds, whois, subdomains, history, runs") + fmt.Println("[*] Use --list-all to see all tables and their columns.") + return + } + + // Validate columns if specified + if exportColumns != "" { + columns := strings.Split(exportColumns, ",") + invalidColumns := validateColumns(exportTableName, columns) + if len(invalidColumns) > 0 { + fmt.Printf("[!] Error: Invalid column(s) for table '%s': %s\n", + exportTableName, strings.Join(invalidColumns, ", ")) + fmt.Println("[*] Available columns for this table:") + for i := 0; i < len(availableTables[exportTableName]); i += 5 { + end := i + 5 + if end > len(availableTables[exportTableName]) { + end = len(availableTables[exportTableName]) + } + fmt.Printf(" %s\n", strings.Join(availableTables[exportTableName][i:end], ", ")) + } + return + } + } + + // Validate not-null fields if specified + if exportNotNull != "" { + notNullFields := strings.Split(exportNotNull, ",") + invalidFields := validateColumns(exportTableName, notNullFields) + if len(invalidFields) > 0 { + fmt.Printf("[!] Error: Invalid not-null field(s) for table '%s': %s\n", + exportTableName, strings.Join(invalidFields, ", ")) + fmt.Println("[*] Available columns for this table:") + for i := 0; i < len(availableTables[exportTableName]); i += 5 { + end := i + 5 + if end > len(availableTables[exportTableName]) { + end = len(availableTables[exportTableName]) + } + fmt.Printf(" %s\n", strings.Join(availableTables[exportTableName][i:end], ", ")) + } + return + } + } + // Determine which table to query based on the tableTypeDBQuery parameter table := sqlite.GetTable(exportTableName) if table == sqlite.UnknownTable { - fmt.Printf("Error: Unknown table type '%s'.\n", exportTableName) - cmd.Help() + fmt.Printf("[!] Error: Unknown table type '%s'.\n", exportTableName) + fmt.Println("[*] Available tables: results, creds, whois, subdomains, history, runs") + fmt.Println("[*] Use --list-all to see all tables and their columns.") return } @@ -99,6 +159,12 @@ func exportTableQuery(table sqlite.Table) { // Get the object for the table object := table.Object() + // Check if object is nil (invalid table) + if object == nil { + fmt.Printf("[!] Error: Table '%s' is not valid or does not exist.\n", exportTableName) + return + } + // Query the database db := sqlite.GetDB() query := db.Model(object).Select(columns) diff --git a/cmd/query.go b/cmd/query.go index 82dc6b8..dc1c3a0 100644 --- a/cmd/query.go +++ b/cmd/query.go @@ -7,28 +7,137 @@ import ( "fmt" "github.com/spf13/cobra" "go.uber.org/zap" - "os" "strings" ) +// Map of available tables and their columns +var availableTables = map[string][]string{ + "results": { + "id", "created_at", "updated_at", "deleted_at", "dehashed_id", "email", "ip_address", "username", + "password", "hashed_password", "hash_type", "name", "vin", "license_plate", "url", "social", + "cryptocurrency_address", "address", "phone", "company", "database_name", + }, + "creds": { + "id", "created_at", "updated_at", "deleted_at", "email", "username", "password", + }, + "whois": { + "id", "created_at", "updated_at", "deleted_at", "contact_email", "created_date", + "domain_name", "domain_name_ext", "expires_date", "status", "whois_server", + }, + "subdomains": { + "id", "created_at", "updated_at", "deleted_at", "domain", "first_seen", "last_seen", + }, + "history": { + "id", "created_at", "updated_at", "deleted_at", "domain_name", "domain_type", + "registrar_name", "whois_server", "created_date_iso8601", "updated_date_iso8601", "expires_date_iso8601", + }, + "runs": { + "id", "created_at", "updated_at", "deleted_at", "max_records", "max_requests", "starting_page", + "output_format", "output_file", "regex_match", "wildcard_match", "username_query", "email_query", + "ip_query", "pass_query", "hash_query", "name_query", "domain_query", "vin_query", "license_plate_query", + "address_query", "phone_query", "social_query", "crypto_address_query", "print_balance", "creds_only", + }, +} + +// Function to list available tables and their columns +func listAvailableTables() { + fmt.Println("Available tables and columns:") + + // Prepare data for pretty.Table + headers := []string{"Table", "Columns"} + var tableRows [][]string + + // Sort tables alphabetically for consistent output + var tableNames []string + for tableName := range availableTables { + tableNames = append(tableNames, tableName) + } + + // Simple bubble sort for table names + for i := 0; i < len(tableNames)-1; i++ { + for j := 0; j < len(tableNames)-i-1; j++ { + if tableNames[j] > tableNames[j+1] { + tableNames[j], tableNames[j+1] = tableNames[j+1], tableNames[j] + } + } + } + + // Create rows for the table + for _, tableName := range tableNames { + columns := availableTables[tableName] + + // Format columns with line breaks for better readability + var formattedColumns string + for i := 0; i < len(columns); i += 5 { + end := i + 5 + if end > len(columns) { + end = len(columns) + } + if i > 0 { + formattedColumns += "\n" + } + formattedColumns += strings.Join(columns[i:end], ", ") + } + + tableRows = append(tableRows, []string{tableName, formattedColumns}) + } + + // Display the table + pretty.Table(headers, tableRows) +} + +// Function to validate table name +func isValidTable(tableName string) bool { + _, exists := availableTables[tableName] + return exists +} + +// Function to validate column names for a specific table +func validateColumns(tableName string, columns []string) []string { + if tableName == "" || columns == nil || len(columns) == 0 || columns[0] == "*" { + return nil + } + + tableColumns, exists := availableTables[tableName] + if !exists { + return []string{fmt.Sprintf("Table '%s' does not exist", tableName)} + } + + var invalidColumns []string + for _, col := range columns { + valid := false + for _, tableCol := range tableColumns { + if col == tableCol { + valid = true + break + } + } + if !valid { + invalidColumns = append(invalidColumns, col) + } + } + + return invalidColumns +} + func init() { // Add whois command to root command rootCmd.AddCommand(queryCmd) // Add flags specific to whois command - queryCmd.Flags().StringVarP(&dbQueryTableName, "table", "t", "", "Table to query (results, creds, whois, subdomains, history, query_options)") + queryCmd.Flags().StringVarP(&dbQueryTableName, "table", "t", "", "Table to query (results, creds, whois, subdomains, history, runs)") queryCmd.Flags().IntVarP(&dbQueryLimitRows, "limit", "l", 100, "Limit number of results") queryCmd.Flags().StringVarP(&dbQueryNotNull, "not-null", "n", "", "Filter for non-null values (comma-separated list, e.g., 'password,email')") queryCmd.Flags().StringVarP(&dbQueryColumns, "columns", "c", "", "Columns to display in output (comma-separated list, e.g., 'username,email,password')") queryCmd.Flags().StringVarP(&dbQueryUserQuery, "user-query", "q", "", "User query to execute") queryCmd.Flags().StringVarP(&dbQueryRawQuery, "raw-query", "r", "", "Raw SQL query to execute") - queryCmd.Flags().BoolVarP(&dbQueryListAll, "list-all", "a", false, "List all columns") + queryCmd.Flags().BoolVarP(&dbQueryListAll, "list-all", "a", false, "List all tables and their columns") // Add mutually exclusive flags to query and raw-query // Cannot use query and raw-query at the same time queryCmd.MarkFlagsMutuallyExclusive("user-query", "raw-query") // Raw query does not require a table - queryCmd.MarkFlagsMutuallyExclusive("user-query", "table") + queryCmd.MarkFlagsMutuallyExclusive("raw-query", "table") // List all columns does not require a query or raw-query queryCmd.MarkFlagsMutuallyExclusive("raw-query", "list-all") } @@ -47,20 +156,81 @@ var ( Short: "Query the database", Long: `Query the database for various information.`, Run: func(cmd *cobra.Command, args []string) { + // If list-all flag is set, list all tables and columns + if dbQueryListAll { + listAvailableTables() + return + } + // If Raw Query is set, execute it and return if dbQueryRawQuery != "" { fmt.Println("[*] Executing Raw Query...") rawDBQuery() - os.Exit(1) + return + } + + // Validate table name + if dbQueryTableName == "" { + fmt.Println("[!] Error: Table name is required. Use -t or --table to specify a table.") + fmt.Println("[*] Available tables: results, creds, whois, subdomains, history, runs") + fmt.Println("[*] Use --list-all to see all tables and their columns.") + return + } + + if !isValidTable(dbQueryTableName) { + fmt.Printf("[!] Error: Unknown table '%s'.\n", dbQueryTableName) + fmt.Println("[*] Available tables: results, creds, whois, subdomains, history, runs") + fmt.Println("[*] Use --list-all to see all tables and their columns.") + return + } + + // Validate columns if specified + if dbQueryColumns != "" { + columns := strings.Split(dbQueryColumns, ",") + invalidColumns := validateColumns(dbQueryTableName, columns) + if len(invalidColumns) > 0 { + fmt.Printf("[!] Error: Invalid column(s) for table '%s': %s\n", + dbQueryTableName, strings.Join(invalidColumns, ", ")) + fmt.Println("[*] Available columns for this table:") + for i := 0; i < len(availableTables[dbQueryTableName]); i += 5 { + end := i + 5 + if end > len(availableTables[dbQueryTableName]) { + end = len(availableTables[dbQueryTableName]) + } + fmt.Printf(" %s\n", strings.Join(availableTables[dbQueryTableName][i:end], ", ")) + } + return + } + } + + // Validate not-null fields if specified + if dbQueryNotNull != "" { + notNullFields := strings.Split(dbQueryNotNull, ",") + invalidFields := validateColumns(dbQueryTableName, notNullFields) + if len(invalidFields) > 0 { + fmt.Printf("[!] Error: Invalid not-null field(s) for table '%s': %s\n", + dbQueryTableName, strings.Join(invalidFields, ", ")) + fmt.Println("[*] Available columns for this table:") + for i := 0; i < len(availableTables[dbQueryTableName]); i += 5 { + end := i + 5 + if end > len(availableTables[dbQueryTableName]) { + end = len(availableTables[dbQueryTableName]) + } + fmt.Printf(" %s\n", strings.Join(availableTables[dbQueryTableName][i:end], ", ")) + } + return + } } // Determine which table to query based on the tableTypeDBQuery parameter table := sqlite.GetTable(dbQueryTableName) if table == sqlite.UnknownTable { - fmt.Printf("Error: Unknown table type '%s'.\n", dbQueryTableName) - cmd.Help() + fmt.Printf("[!] Error: Unknown table type '%s'.\n", dbQueryTableName) + fmt.Println("[*] Available tables: results, creds, whois, subdomains, history, runs") + fmt.Println("[*] Use --list-all to see all tables and their columns.") return } + fmt.Println("[*] Querying Database...") tableQuery(table) }, @@ -93,6 +263,12 @@ func tableQuery(table sqlite.Table) { // Get the object for the table object := table.Object() + // Check if object is nil (invalid table) + if object == nil { + fmt.Printf("[!] Error: Table '%s' is not valid or does not exist.\n", dbQueryTableName) + return + } + // Query the database db := sqlite.GetDB() query := db.Model(object).Select(columns) @@ -114,6 +290,7 @@ func tableQuery(table sqlite.Table) { zap.Error(err), ) fmt.Printf("[!] Error executing query: %v\n", err) + return } defer rows.Close() @@ -125,6 +302,7 @@ func tableQuery(table sqlite.Table) { zap.Error(err), ) fmt.Printf("[!] Error getting columns from query: %v\n", err) + return } // Prepare data for pretty.Table @@ -144,6 +322,7 @@ func tableQuery(table sqlite.Table) { zap.Error(err), ) fmt.Printf("[!] Error scanning row from query: %v\n", err) + continue } // Convert row values to strings @@ -202,6 +381,7 @@ func rawDBQuery() { zap.Error(err), ) fmt.Printf("[!] Error executing raw query: %v\n", err) + return } defer rows.Close() @@ -212,6 +392,7 @@ func rawDBQuery() { zap.Error(err), ) fmt.Printf("[!] Error getting columns from raw query: %v\n", err) + return } // Prepare data for pretty.Table @@ -231,6 +412,7 @@ func rawDBQuery() { zap.Error(err), ) fmt.Printf("[!] Error scanning row from raw query: %v\n", err) + continue } // Convert row values to strings