extensions: Add watchdog thread and open retries (#31)

Author: Teddy Reed

osquery/__init__.py

@@ -4,7 +4,7 @@
 """
 
 __title__ = "osquery"
-__version__ = "1.8.0"
+__version__ = "2.2.0"
 __author__ = "osquery developers"
 __license__ = "BSD"
 __copyright__ = "Copyright 2015 Facebook"

osquery/extension_client.py

@@ -8,6 +8,8 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import time
+
 from thrift.protocol import TBinaryProtocol
 from thrift.transport import TSocket
 from thrift.transport import TTransport
@@ -41,9 +43,18 @@ def close(self):
         if self._transport:
             self._transport.close()
 
-    def open(self):
+    def open(self, timeout=1, interval=0.2):
         """Attempt to open the UNIX domain socket"""
-        self._transport.open()
+        delay = 0
+        while delay < timeout:
+            try:
+                self._transport.open()
+                return True
+            except TTransport.TTransportException:
+                pass
+            delay += interval
+            time.sleep(interval)
+        return False
 
     def extension_manager_client(self):
         """Return an extension manager (osquery core) client."""

osquery/extension_manager.py

@@ -8,12 +8,19 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import sys
+import os
+import threading
+import time
 
 from osquery.extensions.Extension import Iface
 from osquery.extensions.ttypes import ExtensionResponse, ExtensionStatus
 from osquery.singleton import Singleton
 
+def shutdown_request(code, timeout=1):
+    """A delayed shutdown request."""
+    time.sleep(timeout)
+    os._exit(code)
+
 class ExtensionManager(Singleton, Iface):
     """The thrift server for handling extension requests
 
@@ -64,7 +71,10 @@ def add_plugin(self, plugin):
 
     def shutdown(self):
         """The osquery extension manager requested a shutdown"""
-        sys.exit(0)
+        rt = threading.Thread(target=shutdown_request, args=(0,))
+        rt.daemon = True
+        rt.start()
+        return ExtensionStatus(code=0, message="OK")
 
     def registry(self):
         """Accessor for the internal _registry member variable"""

osquery/management.py

@@ -9,11 +9,13 @@
 from __future__ import unicode_literals
 
 import argparse
+import os
 import shutil
 import socket
 import subprocess
 import sys
 import tempfile
+import threading
 import time
 
 from thrift.protocol import TBinaryProtocol
@@ -144,6 +146,19 @@ def parse_cli_params():
         help="Enable verbose informational messages")
     return parser.parse_args()
 
+def start_watcher(client, interval):
+    """Ping the osquery extension manager to detect dirty shutdowns."""
+    try:
+        while True:
+            status = client.extension_manager_client().ping()
+            if status.code is not 0:
+                break
+            time.sleep(interval)
+    except socket.error:
+        # The socket was torn down.
+        pass
+    os._exit(0)
+
 def start_extension(name="<unknown>", version="0.0.0", sdk_version="1.8.0",
                     min_sdk_version="1.8.0"):
     """Start your extension by communicating with osquery core and starting
@@ -157,7 +172,8 @@ def start_extension(name="<unknown>", version="0.0.0", sdk_version="1.8.0",
     """
     args = parse_cli_params()
     client = ExtensionClient(path=args.socket)
-    client.open()
+    if not client.open(args.timeout):
+        return
     ext_manager = ExtensionManager()
 
     # try connecting to the desired osquery core extension manager socket
@@ -184,6 +200,11 @@ def start_extension(name="<unknown>", version="0.0.0", sdk_version="1.8.0",
             message=status.message,
         )
 
+    # Start a watchdog thread to monitor the osquery process.
+    rt = threading.Thread(target=start_watcher, args=(client, args.interval))
+    rt.daemon = True
+    rt.start()
+
     # start a thrift server listening at the path dictated by the uuid returned
     # by the osquery core extension manager
     ext_manager.uuid = status.uuid