Merge pull request #771 from theopolis/extension_query

Add thrift 'query' endpoint

include/osquery/extensions.h

@@ -138,6 +138,8 @@ class ExtensionManagerHandler : virtual public ExtensionManagerIf,
   void deregisterExtension(ExtensionStatus& _return,
                            const ExtensionRouteUUID uuid);
 
+  void query(ExtensionResponse& _return, const std::string& sql);
+
  private:
   /// Check if an extension exists by the name it registered.
   bool exists(const std::string& name);

osquery.thrift

@@ -4,16 +4,22 @@ namespace cpp osquery.extensions
 typedef map<string, string> ExtensionPluginRequest
 typedef list<map<string, string>> ExtensionPluginResponse
 
+/// When communicating extension metadata, use a thrift-internal structure.
 struct InternalExtensionInfo {
   1:string name,
   2:string version,
   3:string sdk_version,
 }
 
+/// Unique ID for each extension.
 typedef i64 ExtensionRouteUUID
+/// An optional key/value pair provided by extension plugins.
 typedef map<string, string> ExtensionRoute
+/// A map from each plugin name to its optional route information.
 typedef map<string, ExtensionRoute> ExtensionRouteTable
+/// A map from each registry name.
 typedef map<string, ExtensionRouteTable> ExtensionRegistry
+/// A map from each extension's unique ID to its map of registries.
 typedef map<ExtensionRouteUUID, InternalExtensionInfo> InternalExtensionList
 
 enum ExtensionCode {
@@ -26,6 +32,7 @@ enum ExtensionCode {
 struct ExtensionStatus {
   1:i32 code,
   2:string message,
+  /// Add a thrift Status parameter identifying the request/response.
   3:ExtensionRouteUUID uuid,
 }
 
@@ -41,19 +48,31 @@ exception ExtensionException {
 }
 
 service Extension {
+  /// Ping to/from an extension and extension manager for metadata.
   ExtensionStatus ping(),
+  /// Call an extension (or core) registry plugin.
   ExtensionResponse call(
+    /// The registry name (e.g., config, logger, table, etc).
     1:string registry,
+    /// The registry item name (plugin name).
     2:string item,
+    /// The thrift-equivilent of an osquery::PluginRequest.
     3:ExtensionPluginRequest request),
 }
 
+/// The extension manager is run by the osquery core process.
 service ExtensionManager extends Extension {
+  /// Return the list of active registered extensions.
   InternalExtensionList extensions(),
+  /// The API endpoint used by an extension to register its plugins.
   ExtensionStatus registerExtension(
     1:InternalExtensionInfo info,
     2:ExtensionRegistry registry),
   ExtensionStatus deregisterExtension(
     1:ExtensionRouteUUID uuid,
   ),
+  /// Allow an extension to query using an SQL string.
+  ExtensionResponse query(
+    1:string sql,
+  ),
 }

osquery/extensions/extensions.cpp

@@ -18,6 +18,7 @@
 #include <osquery/extensions.h>
 #include <osquery/filesystem.h>
 #include <osquery/logger.h>
+#include <osquery/sql.h>
 
 using namespace apache::thrift;
 using namespace apache::thrift::protocol;
@@ -111,6 +112,20 @@ void ExtensionManagerHandler::deregisterExtension(
   extensions_.erase(uuid);
 }
 
+void ExtensionManagerHandler::query(ExtensionResponse& _return,
+                                    const std::string& sql) {
+  QueryData results;
+  auto status = osquery::query(sql, results);
+  _return.status.code = status.getCode();
+  _return.status.message = status.getMessage();
+
+  if (status.ok()) {
+    for (const auto& row : results) {
+      _return.response.push_back(row);
+    }
+  }
+}
+
 bool ExtensionManagerHandler::exists(const std::string& name) {
   for (const auto& extension : extensions_) {
     if (extension.second.name == name) {

osquery/extensions/extensions_tests.cpp

@@ -19,6 +19,7 @@
 
 #include <osquery/extensions.h>
 #include <osquery/filesystem.h>
+#include <osquery/database.h>
 
 using namespace apache::thrift;
 using namespace apache::thrift::protocol;
@@ -71,6 +72,34 @@ class ExtensionsTest : public testing::Test {
     return false;
   }
 
+  QueryData query(const std::string& sql, int attempts = 3) {
+    // Open a socket to the test extension manager.
+    boost::shared_ptr<TSocket> socket(new TSocket(kTestManagerSocket));
+    boost::shared_ptr<TTransport> transport(new TBufferedTransport(socket));
+    boost::shared_ptr<TProtocol> protocol(new TBinaryProtocol(transport));
+
+    ExtensionManagerClient client(protocol);
+
+    // Calling open will except if the socket does not exist.
+    ExtensionResponse response;
+    for (int i = 0; i < attempts; ++i) {
+      try {
+        transport->open();
+        client.query(response, sql);
+        transport->close();
+      } catch (const std::exception& e) {
+        ::usleep(kDelayUS);
+      }
+    }
+
+    QueryData qd;
+    for (const auto& row : response.response) {
+      qd.push_back(row);
+    }
+
+    return qd;
+  }
+
   ExtensionList registeredExtensions(int attempts = 3) {
     ExtensionList extensions;
     for (int i = 0; i < attempts; ++i) {
@@ -156,6 +185,17 @@ TEST_F(ExtensionsTest, test_extension_start) {
   Registry::allowDuplicates(false);
 }
 
+TEST_F(ExtensionsTest, test_extension_query) {
+  auto status = startExtensionManager(kTestManagerSocket);
+  EXPECT_TRUE(status.ok());
+  // Wait for the extension manager to start.
+  EXPECT_TRUE(socketExists(kTestManagerSocket));
+
+  auto qd = query("select seconds from time");
+  EXPECT_EQ(qd.size(), 1);
+  EXPECT_EQ(qd[0].count("seconds"), 1);
+}
+
 class ExtensionPlugin : public Plugin {
  public:
   Status call(const PluginRequest& request, PluginResponse& response) {