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.
This commit is contained in:
vieira 2016-10-23 23:12:12 +00:00 committed by Brian May
parent 15b394da86
commit fbbcc05d58
3 changed files with 44 additions and 0 deletions

View File

@ -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:

View File

@ -221,6 +221,7 @@ def main(method_name, syslog):
break
finally:
try:
sdnotify.send(sdnotify.stop())
debug1('firewall manager: undoing changes.\n')
except:
pass

39
sshuttle/sdnotify.py Normal file
View File

@ -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')