1 files changed, 495 insertions, 0 deletions
diff --git a/scripts/powerline-daemon b/scripts/powerline-daemon
new file mode 100755
index 0000000..f15ac69
--- /dev/null
+++ b/scripts/powerline-daemon
@@ -0,0 +1,495 @@
+#!/usr/bin/env python
+# 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()