tables: Refactor shell_history to use generators (#6541)

osquery/sql/dynamic_table_row.h

@@ -58,9 +58,15 @@ class DynamicTableRowHolder {
   inline std::string& operator[](const std::string& key) {
     return (*row)[key];
   }
+  inline const std::string& operator[](const std::string& key) const {
+    return (*row)[key];
+  }
   inline std::string& operator[](std::string&& key) {
     return (*row)[key];
   }
+  inline const std::string& operator[](std::string&& key) const {
+    return (*row)[key];
+  }
   inline size_t count(const std::string& key) {
     return (*row).count(key);
   }

osquery/tables/system/posix/shell_history.cpp

@@ -12,76 +12,119 @@
 
 #include <osquery/core.h>
 #include <osquery/filesystem/filesystem.h>
-#include <osquery/logger.h>
 #include <osquery/tables.h>
-#include <osquery/tables/system/system_utils.h>
 #include <osquery/tables/system/posix/shell_history.h>
-#include <osquery/utils/conversions/split.h>
+#include <osquery/tables/system/system_utils.h>
 #include <osquery/utils/system/system.h>
 
 namespace osquery {
 namespace tables {
 
 const std::vector<std::string> kShellHistoryFiles = {
-    ".bash_history", ".zsh_history", ".zhistory", ".history", ".sh_history",
+    ".bash_history",
+    ".zsh_history",
+    ".zhistory",
+    ".history",
+    ".sh_history",
 };
 
-void genShellHistoryFromFile(const std::string& uid,
+struct HistoryState {
-                             const boost::filesystem::path& history_file,
+  std::string content;
-                             QueryData& results) {
+  std::regex bash_timestamp_rx{"^#([0-9]+)$"};
-  std::string history_content;
+  std::regex zsh_timestamp_rx{"^: {0,10}([0-9]{1,11}):[0-9]+;(.*)$"};
-  if (forensicReadFile(history_file, history_content).ok()) {
+  std::string prev_bash_timestamp;
-    std::regex bash_timestamp_rx("^#([0-9]+)$");
+};
-    std::regex zsh_timestamp_rx("^: {0,10}([0-9]{1,11}):[0-9]+;(.*)$");
 
-    std::string prev_bash_timestamp;
+void genShellHistoryFromFile(
-    for (const auto& line : split(history_content, "\n")) {
+    const std::string& uid,
-      std::smatch bash_timestamp_matches;
+    const boost::filesystem::path& history_file,
-      std::smatch zsh_timestamp_matches;
+    std::function<void(DynamicTableRowHolder& row)> predicate) {
+  struct HistoryState hState;
 
-      if (prev_bash_timestamp.empty() &&
+  auto parseLine =
-          std::regex_search(line, bash_timestamp_matches, bash_timestamp_rx)) {
+      [&hState, &uid, &history_file, &predicate](std::string& line) {
-        prev_bash_timestamp = bash_timestamp_matches[1];
+        std::smatch bash_timestamp_matches;
-        continue;
+        std::smatch zsh_timestamp_matches;
-      }
 
-      Row r;
+        if (hState.prev_bash_timestamp.empty() &&
+            std::regex_search(
+                line, bash_timestamp_matches, hState.bash_timestamp_rx)) {
+          hState.prev_bash_timestamp = bash_timestamp_matches[1];
+          return;
+        }
 
-      if (!prev_bash_timestamp.empty()) {
+        auto r = make_table_row();
-        r["time"] = INTEGER(prev_bash_timestamp);
-        r["command"] = line;
-        prev_bash_timestamp.clear();
-      } else if (std::regex_search(
-                     line, zsh_timestamp_matches, zsh_timestamp_rx)) {
-        std::string timestamp = zsh_timestamp_matches[1];
-        r["time"] = INTEGER(timestamp);
-        r["command"] = zsh_timestamp_matches[2];
-      } else {
-        r["time"] = INTEGER(0);
-        r["command"] = line;
-      }
 
-      r["uid"] = uid;
+        if (!hState.prev_bash_timestamp.empty()) {
-      r["history_file"] = history_file.string();
+          r["time"] = INTEGER(hState.prev_bash_timestamp);
-      results.push_back(r);
+          r["command"] = std::move(line);
+          hState.prev_bash_timestamp.clear();
+        } else if (std::regex_search(
+                       line, zsh_timestamp_matches, hState.zsh_timestamp_rx)) {
+          std::string timestamp = zsh_timestamp_matches[1];
+          r["time"] = INTEGER(timestamp);
+          r["command"] = zsh_timestamp_matches[2];
+        } else {
+          r["time"] = INTEGER(0);
+          r["command"] = std::move(line);
+        }
+
+        r["uid"] = uid;
+        r["history_file"] = history_file.string();
+        predicate(r);
+      };
+
+  auto parseChunk = [&hState, &parseLine](std::string& buffer, size_t size) {
+    // We may be appending this chunk to the end of the previous.
+    if (buffer.size() == size) {
+      hState.content += std::move(buffer);
+    } else {
+      hState.content += buffer.substr(0, size);
     }
+
+    // Search for newlines and parse each.
+    size_t last_newline = 0;
+    auto newline = hState.content.find('\n');
+    while (newline != std::string::npos) {
+      auto line = hState.content.substr(last_newline, newline - last_newline);
+      parseLine(line);
+
+      last_newline = newline + 1;
+      newline = hState.content.find('\n', last_newline);
+    }
+
+    if (last_newline != hState.content.size() - 1) {
+      // We need to buffer the end of the string.
+      hState.content = hState.content.substr(last_newline);
+    }
+  };
+
+  if (!readFile(history_file, 0, 4096, false, false, parseChunk, false)) {
+    return;
+  }
+
+  // Parse the final line.
+  if (!hState.content.empty()) {
+    parseLine(hState.content);
   }
 }
 
-void genShellHistoryForUser(const std::string& uid,
+void genShellHistoryForUser(
-                            const std::string& gid,
+    const std::string& uid,
-                            const std::string& directory,
+    const std::string& gid,
-                            QueryData& results) {
+    const std::string& directory,
+    std::function<void(DynamicTableRowHolder& row)> predicate) {
   for (const auto& hfile : kShellHistoryFiles) {
     boost::filesystem::path history_file = directory;
     history_file /= hfile;
-    genShellHistoryFromFile(uid, history_file, results);
+    genShellHistoryFromFile(uid, history_file, predicate);
   }
 }
 
-void genShellHistoryFromBashSessions(const std::string& uid,
+void genShellHistoryFromBashSessions(
-                                     const std::string& directory,
+    const std::string& uid,
-                                     QueryData& results) {
+    const std::string& directory,
+    std::function<void(DynamicTableRowHolder& row)> predicate) {
   boost::filesystem::path bash_sessions = directory;
   bash_sessions /= ".bash_sessions";
 
@@ -92,13 +135,13 @@ void genShellHistoryFromBashSessions(const std::string& uid,
 
     for (const auto& hfile : session_hist_files) {
       boost::filesystem::path history_file = hfile;
-      genShellHistoryFromFile(uid, history_file, results);
+      genShellHistoryFromFile(uid, history_file, predicate);
     }
   }
 }
 
-QueryData genShellHistory(QueryContext& context) {
+void genShellHistory(RowYield& yield, QueryContext& context) {
-  QueryData results;
+  auto predicate = [&yield](DynamicTableRowHolder& r) { yield(std::move(r)); };
 
   // Iterate over each user
   QueryData users = usersFromContext(context);
@@ -107,12 +150,10 @@ QueryData genShellHistory(QueryContext& context) {
     auto gid = row.find("gid");
     auto dir = row.find("directory");
     if (uid != row.end() && gid != row.end() && dir != row.end()) {
-      genShellHistoryForUser(uid->second, gid->second, dir->second, results);
+      genShellHistoryForUser(uid->second, gid->second, dir->second, predicate);
-      genShellHistoryFromBashSessions(uid->second, dir->second, results);
+      genShellHistoryFromBashSessions(uid->second, dir->second, predicate);
     }
   }
-
-  return results;
-}
-}
 }
+} // namespace tables
+} // namespace osquery

osquery/tables/system/posix/shell_history.h

@@ -6,7 +6,7 @@
  *  the LICENSE file found in the root directory of this source tree.
  */
 
-#include <osquery/query.h>
+#include <osquery/sql/dynamic_table_row.h>
 #include <osquery/tables.h>
 
 #include <string>
@@ -14,16 +14,18 @@
 namespace osquery {
 namespace tables {
 
-void genShellHistoryFromBashSessions(const std::string& uid,
+/// This takes a predicate function to aide testing.
-                                     const std::string& directory,
+void genShellHistoryFromBashSessions(
-                                     QueryData& results);
+    const std::string& uid,
+    const std::string& directory,
+    std::function<void(DynamicTableRowHolder& row)> predicate);
 
-void genShellHistoryForUser(const std::string& uid,
+/// This takes a predicate function to aide testing.
-                            const std::string& gid,
+void genShellHistoryForUser(
-                            const std::string& directory,
+    const std::string& uid,
-                            QueryData& results);
+    const std::string& gid,
-
+    const std::string& directory,
-QueryData genShellHistory(QueryContext& context);
+    std::function<void(DynamicTableRowHolder& row)> predicate);
 
 } // namespace tables
 } // namespace osquery

osquery/tables/system/tests/posix/shell_history_tests.cpp

@@ -23,7 +23,11 @@ namespace tables {
 class ShellHistoryTests : public testing::Test {};
 
 TEST_F(ShellHistoryTests, empty_timestamp) {
-  auto results = QueryData{};
+  std::vector<DynamicTableRowHolder> results;
+  auto predicate = [&results](DynamicTableRowHolder& r) {
+    results.push_back(std::move(r));
+  };
+
   auto directory =
       fs::temp_directory_path() /
       fs::unique_path("osquery.shell_history_tests.empty_timestamp.%%%%-%%%%");
@@ -37,27 +41,32 @@ TEST_F(ShellHistoryTests, empty_timestamp) {
     fout << first_line << '\n';
     fout << second_line << '\n';
   }
+
   auto const uid = std::to_string(geteuid());
   genShellHistoryForUser(
-      uid, std::to_string(getegid()), directory.native(), results);
+      uid, std::to_string(getegid()), directory.native(), predicate);
   ASSERT_EQ(results.size(), 2u);
 
   const auto& first_row = results[0];
-  EXPECT_EQ(first_row.at("uid"), uid);
+  EXPECT_EQ(first_row["uid"], uid);
-  EXPECT_EQ(first_row.at("time"), "0");
+  EXPECT_EQ(first_row["time"], "0");
-  EXPECT_EQ(first_row.at("command"), first_line);
+  EXPECT_EQ(first_row["command"], first_line);
-  EXPECT_EQ(first_row.at("history_file"), filepath.native());
+  EXPECT_EQ(first_row["history_file"], filepath.native());
 
   const auto& second_row = results[1];
-  EXPECT_EQ(second_row.at("uid"), uid);
+  EXPECT_EQ(second_row["uid"], uid);
-  EXPECT_EQ(second_row.at("time"), "0");
+  EXPECT_EQ(second_row["time"], "0");
-  EXPECT_EQ(second_row.at("command"), second_line);
+  EXPECT_EQ(second_row["command"], second_line);
-  EXPECT_EQ(second_row.at("history_file"), filepath.native());
+  EXPECT_EQ(second_row["history_file"], filepath.native());
   fs::remove_all(directory);
 }
 
 TEST_F(ShellHistoryTests, bash_sessions_no_exist) {
-  auto results = QueryData{};
+  std::vector<DynamicTableRowHolder> results;
+  auto predicate = [&results](DynamicTableRowHolder& r) {
+    results.push_back(std::move(r));
+  };
+
   auto directory =
       fs::temp_directory_path() /
       fs::unique_path(
@@ -66,13 +75,17 @@ TEST_F(ShellHistoryTests, bash_sessions_no_exist) {
   auto const uid = std::to_string(geteuid());
 
   // test non-existent .bash_sessions directory
-  genShellHistoryFromBashSessions(uid, directory.native(), results);
+  genShellHistoryFromBashSessions(uid, directory.native(), predicate);
   ASSERT_EQ(results.size(), 0u);
   fs::remove_all(directory);
 }
 
 TEST_F(ShellHistoryTests, bash_sessions_no_history) {
-  auto results = QueryData{};
+  std::vector<DynamicTableRowHolder> results;
+  auto predicate = [&results](DynamicTableRowHolder& r) {
+    results.push_back(std::move(r));
+  };
+
   auto directory =
       fs::temp_directory_path() /
       fs::unique_path(
@@ -92,13 +105,17 @@ TEST_F(ShellHistoryTests, bash_sessions_no_history) {
   }
   auto const uid = std::to_string(geteuid());
   // test non-existent some_guid_here.history file
-  genShellHistoryFromBashSessions(uid, directory.native(), results);
+  genShellHistoryFromBashSessions(uid, directory.native(), predicate);
   ASSERT_EQ(results.size(), 0u);
   fs::remove_all(directory);
 }
 
 TEST_F(ShellHistoryTests, bash_sessions_empty_ts) {
-  auto results = QueryData{};
+  std::vector<DynamicTableRowHolder> results;
+  auto predicate = [&results](DynamicTableRowHolder& r) {
+    results.push_back(std::move(r));
+  };
+
   auto directory =
       fs::temp_directory_path() /
       fs::unique_path(
@@ -118,20 +135,20 @@ TEST_F(ShellHistoryTests, bash_sessions_empty_ts) {
     fout << second_line << '\n';
   }
   auto const uid = std::to_string(geteuid());
-  genShellHistoryFromBashSessions(uid, directory.native(), results);
+  genShellHistoryFromBashSessions(uid, directory.native(), predicate);
   ASSERT_EQ(results.size(), 2u);
 
   const auto& first_row = results[0];
-  EXPECT_EQ(first_row.at("uid"), uid);
+  EXPECT_EQ(first_row["uid"], uid);
-  EXPECT_EQ(first_row.at("time"), "0");
+  EXPECT_EQ(first_row["time"], "0");
-  EXPECT_EQ(first_row.at("command"), first_line);
+  EXPECT_EQ(first_row["command"], first_line);
-  EXPECT_EQ(first_row.at("history_file"), fs::canonical(filepath).native());
+  EXPECT_EQ(first_row["history_file"], fs::canonical(filepath).native());
 
   const auto& second_row = results[1];
-  EXPECT_EQ(second_row.at("uid"), uid);
+  EXPECT_EQ(second_row["uid"], uid);
-  EXPECT_EQ(second_row.at("time"), "0");
+  EXPECT_EQ(second_row["time"], "0");
-  EXPECT_EQ(second_row.at("command"), second_line);
+  EXPECT_EQ(second_row["command"], second_line);
-  EXPECT_EQ(second_row.at("history_file"), fs::canonical(filepath).native());
+  EXPECT_EQ(second_row["history_file"], fs::canonical(filepath).native());
   fs::remove_all(directory);
 }
 

specs/posix/shell_history.table

@@ -8,7 +8,7 @@ schema([
     ForeignKey(column="uid", table="users"),
 ])
 attributes(user_data=True, no_pkey=True)
-implementation("shell_history@genShellHistory")
+implementation("shell_history@genShellHistory", generator=True)
 examples([
     "select * from users join shell_history using (uid)",
 ])