Speed up shell and add max value size

This commit is contained in:
Teddy Reed 2015-03-18 12:01:58 -07:00
parent b8c658ec71
commit 91dce32095
10 changed files with 192 additions and 283 deletions

View File

@ -15,9 +15,13 @@
#include <vector>
#include <osquery/database/results.h>
#include <osquery/flags.h>
#include <osquery/tables.h>
namespace osquery {
DECLARE_int32(value_max);
/**
* @brief The core interface to executing osquery SQL commands
*

View File

@ -37,24 +37,19 @@ namespace osquery {
*/
int launchIntoShell(int argc, char** argv);
/**
* @brief Generate a pretty representation of a QueryData object
*
* @return The beautified string representation of the supplied QueryData
* @param order The order of the keys (since maps are unordered)
*/
std::string beautify(const QueryData& q, const std::vector<std::string>& order);
/**
* @brief Pretty print a QueryData object
*
* This is a helper method which called osquery::beautify on the supplied
* QueryData object and prints the results to stdout.
*
* @param q The QueryData object to print
* @param order The order of the keys (since maps are unordered)
* @param results The QueryData object to print
* @param columns The order of the keys (since maps are unordered)
* @param lengths A mutable set of column lengths
*/
void prettyPrint(const QueryData& q, const std::vector<std::string>& order);
void prettyPrint(const QueryData& results,
const std::vector<std::string>& columns,
std::map<std::string, size_t>& lengths);
/**
* @brief JSON print a QueryData object
@ -69,44 +64,49 @@ void jsonPrint(const QueryData& q);
/**
* @brief Compute a map of metadata about the supplied QueryData object
*
* @param q The QueryData object to analyze
* @param r A row to analyze
* @param lengths A mutable set of column lengths
* @param use_columns Calulate lengths of column names or values
*
* @return A map of string to int such that the key represents the "column" in
* the supplied QueryData and the int represents the length of the longest key
*/
std::map<std::string, int> computeQueryDataLengths(const QueryData& q);
void computeRowLengths(const Row& r,
std::map<std::string, size_t>& lengths,
bool use_columns = false);
/**
* @brief Generate the separator string for query results
*
* @param lengths The data returned from computeQueryDataLengths
* @param order The order of the keys (since maps are unordered)
* @param columns The order of the keys (since maps are unordered)
*
* @return A string, with a newline, representing your separator
*/
std::string generateSeparator(const std::map<std::string, int>& lengths,
const std::vector<std::string>& order);
std::string generateToken(const std::map<std::string, size_t>& lengths,
const std::vector<std::string>& columns);
/**
* @brief Generate the header string for query results
*
* @param lengths The data returned from computeQueryDataLengths
* @param order The order of the keys (since maps are unordered)
* @param columns The order of the keys (since maps are unordered)
*
* @return A string, with a newline, representing your header
*/
std::string generateHeader(const std::map<std::string, int>& lengths,
const std::vector<std::string>& order);
std::string generateHeader(const std::map<std::string, size_t>& lengths,
const std::vector<std::string>& columns);
/**
* @brief Generate a row string for query results
*
* @param r A row to analyze
* @param lengths The data returned from computeQueryDataLengths
* @param order The order of the keys (since maps are unordered)
* @param columns The order of the keys (since maps are unordered)
*
* @return A string, with a newline, representing your row
*/
std::string generateRow(const Row& r,
const std::map<std::string, int>& lengths,
const std::vector<std::string>& order);
const std::map<std::string, size_t>& lengths,
const std::vector<std::string>& columns);
}

View File

@ -12,112 +12,88 @@
#include <sstream>
#include <osquery/core.h>
#include <osquery/devtools.h>
#include <osquery/logger.h>
#include "osquery/devtools/devtools.h"
namespace osquery {
std::string beautify(const QueryData& q,
const std::vector<std::string>& order) {
auto lengths = computeQueryDataLengths(q);
if (q.size() == 0) {
return std::string();
std::string generateToken(const std::map<std::string, size_t>& lengths,
const std::vector<std::string>& columns) {
std::string output = "+";
for (const auto& col : columns) {
if (lengths.count(col) > 0) {
output += std::string(lengths.at(col) + 2, '-');
}
output += "+";
}
auto separator = generateSeparator(lengths, order);
std::ostringstream results;
results << "\n";
results << separator;
results << generateHeader(lengths, order);
results << separator;
for (const auto& r : q) {
results << generateRow(r, lengths, order);
}
results << separator;
return results.str();
output += "\n";
return output;
}
std::string generateSeparator(const std::map<std::string, int>& lengths,
const std::vector<std::string>& order) {
std::ostringstream separator;
separator << "+";
for (const auto& each : order) {
try {
for (int i = 0; i < lengths.at(each) + 2; ++i) {
separator << "-";
std::string generateHeader(const std::map<std::string, size_t>& lengths,
const std::vector<std::string>& columns) {
std::string output = "|";
for (const auto& col : columns) {
output += " " + col;
if (lengths.count(col) > 0) {
int buffer_size = lengths.at(col) - utf8StringSize(col) + 1;
if (buffer_size > 0) {
output += std::string(buffer_size, ' ');
} else {
output += ' ';
}
} catch (const std::out_of_range& e) {
LOG(ERROR) << "Error retrieving the \"" << each
<< "\" key in generateSeparator: " << e.what();
}
separator << "+";
output += "|";
}
separator << "\n";
return separator.str();
}
std::string generateHeader(const std::map<std::string, int>& lengths,
const std::vector<std::string>& order) {
std::ostringstream header;
header << "|";
for (const auto& each : order) {
header << " ";
header << each;
try {
for (int i = 0; i < (lengths.at(each) - utf8StringSize(each) + 1); ++i) {
header << " ";
}
} catch (const std::out_of_range& e) {
LOG(ERROR) << "Error retrieving the \"" << each
<< "\" key in generateHeader: " << e.what();
}
header << "|";
}
header << "\n";
return header.str();
output += "\n";
return output;
}
std::string generateRow(const Row& r,
const std::map<std::string, int>& lengths,
const std::map<std::string, size_t>& lengths,
const std::vector<std::string>& order) {
std::ostringstream row;
std::string value;
for (const auto& each : order) {
try {
value = r.at(each);
row << "| " << value;
for (int i = 0; i < (lengths.at(each) - utf8StringSize(r.at(each)) + 1);
++i) {
row << " ";
}
} catch (const std::out_of_range& e) {
for (const auto& foo : r) {
VLOG(1) << foo.first << " => " << foo.second;
}
LOG(ERROR) << "Error retrieving the \"" << each
<< "\" key in generateRow: " << e.what();
std::string output;
for (const auto& column : order) {
if (r.count(column) == 0 || lengths.count(column) == 0) {
continue;
}
// Print a terminator for the previous value or lhs, followed by spaces.
int buffer_size = lengths.at(column) - utf8StringSize(r.at(column)) + 1;
if (buffer_size > 0) {
output += "| " + r.at(column) + std::string(buffer_size, ' ');
}
}
if (row.str().size() > 0) {
if (output.size() > 0) {
// Only append if a row was added.
row << "|\n";
output += "|\n";
}
return row.str();
return output;
}
void prettyPrint(const QueryData& q, const std::vector<std::string>& order) {
std::cout << beautify(q, order);
void prettyPrint(const QueryData& results,
const std::vector<std::string>& columns,
std::map<std::string, size_t>& lengths) {
if (results.size() == 0) {
return;
}
// Call a final compute using the column names as minimum lengths.
computeRowLengths(results.front(), lengths, true);
// Output a nice header wrapping the column names.
auto separator = generateToken(lengths, columns);
auto header = separator + generateHeader(lengths, columns) + separator;
printf("%s", header.c_str());
// Iterate each row and pretty print.
for (const auto& row : results) {
printf("%s", generateRow(row, lengths, columns).c_str());
}
printf("%s", separator.c_str());
}
void jsonPrint(const QueryData& q) {
@ -135,31 +111,14 @@ void jsonPrint(const QueryData& q) {
printf("\n]\n");
}
std::map<std::string, int> computeQueryDataLengths(const QueryData& q) {
std::map<std::string, int> results;
if (q.size() == 0) {
return results;
void computeRowLengths(const Row& r,
std::map<std::string, size_t>& lengths,
bool use_columns) {
for (const auto& col : r) {
size_t current = (lengths.count(col.first) > 0) ? lengths.at(col.first) : 0;
size_t size =
(use_columns) ? utf8StringSize(col.first) : utf8StringSize(col.second);
lengths[col.first] = (size > current) ? size : current;
}
for (const auto& it : q.front()) {
results[it.first] = utf8StringSize(it.first);
}
for (const auto& row : q) {
for (const auto& it : row) {
try {
auto s = utf8StringSize(it.second);
if (s > results[it.first]) {
results[it.first] = s;
}
} catch (const std::out_of_range& e) {
LOG(ERROR) << "Error retrieving the \"" << it.first
<< "\" key in computeQueryDataLength: " << e.what();
}
}
}
return results;
}
}

View File

@ -10,9 +10,10 @@
#include <gtest/gtest.h>
#include <osquery/devtools.h>
#include <osquery/logger.h>
#include "osquery/devtools/devtools.h"
namespace osquery {
class PrinterTests : public testing::Test {
@ -20,132 +21,87 @@ class PrinterTests : public testing::Test {
QueryData q;
std::vector<std::string> order;
void SetUp() {
order = {"name", "age", "favorite_food", "lucky_number"};
order = {"name", "age", "food", "number"};
q = {
{
{"name", "Mike Jones"},
{"age", "39"},
{"favorite_food", "mac and cheese"},
{"lucky_number", "1"},
{"food", "mac and cheese"},
{"number", "1"},
},
{
{"name", "John Smith"},
{"age", "44"},
{"favorite_food", "peanut butter and jelly"},
{"lucky_number", "2"},
{"food", "peanut butter and jelly"},
{"number", "2"},
},
{
{"name", "Doctor Who"},
{"age", "2000"},
{"favorite_food", "fish sticks and custard"},
{"lucky_number", "11"},
{"food", "fish sticks and custard"},
{"number", "11"},
},
};
}
};
TEST_F(PrinterTests, test_compute_query_data_lengths) {
auto results = computeQueryDataLengths(q);
std::map<std::string, int> expected = {
{"name", 10}, {"age", 4}, {"favorite_food", 23}, {"lucky_number", 12},
};
EXPECT_EQ(results, expected);
std::map<std::string, size_t> lengths;
for (const auto& row : q) {
computeRowLengths(row, lengths);
}
// Check that all value lengths were maxed.
std::map<std::string, size_t> expected = {
{"name", 10}, {"age", 4}, {"food", 23}, {"number", 2}};
EXPECT_EQ(lengths, expected);
// Then compute lengths of column names.
computeRowLengths(q.front(), lengths, true);
expected = {{"name", 10}, {"age", 4}, {"food", 23}, {"number", 6}};
EXPECT_EQ(lengths, expected);
}
TEST_F(PrinterTests, test_generate_separator) {
auto results = generateSeparator(computeQueryDataLengths(q), order);
auto expected =
"+------------+------+-------------------------+--------------+\n";
EXPECT_EQ(results, expected);
}
std::map<std::string, size_t> lengths;
for (const auto& row : q) {
computeRowLengths(row, lengths);
}
TEST_F(PrinterTests, test_generate_separator_2) {
auto results =
generateSeparator(computeQueryDataLengths(q),
{"lucky_number", "age", "name", "favorite_food"});
auto expected =
"+--------------+------+------------+-------------------------+\n";
auto results = generateToken(lengths, order);
auto expected = "+------------+------+-------------------------+----+\n";
EXPECT_EQ(results, expected);
}
TEST_F(PrinterTests, test_generate_header) {
auto results = generateHeader(computeQueryDataLengths(q), order);
auto expected =
"| name | age | favorite_food | lucky_number |\n";
EXPECT_EQ(results, expected);
}
std::map<std::string, size_t> lengths;
for (const auto& row : q) {
computeRowLengths(row, lengths);
}
TEST_F(PrinterTests, test_generate_header_2) {
auto results =
generateHeader(computeQueryDataLengths(q),
{"lucky_number", "age", "name", "favorite_food"});
auto expected =
"| lucky_number | age | name | favorite_food |\n";
auto results = generateHeader(lengths, order);
auto expected = "| name | age | food | number |\n";
EXPECT_EQ(results, expected);
}
TEST_F(PrinterTests, test_generate_row) {
auto results = generateRow(q.back(), computeQueryDataLengths(q), order);
auto expected =
"| Doctor Who | 2000 | fish sticks and custard | 11 |\n";
EXPECT_EQ(results, expected);
}
std::map<std::string, size_t> lengths;
for (const auto& row : q) {
computeRowLengths(row, lengths);
}
TEST_F(PrinterTests, test_generate_row_2) {
auto results = generateRow(q.back(),
computeQueryDataLengths(q),
{"lucky_number", "age", "name", "favorite_food"});
auto expected =
"| 11 | 2000 | Doctor Who | fish sticks and custard |\n";
auto results = generateRow(q.front(), lengths, order);
auto expected = "| Mike Jones | 39 | mac and cheese | 1 |\n";
EXPECT_EQ(results, expected);
}
TEST_F(PrinterTests, test_beautify) {
auto result = beautify(q, order);
std::string expected = R"(
+------------+------+-------------------------+--------------+
| name | age | favorite_food | lucky_number |
+------------+------+-------------------------+--------------+
| Mike Jones | 39 | mac and cheese | 1 |
| John Smith | 44 | peanut butter and jelly | 2 |
| Doctor Who | 2000 | fish sticks and custard | 11 |
+------------+------+-------------------------+--------------+
)";
EXPECT_EQ(result, expected);
}
TEST_F(PrinterTests, test_unicode) {
QueryData augmented = {
{
{"name", "Mike Jones"},
{"age", "39"},
{"favorite_food", "mac and cheese"},
{"lucky_number", "1"},
},
{
{"name", "Àlex Smith"},
{"age", "44"},
{"favorite_food", "peanut butter and jelly"},
{"lucky_number", "2"},
},
{
{"name", "Doctor Who"},
{"age", "2000"},
{"favorite_food", "fish sticks and custard"},
{"lucky_number", "11"},
},
};
auto result = beautify(augmented, order);
std::string expected = R"(
+------------+------+-------------------------+--------------+
| name | age | favorite_food | lucky_number |
+------------+------+-------------------------+--------------+
| Mike Jones | 39 | mac and cheese | 1 |
| Àlex Smith | 44 | peanut butter and jelly | 2 |
| Doctor Who | 2000 | fish sticks and custard | 11 |
+------------+------+-------------------------+--------------+
)";
EXPECT_EQ(result, expected);
Row r = {{"name", "Àlex Smith"}};
std::map<std::string, size_t> lengths;
computeRowLengths(r, lengths);
std::map<std::string, size_t> expected = {{"name", 10}};
EXPECT_EQ(lengths, expected);
}
}

View File

@ -81,9 +81,9 @@
#else
#include <osquery/database/results.h>
#include <osquery/devtools.h>
#include <osquery/flags.h>
#include "osquery/devtools/devtools.h"
#include "osquery/sql/virtual_table.h"
// Json is a specific form of pretty printing.
@ -479,8 +479,9 @@ struct previous_mode_data {
** Pretty print structure
*/
struct prettyprint_data {
osquery::QueryData queryData;
std::vector<std::string> resultsOrder;
osquery::QueryData results;
std::vector<std::string> columns;
std::map<std::string, size_t> lengths;
};
/*
@ -999,21 +1000,6 @@ static void interrupt_handler(int NotUsed) {
}
#endif
void callback_row(int nArg, char **azArg, char **azCol, osquery::Row &r) {
for (int i = 0; i < nArg; i++) {
std::string header;
if (azCol[i] != nullptr) {
header = std::string(azCol[i]);
}
std::string result;
if (azArg[i] != nullptr) {
result = std::string(azArg[i]);
}
r[header] = result;
}
}
/*
** This is the callback routine that the shell
** invokes for each row of a query result.
@ -1025,15 +1011,20 @@ static int shell_callback(
switch (p->mode) {
case MODE_Pretty: {
if (p->prettyPrint->resultsOrder.size() == 0) {
if (p->prettyPrint->columns.size() == 0) {
for (i = 0; i < nArg; i++) {
p->prettyPrint->resultsOrder.push_back(std::string(azCol[i]));
p->prettyPrint->columns.push_back(std::string(azCol[i]));
}
}
osquery::Row r;
callback_row(nArg, azArg, azCol, r);
p->prettyPrint->queryData.push_back(r);
for (int i = 0; i < nArg; ++i) {
if (azCol[i] != nullptr && azArg[i] != nullptr) {
r[std::string(azCol[i])] = std::string(azArg[i]);
}
}
osquery::computeRowLengths(r, p->prettyPrint->lengths);
p->prettyPrint->results.push_back(r);
break;
}
case MODE_Line: {
@ -1845,13 +1836,15 @@ static int shell_exec(
if (pArg->mode == MODE_Pretty) {
if (osquery::FLAGS_json) {
osquery::jsonPrint(pArg->prettyPrint->queryData);
osquery::jsonPrint(pArg->prettyPrint->results);
} else {
osquery::prettyPrint(pArg->prettyPrint->queryData,
pArg->prettyPrint->resultsOrder);
osquery::prettyPrint(pArg->prettyPrint->results,
pArg->prettyPrint->columns,
pArg->prettyPrint->lengths);
}
pArg->prettyPrint->queryData.clear();
pArg->prettyPrint->resultsOrder.clear();
pArg->prettyPrint->results.clear();
pArg->prettyPrint->columns.clear();
pArg->prettyPrint->lengths.clear();
}
return rc;

View File

@ -10,7 +10,8 @@
#include <osquery/core.h>
#include <osquery/devtools.h>
#include "osquery/devtools/devtools.h"
int main(int argc, char *argv[]) {
// Parse/apply flags, start registry, load logger/config plugins.

View File

@ -18,6 +18,8 @@
namespace osquery {
FLAG(int32, value_max, 512, "Maximum returned row value size");
const std::map<tables::ConstraintOperator, std::string> kSQLOperatorRepr = {
{tables::EQUALS, "="},
{tables::GREATER_THAN, ">"},

View File

@ -208,17 +208,24 @@ static int xFilter(sqlite3_vtab_cursor *pVtabCursor,
Registry::call("table", pVtab->content->name, request, response);
// Now organize the response rows by column instead of row.
auto &data = pVtab->content->data;
for (const auto &row : response) {
for (const auto &column : pVtab->content->columns) {
try {
pVtab->content->data[column.first].push_back(row.at(column.first));
auto &value = row.at(column.first);
if (value.size() > FLAGS_value_max) {
data[column.first].push_back(value.substr(0, FLAGS_value_max));
} else {
data[column.first].push_back(value);
}
} catch (const std::out_of_range &e) {
VLOG(1) << "Table " << pVtab->content->name << " row "
<< pVtab->content->n << " did not include column "
<< column.first;
pVtab->content->data[column.first].push_back("");
data[column.first].push_back("");
}
}
pVtab->content->n++;
}

View File

@ -26,21 +26,9 @@ const std::vector<std::string> kShellHistoryFiles = {
".bash_history", ".zsh_history", ".zhistory", ".history",
};
Status genShellHistoryForUser(const Row& row, QueryData& results) {
std::string username;
std::string directory;
try {
username = row.at("username");
directory = row.at("directory");
} catch (const std::out_of_range& e) {
return Status(1, "Error retrieving query column");
}
std::vector<std::string> history_files;
if (!osquery::listFilesInDirectory(directory, history_files).ok()) {
return Status(1, "Cannot list files in: " + directory);
}
void genShellHistoryForUser(const std::string& username,
const std::string& directory,
QueryData& results) {
for (const auto& hfile : kShellHistoryFiles) {
boost::filesystem::path history_file = directory;
history_file /= hfile;
@ -59,32 +47,30 @@ Status genShellHistoryForUser(const Row& row, QueryData& results) {
results.push_back(r);
}
}
return Status(0, "OK");
}
QueryData genShellHistory(QueryContext& context) {
QueryData results;
std::string sql_str;
QueryData users;
if (!getuid()) {
sql_str = "SELECT username,directory FROM users";
// No uid is available, attempt to select from all users.
users = SQL::selectAllFrom("users");
} else {
struct passwd* pwd = nullptr;
pwd = getpwuid(getuid());
// TODO: https://github.com/facebook/osquery/issues/244
sql_str = "SELECT username,directory FROM users WHERE username = '" +
std::string(pwd->pw_name) + "';";
// A uid is available, select only the home directory for this user.
struct passwd* pwd = getpwuid(getuid());
if (pwd != nullptr && pwd->pw_name != nullptr) {
users = SQL::selectAllFrom(
"users", "username", EQUALS, std::string(pwd->pw_name));
}
}
auto sql = SQL(sql_str);
if (!sql.ok()) {
LOG(ERROR) << "Error executing SQL: " << sql.getMessageString();
return results;
}
for (const auto& row : sql.rows()) {
auto status = genShellHistoryForUser(row, results);
// Iterate over each user
for (const auto& row : users) {
if (row.count("username") > 0 && row.count("directory") > 0) {
genShellHistoryForUser(row.at("username"), row.at("directory"), results);
}
}
return results;

View File

@ -34,6 +34,7 @@ CANONICAL_PLATFORMS = {
"linux": "Ubuntu, CentOS",
"centos": "CentOS",
"ubuntu": "Ubuntu",
"utility": "Utility",
}
TEMPLATE_API_DEFINITION = """