Support sdnotify for better systemd integration

These changes introduce support for sdnotify allowing sshuttle to notify
systemd when it finishes connecting to the server and installing
firewall rules, and is ready to tunnel requests.

3 files changed, 44 insertions, 0 deletions

diff --git a/sshuttle/client.py b/sshuttle/client.py
index 0711912..45788d2 100644
--- a/sshuttle/client.py
+++ b/sshuttle/client.py
@@ -9,6 +9,7 @@ import os
 import sshuttle.ssnet as ssnet
 import sshuttle.ssh as ssh
 import sshuttle.ssyslog as ssyslog
+import sshuttle.sdnotify as sdnotify
 import sys
 import platform
 from sshuttle.ssnet import SockWrapper, Handler, Proxy, Mux, MuxWrapper
@@ -503,6 +504,9 @@ def _main(tcp_listener, udp_listener, fw, ssh_cmd, remotename,
         debug1('seed_hosts: %r\n' % seed_hosts)
         mux.send(0, ssnet.CMD_HOST_REQ, str.encode('\n'.join(seed_hosts)))
 
+    sdnotify.send(sdnotify.ready(),
+                  sdnotify.status('Connected to %s.' % remotename))
+
     while 1:
         rv = serverproc.poll()
         if rv:
diff --git a/sshuttle/firewall.py b/sshuttle/firewall.py
index f16aac6..0e30e9c 100644
--- a/sshuttle/firewall.py
+++ b/sshuttle/firewall.py
@@ -221,6 +221,7 @@ def main(method_name, syslog):
                 break
     finally:
         try:
+            sdnotify.send(sdnotify.stop())
             debug1('firewall manager: undoing changes.\n')
         except:
             pass
diff --git a/sshuttle/sdnotify.py b/sshuttle/sdnotify.py
new file mode 100644
index 0000000..665d953
--- /dev/null
+++ b/sshuttle/sdnotify.py
@@ -0,0 +1,39 @@
+import socket
+import os
+from sshuttle.helpers import debug1, debug2, debug3
+
+def _notify(message):
+    addr = os.environ.get("NOTIFY_SOCKET", None)
+
+    if not addr or len(addr) == 1 or addr[0] not in ('/', '@'):
+        return False
+
+    addr = '\0' + addr[1:] if addr[0] == '@' else addr
+    
+    try:
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+    except socket.error as e:
+        debug1("Error creating socket to notify to systemd: %s\n" % e)
+
+    if not (sock and message): 
+        return False
+
+    assert isinstance(message, bytes)
+
+    try:
+        return (sock.sendto(message, addr) > 0)
+    except socket.error as e:
+        debug1("Error notifying systemd: %s\n" % e)
+        return False
+
+def send(*messages):
+    return _notify(b'\n'.join(messages))
+
+def ready():
+    return b"READY=1"
+
+def stop():
+    return b"STOPPING=1"
+
+def status(message):
+    return b"STATUS=%s" % message.encode('utf8')