powerline/scripts/powerline-daemon

#!/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

from argparse import ArgumentParser
from select import select
from signal import signal, SIGTERM
from time import sleep
from functools import partial
from locale import getpreferredencoding
from io import BytesIO

from powerline.shell import get_argparser, finish_args, ShellPowerline, write_output
from powerline.lib.monotonic import monotonic


is_daemon = False
platform = sys.platform.lower()
use_filesystem = 'darwin' in platform

address = None
pidfile = None


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())


parser = get_argparser(parser=NonInteractiveArgParser, description='powerline daemon')

EOF = b'EOF\0\0'

powerlines = {}
logger = None
config_loader = None
home = os.path.expanduser('~')


class PowerlineDaemon(ShellPowerline):
	def get_log_handler(self):
		if not is_daemon:
			import logging
			return logging.StreamHandler()
		return super(PowerlineDaemon, self).get_log_handler()


def render(args, environ, cwd):
	global logger
	global config_loader
	cwd = cwd or environ.get('PWD', '/')
	segment_info = {
		'getcwd': lambda: cwd,
		'home': environ.get('HOME', home),
		'environ': environ,
		'args': args,
	}
	key = (
		args.ext[0],
		args.renderer_module,
		tuple(args.config) if args.config else None,
		tuple(args.theme_option) if args.theme_option else None,
	)
	finish_args(args)
	powerline = None
	try:
		powerline = powerlines[key]
	except KeyError:
		try:
			powerline = powerlines[key] = PowerlineDaemon(
				args,
				logger=logger,
				config_loader=config_loader,
				run_once=False,
			)
			if logger is None:
				logger = powerline.logger
			if config_loader is None:
				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, s.write, encoding)
	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


encoding = getpreferredencoding() or 'UTF-8'


def safe_bytes(o, encoding=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):
	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:])))
	return shell_args, environ, cwd


def do_render(req):
	try:
		return safe_bytes(render(*parse_args(req)))
	except Exception as e:
		return safe_bytes(str(e))


def do_one(sock, read_sockets, write_sockets, result_map):
	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 = do_render(req)
				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 main_loop(sock):
	sock.listen(1)
	sock.setblocking(0)

	read_sockets, write_sockets = set(), set()
	result_map = {}
	try:
		while True:
			do_one(sock, read_sockets, write_sockets, result_map)
	except KeyboardInterrupt:
		raise SystemExit(0)
	return 0


def daemonize(stdin=os.devnull, stdout=os.devnull, stderr=os.devnull):
	try:
		pid = os.fork()
		if pid > 0:
			# exit first parent
			sys.exit(0)
	except OSError as e:
		sys.stderr.write("fork #1 failed: %d (%s)\n" % (e.errno, e.strerror))
		sys.exit(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
			sys.exit(0)
	except OSError as e:
		sys.stderr.write("fork #2 failed: %d (%s)\n" % (e.errno, e.strerror))
		sys.exit(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())
	global is_daemon
	is_daemon = True


def check_existing():
	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():
	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(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():
	import fcntl
	import atexit
	import stat
	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, fd)
	signal(SIGTERM, cleanup)
	atexit.register(cleanup)
	return fd


def main():
	global address
	global pidfile
	p = ArgumentParser(description='Daemon to improve the performance of powerline')
	p.add_argument('--quiet', '-q', action='store_true', help='Without other options: do not complain about already running powerline-daemon instance. Will still exit with 1. With `--kill\' and `--replace\': do not show any messages. With `--foreground\': ignored. Does not silence exceptions in any case.')
	p.add_argument('--socket', '-s', help='Specify socket which will be used for connecting to daemon.')
	a = p.add_mutually_exclusive_group().add_argument
	a('--kill', '-k', action='store_true', help='Kill an already running instance')
	a('--foreground', '-f', action='store_true', help='Run in the foreground (dont daemonize)')
	a('--replace', '-r', action='store_true', help='Replace an already running instance')
	args = p.parse_args()

	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 kill_daemon():
			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():
			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
		daemonize()

	if use_filesystem:
		# Create a locked pid file containing the daemon's PID
		if lockpidfile() 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()
	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 args.foreground:
		return main_loop(sock)

	if not use_filesystem:
		# We daemonize on linux
		daemonize()

	main_loop(sock)


if __name__ == '__main__':
	main()