TEST_F(QueryTests, test_add_and_get_current_results) {
// Test adding a "current" set of results to a scheduled query instance.
auto query = getOsqueryScheduledQuery();
auto cf = Query("foobar", query);
auto status = cf.addNewResults(getTestDBExpectedResults(), db_);
EXPECT_TRUE(status.ok());
EXPECT_EQ(status.toString(), "OK");
// Simulate results from several schedule runs, calculate differentials.
for (auto result : getTestDBResultStream()) {
// Get the results from the previous query execution (from RocksDB).
QueryData previous_qd;
auto status = cf.getPreviousQueryResults(previous_qd, db_);
EXPECT_TRUE(status.ok());
EXPECT_EQ(status.toString(), "OK");
// Add the "current" results and output the differentials.
DiffResults dr;
auto s = cf.addNewResults(result.second, dr, true, db_);
EXPECT_TRUE(s.ok());
// Call the diffing utility directly.
DiffResults expected = diff(previous_qd, result.second);
EXPECT_EQ(dr, expected);
// After Query::addNewResults the previous results are now current.
QueryData qd;
cf.getPreviousQueryResults(qd, db_);
EXPECT_EQ(qd, result.second);
}
}
Status Query::addNewResults(const QueryData& current_qd,
DiffResults& dr,
bool calculate_diff,
DBHandleRef db) {
// Get the rows from the last run of this query name.
QueryData previous_qd;
auto status = getPreviousQueryResults(previous_qd);
if (!status.ok()) {
return status;
}
// Calculate the differential between previous and current query results.
if (calculate_diff) {
dr = diff(previous_qd, current_qd);
}
if (previous_qd.size() == 0 || dr.added.size() != 0 ||
dr.removed.size() != 0) {
// Replace the "previous" query data with the current.
std::string json;
status = serializeQueryDataJSON(current_qd, json);
if (!status.ok()) {
return status;
}
status = db->Put(kQueries, name_, json);
if (!status.ok()) {
return status;
}
}
return Status(0, "OK");
}
Status Query::addNewResults(const QueryData& current_qd,
DiffResults& dr,
bool calculate_diff,
DBHandleRef db) {
// Get the rows from the last run of this query name.
QueryData previous_qd;
auto status = getPreviousQueryResults(previous_qd);
// Sanitize all non-ASCII characters from the query data values.
QueryData escaped_current_qd;
escapeQueryData(current_qd, escaped_current_qd);
// Calculate the differential between previous and current query results.
if (calculate_diff) {
dr = diff(previous_qd, escaped_current_qd);
}
// Replace the "previous" query data with the current.
std::string json;
status = serializeQueryDataJSON(escaped_current_qd, json);
if (!status.ok()) {
return status;
}
status = db->Put(kQueries, name_, json);
if (!status.ok()) {
return status;
}
return Status(0, "OK");
}
TEST_F(QueryTests, test_query_name_not_found_in_db) {
// Try to retrieve results from a query that has not executed.
QueryData previous_qd;
auto query = getOsqueryScheduledQuery();
auto cf = Query("not_a_real_query", query);
auto status = cf.getPreviousQueryResults(previous_qd, db_);
EXPECT_FALSE(status.ok());
}
TEST_F(QueryTests, test_get_query_results) {
// Grab an expected set of query data and add it as the previous result.
auto encoded_qd = getSerializedQueryDataJSON();
auto query = getOsqueryScheduledQuery();
auto status = db_->Put(kQueries, "foobar", encoded_qd.first);
EXPECT_TRUE(status.ok());
// Use the Query retrieval API to check the now "previous" result.
QueryData previous_qd;
auto cf = Query("foobar", query);
status = cf.getPreviousQueryResults(previous_qd, db_);
EXPECT_TRUE(status.ok());
}
Status Query::getPreviousQueryResults(QueryData& results) {
return getPreviousQueryResults(results, DBHandle::getInstance());
}
Status Query::addNewResults(QueryData current_qd,
const uint64_t current_epoch,
uint64_t& counter,
DiffResults& dr,
bool calculate_diff) const {
// The current results are 'fresh' when not calculating a differential.
bool fresh_results = !calculate_diff;
bool new_query = false;
if (!isQueryNameInDatabase()) {
// This is the first encounter of the scheduled query.
fresh_results = true;
LOG(INFO) << "Storing initial results for new scheduled query: " << name_;
saveQuery(name_, query_);
} else if (getPreviousEpoch() != current_epoch) {
fresh_results = true;
LOG(INFO) << "New Epoch " << current_epoch << " for scheduled query "
<< name_;
} else if (isNewQuery()) {
// This query is 'new' in that the previous results may be invalid.
new_query = true;
LOG(INFO) << "Scheduled query has been updated: " + name_;
saveQuery(name_, query_);
}
// Use a 'target' avoid copying the query data when serializing and saving.
// If a differential is requested and needed the target remains the original
// query data, otherwise the content is moved to the differential's added set.
const auto* target_gd = ¤t_qd;
bool update_db = true;
if (!fresh_results && calculate_diff) {
// Get the rows from the last run of this query name.
QueryDataSet previous_qd;
auto status = getPreviousQueryResults(previous_qd);
if (!status.ok()) {
return status;
}
// Calculate the differential between previous and current query results.
dr = diff(previous_qd, current_qd);
update_db = (!dr.added.empty() || !dr.removed.empty());
} else {
dr.added = std::move(current_qd);
target_gd = &dr.added;
}
counter = getQueryCounter(fresh_results || new_query);
auto status =
setDatabaseValue(kQueries, name_ + "counter", std::to_string(counter));
if (!status.ok()) {
return status;
}
if (update_db) {
// Replace the "previous" query data with the current.
std::string json;
status = serializeQueryDataJSON(*target_gd, json);
if (!status.ok()) {
return status;
}
status = setDatabaseValue(kQueries, name_, json);
if (!status.ok()) {
return status;
}
status = setDatabaseValue(
kQueries, name_ + "epoch", std::to_string(current_epoch));
if (!status.ok()) {
return status;
}
}
return Status(0, "OK");
}
