496 lines
12 KiB
Plaintext
496 lines
12 KiB
Plaintext
|
#!/opt/python/bin/python3
|
|||
|
# vim:fileencoding=utf-8:noet
|
|||
|
from __future__ import (unicode_literals, division, absolute_import, print_function)
|
|||
|
|
|||
|
import socket
|
|||
|
import os
|
|||
|
import errno
|
|||
|
import sys
|
|||
|
import fcntl
|
|||
|
import atexit
|
|||
|
import stat
|
|||
|
|
|||
|
from argparse import ArgumentParser
|
|||
|
from select import select
|
|||
|
from signal import signal, SIGTERM
|
|||
|
from time import sleep
|
|||
|
from functools import partial
|
|||
|
from io import BytesIO
|
|||
|
from threading import Event
|
|||
|
from itertools import chain
|
|||
|
from logging import StreamHandler
|
|||
|
|
|||
|
from powerline.shell import ShellPowerline
|
|||
|
from powerline.commands.main import finish_args, write_output
|
|||
|
from powerline.lib.monotonic import monotonic
|
|||
|
from powerline.lib.encoding import get_preferred_output_encoding, get_preferred_arguments_encoding, get_unicode_writer
|
|||
|
from powerline.bindings.wm import wm_threads
|
|||
|
|
|||
|
from powerline.commands.main import get_argparser as get_main_argparser
|
|||
|
from powerline.commands.daemon import get_argparser as get_daemon_argparser
|
|||
|
|
|||
|
|
|||
|
USE_FILESYSTEM = not sys.platform.lower().startswith('linux')
|
|||
|
|
|||
|
|
|||
|
class NonInteractiveArgParser(ArgumentParser):
|
|||
|
def print_usage(self, file=None):
|
|||
|
raise Exception(self.format_usage())
|
|||
|
|
|||
|
def print_help(self, file=None):
|
|||
|
raise Exception(self.format_help())
|
|||
|
|
|||
|
def exit(self, status=0, message=None):
|
|||
|
pass
|
|||
|
|
|||
|
def error(self, message):
|
|||
|
raise Exception(self.format_usage())
|
|||
|
|
|||
|
|
|||
|
EOF = b'EOF\0\0'
|
|||
|
|
|||
|
|
|||
|
class State(object):
|
|||
|
__slots__ = ('powerlines', 'logger', 'config_loader', 'started_wm_threads',
|
|||
|
'ts_shutdown_event')
|
|||
|
|
|||
|
def __init__(self, **kwargs):
|
|||
|
self.logger = None
|
|||
|
self.config_loader = None
|
|||
|
self.started_wm_threads = {}
|
|||
|
self.powerlines = {}
|
|||
|
self.ts_shutdown_event = Event()
|
|||
|
|
|||
|
|
|||
|
HOME = os.path.expanduser('~')
|
|||
|
|
|||
|
|
|||
|
class NonDaemonShellPowerline(ShellPowerline):
|
|||
|
def get_log_handler(self):
|
|||
|
return StreamHandler()
|
|||
|
|
|||
|
|
|||
|
def start_wm(args, environ, cwd, is_daemon, state):
|
|||
|
wm_name = args.ext[0][3:]
|
|||
|
if wm_name in state.started_wm_threads:
|
|||
|
return b''
|
|||
|
thread_shutdown_event = Event()
|
|||
|
thread = wm_threads[wm_name](
|
|||
|
thread_shutdown_event=thread_shutdown_event,
|
|||
|
pl_shutdown_event=state.ts_shutdown_event,
|
|||
|
pl_config_loader=state.config_loader,
|
|||
|
)
|
|||
|
thread.start()
|
|||
|
state.started_wm_threads[wm_name] = (thread, thread_shutdown_event)
|
|||
|
return b''
|
|||
|
|
|||
|
|
|||
|
def render(args, environ, cwd, is_daemon, state):
|
|||
|
segment_info = {
|
|||
|
'getcwd': lambda: cwd,
|
|||
|
'home': environ.get('HOME', HOME),
|
|||
|
'environ': environ,
|
|||
|
'args': args,
|
|||
|
}
|
|||
|
key = (
|
|||
|
args.ext[0],
|
|||
|
args.renderer_module,
|
|||
|
tuple(args.config_override) if args.config_override else None,
|
|||
|
tuple(args.theme_override) if args.theme_override else None,
|
|||
|
tuple(args.config_path) if args.config_path else None,
|
|||
|
environ.get('POWERLINE_THEME_OVERRIDES', ''),
|
|||
|
environ.get('POWERLINE_CONFIG_OVERRIDES', ''),
|
|||
|
environ.get('POWERLINE_CONFIG_PATHS', ''),
|
|||
|
)
|
|||
|
|
|||
|
PowerlineClass = ShellPowerline if is_daemon else NonDaemonShellPowerline
|
|||
|
powerline = None
|
|||
|
try:
|
|||
|
powerline = state.powerlines[key]
|
|||
|
except KeyError:
|
|||
|
try:
|
|||
|
powerline = state.powerlines[key] = PowerlineClass(
|
|||
|
args,
|
|||
|
logger=state.logger,
|
|||
|
config_loader=state.config_loader,
|
|||
|
run_once=False,
|
|||
|
shutdown_event=state.ts_shutdown_event,
|
|||
|
)
|
|||
|
if state.logger is None:
|
|||
|
state.logger = powerline.logger
|
|||
|
if state.config_loader is None:
|
|||
|
state.config_loader = powerline.config_loader
|
|||
|
except SystemExit:
|
|||
|
# Somebody thought raising system exit was a good idea,
|
|||
|
return ''
|
|||
|
except Exception as e:
|
|||
|
if powerline:
|
|||
|
powerline.pl.exception('Failed to render {0}: {1}', str(key), str(e))
|
|||
|
else:
|
|||
|
return 'Failed to render {0}: {1}'.format(str(key), str(e))
|
|||
|
s = BytesIO()
|
|||
|
write_output(args, powerline, segment_info, get_unicode_writer(stream=s))
|
|||
|
s.seek(0)
|
|||
|
return s.read()
|
|||
|
|
|||
|
|
|||
|
def eintr_retry_call(func, *args, **kwargs):
|
|||
|
while True:
|
|||
|
try:
|
|||
|
return func(*args, **kwargs)
|
|||
|
except EnvironmentError as e:
|
|||
|
if getattr(e, 'errno', None) == errno.EINTR:
|
|||
|
continue
|
|||
|
raise
|
|||
|
|
|||
|
|
|||
|
def do_read(conn, timeout=2.0):
|
|||
|
''' Read data from the client. If the client fails to send data within
|
|||
|
timeout seconds, abort. '''
|
|||
|
read = []
|
|||
|
end_time = monotonic() + timeout
|
|||
|
while not read or not read[-1].endswith(b'\0\0'):
|
|||
|
r, w, e = select((conn,), (), (conn,), timeout)
|
|||
|
if e:
|
|||
|
return
|
|||
|
if monotonic() > end_time:
|
|||
|
return
|
|||
|
if not r:
|
|||
|
continue
|
|||
|
x = eintr_retry_call(conn.recv, 4096)
|
|||
|
if x:
|
|||
|
read.append(x)
|
|||
|
else:
|
|||
|
break
|
|||
|
return b''.join(read)
|
|||
|
|
|||
|
|
|||
|
def do_write(conn, result):
|
|||
|
try:
|
|||
|
eintr_retry_call(conn.sendall, result)
|
|||
|
except Exception:
|
|||
|
pass
|
|||
|
|
|||
|
|
|||
|
def safe_bytes(o, encoding=get_preferred_output_encoding()):
|
|||
|
'''Return bytes instance without ever throwing an exception.'''
|
|||
|
try:
|
|||
|
try:
|
|||
|
# We are assuming that o is a unicode object
|
|||
|
return o.encode(encoding, 'replace')
|
|||
|
except Exception:
|
|||
|
# Object may have defined __bytes__ (python 3) or __str__ method
|
|||
|
# (python 2)
|
|||
|
# This also catches problem with non_ascii_bytes.encode('utf-8')
|
|||
|
# that first tries to decode to UTF-8 using ascii codec (and fails
|
|||
|
# in this case) and then encode to given encoding: errors= argument
|
|||
|
# is not used in the first stage.
|
|||
|
return bytes(o)
|
|||
|
except Exception as e:
|
|||
|
return safe_bytes(str(e), encoding)
|
|||
|
|
|||
|
|
|||
|
def parse_args(req, parser, encoding=get_preferred_arguments_encoding()):
|
|||
|
args = [x.decode(encoding) for x in req.split(b'\0') if x]
|
|||
|
numargs = int(args[0], 16)
|
|||
|
shell_args = parser.parse_args(args[1:numargs + 1])
|
|||
|
cwd = args[numargs + 1]
|
|||
|
environ = dict(((k, v) for k, v in (x.partition('=')[0::2] for x in args[numargs + 2:])))
|
|||
|
cwd = cwd or environ.get('PWD', '/')
|
|||
|
return shell_args, environ, cwd
|
|||
|
|
|||
|
|
|||
|
def get_answer(req, is_daemon, argparser, state):
|
|||
|
try:
|
|||
|
args, environ, cwd = parse_args(req, argparser)
|
|||
|
finish_args(argparser, environ, args, is_daemon=True)
|
|||
|
if args.ext[0].startswith('wm.'):
|
|||
|
return safe_bytes(start_wm(args, environ, cwd, is_daemon, state))
|
|||
|
else:
|
|||
|
return safe_bytes(render(args, environ, cwd, is_daemon, state))
|
|||
|
except Exception as e:
|
|||
|
return safe_bytes(str(e))
|
|||
|
|
|||
|
|
|||
|
def do_one(sock, read_sockets, write_sockets, result_map, is_daemon, argparser,
|
|||
|
state):
|
|||
|
r, w, e = select(
|
|||
|
tuple(read_sockets) + (sock,),
|
|||
|
tuple(write_sockets),
|
|||
|
tuple(read_sockets) + tuple(write_sockets) + (sock,),
|
|||
|
60.0
|
|||
|
)
|
|||
|
|
|||
|
if sock in e:
|
|||
|
# We cannot accept any more connections, so we exit
|
|||
|
raise SystemExit(1)
|
|||
|
|
|||
|
for s in e:
|
|||
|
# Discard all broken connections to clients
|
|||
|
s.close()
|
|||
|
read_sockets.discard(s)
|
|||
|
write_sockets.discard(s)
|
|||
|
|
|||
|
for s in r:
|
|||
|
if s == sock:
|
|||
|
# A client wants to connect
|
|||
|
conn, _ = eintr_retry_call(sock.accept)
|
|||
|
read_sockets.add(conn)
|
|||
|
else:
|
|||
|
# A client has sent some data
|
|||
|
read_sockets.discard(s)
|
|||
|
req = do_read(s)
|
|||
|
if req == EOF:
|
|||
|
raise SystemExit(0)
|
|||
|
elif req:
|
|||
|
ans = get_answer(req, is_daemon, argparser, state)
|
|||
|
result_map[s] = ans
|
|||
|
write_sockets.add(s)
|
|||
|
else:
|
|||
|
s.close()
|
|||
|
|
|||
|
for s in w:
|
|||
|
# A client is ready to receive the result
|
|||
|
write_sockets.discard(s)
|
|||
|
result = result_map.pop(s)
|
|||
|
try:
|
|||
|
do_write(s, result)
|
|||
|
finally:
|
|||
|
s.close()
|
|||
|
|
|||
|
|
|||
|
def shutdown(sock, read_sockets, write_sockets, state):
|
|||
|
'''Perform operations necessary for nicely shutting down daemon
|
|||
|
|
|||
|
Specifically it
|
|||
|
|
|||
|
#. Closes all sockets.
|
|||
|
#. Notifies segments based on
|
|||
|
:py:class:`powerline.lib.threaded.ThreadedSegment` and WM-specific
|
|||
|
threads that daemon is shutting down.
|
|||
|
#. Waits for threads to finish, but no more then 2 seconds total.
|
|||
|
#. Waits so that total execution time of this function is 2 seconds in order
|
|||
|
to allow ThreadedSegments to finish.
|
|||
|
'''
|
|||
|
total_wait_time = 2
|
|||
|
shutdown_start_time = monotonic()
|
|||
|
|
|||
|
for s in chain((sock,), read_sockets, write_sockets):
|
|||
|
s.close()
|
|||
|
|
|||
|
# Notify ThreadedSegments
|
|||
|
state.ts_shutdown_event.set()
|
|||
|
for thread, shutdown_event in state.started_wm_threads.values():
|
|||
|
shutdown_event.set()
|
|||
|
|
|||
|
for thread, shutdown_event in state.started_wm_threads.values():
|
|||
|
wait_time = total_wait_time - (monotonic() - shutdown_start_time)
|
|||
|
if wait_time > 0:
|
|||
|
thread.join(wait_time)
|
|||
|
|
|||
|
wait_time = total_wait_time - (monotonic() - shutdown_start_time)
|
|||
|
sleep(wait_time)
|
|||
|
|
|||
|
|
|||
|
def main_loop(sock, is_daemon):
|
|||
|
sock.listen(128)
|
|||
|
sock.setblocking(0)
|
|||
|
|
|||
|
read_sockets, write_sockets = set(), set()
|
|||
|
result_map = {}
|
|||
|
parser = get_main_argparser(NonInteractiveArgParser)
|
|||
|
state = State()
|
|||
|
try:
|
|||
|
try:
|
|||
|
while True:
|
|||
|
do_one(
|
|||
|
sock, read_sockets, write_sockets, result_map,
|
|||
|
is_daemon=is_daemon,
|
|||
|
argparser=parser,
|
|||
|
state=state,
|
|||
|
)
|
|||
|
except KeyboardInterrupt:
|
|||
|
raise SystemExit(0)
|
|||
|
except SystemExit as e:
|
|||
|
shutdown(sock, read_sockets, write_sockets, state)
|
|||
|
raise e
|
|||
|
return 0
|
|||
|
|
|||
|
|
|||
|
def daemonize(stdin=os.devnull, stdout=os.devnull, stderr=os.devnull):
|
|||
|
try:
|
|||
|
pid = os.fork()
|
|||
|
if pid > 0:
|
|||
|
# exit first parent
|
|||
|
raise SystemExit(0)
|
|||
|
except OSError as e:
|
|||
|
sys.stderr.write("fork #1 failed: %d (%s)\n" % (e.errno, e.strerror))
|
|||
|
raise SystemExit(1)
|
|||
|
|
|||
|
# decouple from parent environment
|
|||
|
os.chdir("/")
|
|||
|
os.setsid()
|
|||
|
os.umask(0)
|
|||
|
|
|||
|
# do second fork
|
|||
|
try:
|
|||
|
pid = os.fork()
|
|||
|
if pid > 0:
|
|||
|
# exit from second parent
|
|||
|
raise SystemExit(0)
|
|||
|
except OSError as e:
|
|||
|
sys.stderr.write("fork #2 failed: %d (%s)\n" % (e.errno, e.strerror))
|
|||
|
raise SystemExit(1)
|
|||
|
|
|||
|
# Redirect standard file descriptors.
|
|||
|
si = open(stdin, 'rb')
|
|||
|
so = open(stdout, 'a+b')
|
|||
|
se = open(stderr, 'a+b', 0)
|
|||
|
os.dup2(si.fileno(), sys.stdin.fileno())
|
|||
|
os.dup2(so.fileno(), sys.stdout.fileno())
|
|||
|
os.dup2(se.fileno(), sys.stderr.fileno())
|
|||
|
return True
|
|||
|
|
|||
|
|
|||
|
def check_existing(address):
|
|||
|
if USE_FILESYSTEM:
|
|||
|
# We cannot bind if the socket file already exists so remove it, we
|
|||
|
# already have a lock on pidfile, so this should be safe.
|
|||
|
try:
|
|||
|
os.unlink(address)
|
|||
|
except EnvironmentError:
|
|||
|
pass
|
|||
|
|
|||
|
sock = socket.socket(family=socket.AF_UNIX)
|
|||
|
try:
|
|||
|
sock.bind(address)
|
|||
|
except socket.error as e:
|
|||
|
if getattr(e, 'errno', None) == errno.EADDRINUSE:
|
|||
|
return None
|
|||
|
raise
|
|||
|
return sock
|
|||
|
|
|||
|
|
|||
|
def kill_daemon(address):
|
|||
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|||
|
try:
|
|||
|
try:
|
|||
|
eintr_retry_call(sock.connect, address)
|
|||
|
except socket.error:
|
|||
|
return False
|
|||
|
else:
|
|||
|
eintr_retry_call(sock.sendall, EOF)
|
|||
|
finally:
|
|||
|
sock.close()
|
|||
|
return True
|
|||
|
|
|||
|
|
|||
|
def cleanup_lockfile(pidfile, fd, *args):
|
|||
|
try:
|
|||
|
# Remove the directory entry for the lock file
|
|||
|
os.unlink(pidfile)
|
|||
|
# Close the file descriptor
|
|||
|
os.close(fd)
|
|||
|
except EnvironmentError:
|
|||
|
pass
|
|||
|
if args:
|
|||
|
# Called in signal handler
|
|||
|
raise SystemExit(1)
|
|||
|
|
|||
|
|
|||
|
def lockpidfile(pidfile):
|
|||
|
fd = os.open(
|
|||
|
pidfile,
|
|||
|
os.O_WRONLY | os.O_CREAT,
|
|||
|
stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH
|
|||
|
)
|
|||
|
try:
|
|||
|
fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|||
|
except EnvironmentError:
|
|||
|
os.close(fd)
|
|||
|
return None
|
|||
|
os.lseek(fd, 0, os.SEEK_SET)
|
|||
|
os.ftruncate(fd, 0)
|
|||
|
os.write(fd, ('%d' % os.getpid()).encode('ascii'))
|
|||
|
os.fsync(fd)
|
|||
|
cleanup = partial(cleanup_lockfile, pidfile, fd)
|
|||
|
signal(SIGTERM, cleanup)
|
|||
|
atexit.register(cleanup)
|
|||
|
return fd
|
|||
|
|
|||
|
|
|||
|
def main():
|
|||
|
parser = get_daemon_argparser()
|
|||
|
args = parser.parse_args()
|
|||
|
is_daemon = False
|
|||
|
address = None
|
|||
|
pidfile = None
|
|||
|
|
|||
|
if args.socket:
|
|||
|
address = args.socket
|
|||
|
if not USE_FILESYSTEM:
|
|||
|
address = '\0' + address
|
|||
|
else:
|
|||
|
if USE_FILESYSTEM:
|
|||
|
address = '/tmp/powerline-ipc-%d'
|
|||
|
else:
|
|||
|
# Use the abstract namespace for sockets rather than the filesystem
|
|||
|
# (Available only in linux)
|
|||
|
address = '\0powerline-ipc-%d'
|
|||
|
|
|||
|
address = address % os.getuid()
|
|||
|
|
|||
|
if USE_FILESYSTEM:
|
|||
|
pidfile = address + '.pid'
|
|||
|
|
|||
|
if args.kill:
|
|||
|
if args.foreground or args.replace:
|
|||
|
parser.error('--kill and --foreground/--replace cannot be used together')
|
|||
|
if kill_daemon(address):
|
|||
|
if not args.quiet:
|
|||
|
print ('Kill command sent to daemon, if it does not die in a couple of seconds use kill to kill it')
|
|||
|
raise SystemExit(0)
|
|||
|
else:
|
|||
|
if not args.quiet:
|
|||
|
print ('No running daemon found')
|
|||
|
raise SystemExit(1)
|
|||
|
|
|||
|
if args.replace:
|
|||
|
while kill_daemon(address):
|
|||
|
if not args.quiet:
|
|||
|
print ('Kill command sent to daemon, waiting for daemon to exit, press Ctrl-C to terminate wait and exit')
|
|||
|
sleep(2)
|
|||
|
|
|||
|
if USE_FILESYSTEM and not args.foreground:
|
|||
|
# We must daemonize before creating the locked pidfile, unfortunately,
|
|||
|
# this means further print statements are discarded
|
|||
|
is_daemon = daemonize()
|
|||
|
|
|||
|
if USE_FILESYSTEM:
|
|||
|
# Create a locked pid file containing the daemon’s PID
|
|||
|
if lockpidfile(pidfile) is None:
|
|||
|
if not args.quiet:
|
|||
|
sys.stderr.write(
|
|||
|
'The daemon is already running. Use %s -k to kill it.\n' % (
|
|||
|
os.path.basename(sys.argv[0])))
|
|||
|
raise SystemExit(1)
|
|||
|
|
|||
|
# Bind to address or bail if we cannot bind
|
|||
|
sock = check_existing(address)
|
|||
|
if sock is None:
|
|||
|
if not args.quiet:
|
|||
|
sys.stderr.write(
|
|||
|
'The daemon is already running. Use %s -k to kill it.\n' % (
|
|||
|
os.path.basename(sys.argv[0])))
|
|||
|
raise SystemExit(1)
|
|||
|
|
|||
|
if not USE_FILESYSTEM and not args.foreground:
|
|||
|
# We daemonize on linux
|
|||
|
is_daemon = daemonize()
|
|||
|
|
|||
|
return main_loop(sock, is_daemon)
|
|||
|
|
|||
|
|
|||
|
if __name__ == '__main__':
|
|||
|
main()
|