Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 825b983

Browse files
directionlessdirectionless
authored
Command line flags for query input and output (#8786)
1 parent ee880d7 commit 825b983

3 files changed

Lines changed: 86 additions & 28 deletions

File tree

‎osquery/devtools/devtools.h‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
#pragma once
1111

12+
#include <cstdio>
1213
#include <map>
1314
#include <string>
1415
#include <vector>
@@ -59,15 +60,17 @@ int launchIntoShell(int argc, char** argv);
5960
* @brief Pretty print a QueryData object
6061
*
6162
* This is a helper method which called osquery::beautify on the supplied
62-
* QueryData object and prints the results to stdout.
63+
* QueryData object and prints the results to the specified output.
6364
*
6465
* @param results The QueryData object to print
6566
* @param columns The order of the keys (since maps are unordered)
6667
* @param lengths A mutable set of column lengths
68+
* @param out The output file stream (defaults to stdout)
6769
*/
6870
void prettyPrint(const QueryData& results,
6971
const std::vector<std::string>& columns,
70-
std::map<std::string, size_t>& lengths);
72+
std::map<std::string, size_t>& lengths,
73+
FILE* out = stdout);
7174

7275
/**
7376
* @brief JSON print a QueryData object
@@ -76,8 +79,9 @@ void prettyPrint(const QueryData& results,
7679
* in a JSON format.
7780
*
7881
* @param q The QueryData object to print
82+
* @param out The output file stream (defaults to stdout)
7983
*/
80-
void jsonPrint(const QueryData& q);
84+
void jsonPrint(const QueryData& q, FILE* out = stdout);
8185

8286
/**
8387
* @brief JSON pretty print a QueryData object
@@ -86,8 +90,9 @@ void jsonPrint(const QueryData& q);
8690
* in a pretty JSON format.
8791
*
8892
* @param q The QueryData object to print
93+
* @param out The output file stream (defaults to stdout)
8994
*/
90-
void jsonPrettyPrint(const QueryData& q);
95+
void jsonPrettyPrint(const QueryData& q, FILE* out = stdout);
9196

9297
/**
9398
* @brief Compute a map of metadata about the supplied QueryData object

‎osquery/devtools/printer.cpp‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,8 @@ std::string generateRow(const Row& r,
111111

112112
void prettyPrint(const QueryData& results,
113113
const std::vector<std::string>& columns,
114-
std::map<std::string, size_t>& lengths) {
114+
std::map<std::string, size_t>& lengths,
115+
FILE* out) {
115116
if (results.size() == 0) {
116117
return;
117118
}
@@ -122,31 +123,31 @@ void prettyPrint(const QueryData& results,
122123
// Output a nice header wrapping the column names.
123124
auto separator = generateToken(lengths, columns);
124125
auto header = separator + generateHeader(lengths, columns) + separator;
125-
printf("%s", header.c_str());
126+
fprintf(out, "%s", header.c_str());
126127

127128
// Iterate each row and pretty print.
128129
for (const auto& row : results) {
129-
printf("%s", generateRow(row, lengths, columns).c_str());
130+
fprintf(out, "%s", generateRow(row, lengths, columns).c_str());
130131
}
131-
printf("%s", separator.c_str());
132+
fprintf(out, "%s", separator.c_str());
132133
}
133134

134-
void jsonPrint(const QueryData& q) {
135-
printf("[\n");
135+
void jsonPrint(const QueryData& q, FILE* out) {
136+
fprintf(out, "[\n");
136137
for (size_t i = 0; i < q.size(); ++i) {
137138
std::string row_string;
138139

139140
if (serializeRowJSON(q[i], row_string).ok()) {
140-
printf(" %s", row_string.c_str());
141+
fprintf(out, " %s", row_string.c_str());
141142
if (i < q.size() - 1) {
142-
printf(",\n");
143+
fprintf(out, ",\n");
143144
}
144145
}
145146
}
146-
printf("\n]\n");
147+
fprintf(out, "\n]\n");
147148
}
148149

149-
void jsonPrettyPrint(const QueryData& q) {
150+
void jsonPrettyPrint(const QueryData& q, FILE* out) {
150151
auto doc = JSON::newArray();
151152
for (const auto& row : q) {
152153
auto row_json = doc.getObject();
@@ -156,7 +157,7 @@ void jsonPrettyPrint(const QueryData& q) {
156157
}
157158
std::string doc_string;
158159
doc.toPrettyString(doc_string);
159-
printf("%s\n", doc_string.c_str());
160+
fprintf(out, "%s\n", doc_string.c_str());
160161
}
161162

162163
void computeRowLengths(const Row& r,

‎osquery/devtools/shell.cpp‎

Lines changed: 65 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
#include <osquery/utils/conversions/join.h>
4747
#include <osquery/utils/conversions/tryto.h>
4848
#include <osquery/utils/info/version.h>
49+
#include <osquery/utils/scope_guard.h>
4950

5051
#if defined(SQLITE_ENABLE_WHERETRACE)
5152
extern int sqlite3WhereTrace;
@@ -72,6 +73,14 @@ SHELL_FLAG(bool, L, false, "List all table names");
7273
SHELL_FLAG(string, A, "", "Select all from a table");
7374
SHELL_FLAG(string, connect, "", "Connect to an extension socket");
7475

76+
/// One-shot query execution flags.
77+
SHELL_FLAG(string, query, "", "Execute a single SQL query and exit");
78+
SHELL_FLAG(string, query_file, "", "Execute SQL query from a file and exit");
79+
SHELL_FLAG(string,
80+
output,
81+
"",
82+
"Write results to a file (omit or use '-' for stdout)");
83+
7584
DECLARE_string(nullvalue);
7685
DECLARE_string(extensions_socket);
7786
DECLARE_string(tls_hostname);
@@ -799,14 +808,16 @@ static void set_table_name(struct callback_data* p, const char* zName) {
799808

800809
static void pretty_print_if_needed(struct callback_data* pArg) {
801810
if ((pArg != nullptr) && pArg->mode == MODE_Pretty) {
811+
FILE* out = (pArg->out != nullptr) ? pArg->out : stdout;
802812
if (osquery::FLAGS_json_pretty) {
803-
osquery::jsonPrettyPrint(pArg->prettyPrint->results);
813+
osquery::jsonPrettyPrint(pArg->prettyPrint->results, out);
804814
} else if (osquery::FLAGS_json) {
805-
osquery::jsonPrint(pArg->prettyPrint->results);
815+
osquery::jsonPrint(pArg->prettyPrint->results, out);
806816
} else {
807817
osquery::prettyPrint(pArg->prettyPrint->results,
808818
pArg->prettyPrint->columns,
809-
pArg->prettyPrint->lengths);
819+
pArg->prettyPrint->lengths,
820+
out);
810821
}
811822
pArg->prettyPrint->results.clear();
812823
pArg->prettyPrint->columns.clear();
@@ -1740,7 +1751,19 @@ static int process_input(struct callback_data* p, FILE* in) {
17401751

17411752
if (nSql != 0) {
17421753
if (_all_whitespace(zSql) == 0) {
1743-
fprintf(stderr, "Error: incomplete SQL: %s\n", zSql);
1754+
// At EOF with non-whitespace SQL remaining - try to execute it.
1755+
// This allows piped input like "echo 'select 1' | osqueryd -S" to work
1756+
// without requiring a trailing semicolon.
1757+
p->cnt = 0;
1758+
rc = shell_exec(zSql, shell_callback, p, &zErrMsg);
1759+
if (rc != 0 || zErrMsg != nullptr) {
1760+
if (zErrMsg != nullptr) {
1761+
fprintf(stderr, "Error: %s\n", zErrMsg);
1762+
sqlite3_free(zErrMsg);
1763+
zErrMsg = nullptr;
1764+
}
1765+
errCnt++;
1766+
}
17441767
}
17451768
}
17461769
if (zSql != nullptr) {
@@ -1833,6 +1856,8 @@ int runPack(struct callback_data* data) {
18331856

18341857
int launchIntoShell(int argc, char** argv) {
18351858
struct callback_data data {};
1859+
auto pretty_print_guard =
1860+
scope_guard::create([&data]() { delete data.prettyPrint; });
18361861
main_init(&data);
18371862

18381863
#if defined(SQLITE_ENABLE_WHERETRACE)
@@ -1848,6 +1873,26 @@ int launchIntoShell(int argc, char** argv) {
18481873
signal(SIGINT, interrupt_handler);
18491874

18501875
data.out = stdout;
1876+
FILE* output_file = nullptr;
1877+
1878+
// Handle --output flag: write results to a file instead of stdout.
1879+
// Use "-" to explicitly write to stdout.
1880+
if (!FLAGS_output.empty() && FLAGS_output != "-") {
1881+
output_file = fopen(FLAGS_output.c_str(), "w");
1882+
if (output_file == nullptr) {
1883+
fprintf(stderr,
1884+
"Error: Cannot open output file '%s'\n",
1885+
FLAGS_output.c_str());
1886+
return 1;
1887+
}
1888+
data.out = output_file;
1889+
}
1890+
1891+
auto output_file_guard = scope_guard::create([&output_file]() {
1892+
if (output_file != nullptr) {
1893+
fclose(output_file);
1894+
}
1895+
});
18511896

18521897
// Set modes and settings from CLI flags.
18531898
data.showHeader = static_cast<int>(FLAGS_header);
@@ -1882,6 +1927,22 @@ int launchIntoShell(int argc, char** argv) {
18821927
delete[] cmd;
18831928
} else if (!FLAGS_pack.empty()) {
18841929
rc = runPack(&data);
1930+
} else if (!FLAGS_query.empty()) {
1931+
// Run a single query from --query flag
1932+
rc = runQuery(&data, FLAGS_query.c_str());
1933+
} else if (!FLAGS_query_file.empty()) {
1934+
// Read query from file and execute
1935+
std::string query_content;
1936+
auto status = readFile(FLAGS_query_file, query_content);
1937+
if (!status.ok()) {
1938+
fprintf(stderr,
1939+
"Error reading query file '%s': %s\n",
1940+
FLAGS_query_file.c_str(),
1941+
status.getMessage().c_str());
1942+
rc = 1;
1943+
} else {
1944+
rc = runQuery(&data, query_content.c_str());
1945+
}
18851946
} else if (argc > 1 && argv[1] != nullptr) {
18861947
// Run a command or statement from CLI
18871948
char* query = argv[1];
@@ -1890,12 +1951,6 @@ int launchIntoShell(int argc, char** argv) {
18901951
rc = (rc == 2) ? 0 : rc;
18911952
} else {
18921953
rc = runQuery(&data, query);
1893-
if (rc != 0) {
1894-
if (data.prettyPrint != nullptr) {
1895-
delete data.prettyPrint;
1896-
}
1897-
return rc;
1898-
}
18991954
}
19001955
} else {
19011956
// Run commands received from standard input
@@ -1926,9 +1981,6 @@ int launchIntoShell(int argc, char** argv) {
19261981

19271982
set_table_name(&data, nullptr);
19281983

1929-
if (data.prettyPrint != nullptr) {
1930-
delete data.prettyPrint;
1931-
}
19321984
return rc;
19331985
}
19341986
} // namespace osquery

0 commit comments

Comments
 (0)