413 lines
10 KiB
Python
Executable File
413 lines
10 KiB
Python
Executable File
#!/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 StringIO
|
|
|
|
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
|
|
# use_filesystem = True
|
|
del platform
|
|
|
|
if use_filesystem:
|
|
address = '/tmp/powerline-ipc-%d'
|
|
pidfile = address + '.pid'
|
|
else:
|
|
# Use the abstract namespace for sockets rather than the filesystem
|
|
# (Available only in linux)
|
|
address = '\0powerline-ipc-%d'
|
|
address = address % os.getuid()
|
|
|
|
|
|
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 = StringIO()
|
|
write_output(args, powerline, segment_info, s.write)
|
|
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 + b'\0')
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
encoding = getpreferredencoding() or 'UTF-8'
|
|
if encoding.lower() == 'ascii':
|
|
encoding = '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:
|
|
print ("fork #1 failed: %d (%s)" % (e.errno, e.strerror), file=sys.stderr)
|
|
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:
|
|
print ("fork #2 failed: %d (%s)" % (e.errno, e.strerror), file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Redirect standard file descriptors.
|
|
si = file(stdin, 'r')
|
|
so = file(stdout, 'a+')
|
|
se = file(stderr, 'a+', 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():
|
|
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.')
|
|
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.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:
|
|
print ('The daemon is already running. Use %s -k to kill it.' % os.path.basename(sys.argv[0]),
|
|
file=sys.stderr)
|
|
raise SystemExit(1)
|
|
|
|
# Bind to address or bail if we cannot bind
|
|
sock = check_existing()
|
|
if sock is None:
|
|
if not args.quiet:
|
|
print ('The daemon is already running. Use %s -k to kill it.' % os.path.basename(sys.argv[0]),
|
|
file=sys.stderr)
|
|
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()
|