496 lines
12 KiB
Python
496 lines
12 KiB
Python
#!/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()
|