alarm

#!/usr/bin/env python3

import itertools as it, operator as op, functools as ft
from datetime import datetime, timedelta
from xml.sax.saxutils import escape as xml_escape
from html.parser import HTMLParser
from os.path import join, dirname, basename
import os, sys, errno, logging, stat, types, tempfile, math, re, atexit
import contextlib, time, subprocess, traceback, signal, fcntl


class LogMessage:
	def __init__(self, fmt, a, k): self.fmt, self.a, self.k = fmt, a, k
	def __str__(self): return self.fmt.format(*self.a, **self.k) if self.a or self.k else self.fmt

class LogStyleAdapter(logging.LoggerAdapter):
	def __init__(self, logger, extra=None):
		super(LogStyleAdapter, self).__init__(logger, extra or {})
	def log(self, level, msg, *args, **kws):
		if not self.isEnabledFor(level): return
		log_kws = {} if 'exc_info' not in kws else dict(exc_info=kws.pop('exc_info'))
		msg, kws = self.process(msg, kws)
		self.logger._log(level, LogMessage(msg, args, kws), (), **log_kws)

get_logger = lambda name: LogStyleAdapter(logging.getLogger(name))


_short_ts_days = dict(
	y=365.25, yr=365.25, year=365.25,
	mo=30.5, month=30.5, w=7, week=7, d=1, day=1 )
_short_ts_s = dict(
	h=3600, hr=3600, hour=3600,
	m=60, min=60, minute=60,
	s=1, sec=1, second=1 )

def _short_ts_regexp():
	sub_sort = lambda d: sorted(
		d.items(), key=lambda kv: (kv[1], len(kv[0])), reverse=True )
	ts_re = ['^'] + [
		r'(?P<{0}>\d+{0}\s*)?'.format(k)
		for k, v in it.chain.from_iterable(
			map(sub_sort, [_short_ts_days, _short_ts_s]) ) ]
	return re.compile(''.join(ts_re), re.I | re.U)
_short_ts_regexp = _short_ts_regexp()

def parse_timestamp(ts_str):
	assert isinstance(ts_str, str), [type(ts_str), repr(ts_str)]
	ts_str = ts_str.replace('_', ' ')

	# Try to parse time offset in short format
	match = _short_ts_regexp.search(ts_str)
	if match and any(match.groups()):
		delta = list()
		parse_int = lambda v: int(''.join(c for c in v if c.isdigit()))
		for units in [_short_ts_days, _short_ts_s]:
			val = 0
			for k, v in units.items():
				try:
					if not match.group(k): continue
					n = parse_int(match.group(k))
				except IndexError: continue
				val += n * v
			delta.append(val)
		return datetime.now() + timedelta(*delta)

	# Fallback to other generic formats
	ts = None
	if not ts:
		match = re.search( # common BE format
			r'^(?P<date>(?:\d{2}|(?P<Y>\d{4}))-\d{2}-\d{2})'
			r'(?:[ T](?P<time>\d{2}(?::\d{2}(?::\d{2})?)?)?)?$', ts_str )
		if match:
			tpl = 'y' if not match.group('Y') else 'Y'
			tpl, ts_str = '%{}-%m-%d'.format(tpl), match.group('date')
			if match.group('time'):
				tpl_time = ['%H', '%M', '%S']
				ts_str_time = match.group('time').split(':')
				ts_str += ' ' + ':'.join(ts_str_time)
				tpl += ' ' + ':'.join(tpl_time[:len(ts_str_time)])
			try: ts = datetime.strptime(ts_str, tpl)
			except ValueError: pass
	if not ts:
		# coreutils' "date" parses virtually everything, but is more expensive to use
		proc = subprocess.Popen( ['date', '+%s', '-d', ts_str],
				stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, close_fds=True )
		val = proc.stdout.read()
		if not proc.wait():
			val = int(val.strip())
			# Try to add +1 day to simple timestamps like 3:00 if they're in the past
			# Whitelisted cases: 1:00, 4am, 5am GMT, 3:30 UTC-4
			if 0 < time.time() - val <= 24*3600 and re.search(
					r'(?i)^[\d:]+\s*(am|pm)?\s*([-+][\d:]+|\w+|\w+[-+][\d:]+)?$', ts_str.strip() ):
				val += 24*3600
			ts = datetime.fromtimestamp(val)

	if ts: return ts
	raise ValueError('Unable to parse date/time string: {0}'.format(ts_str))


def naturaltime_diff( ts, ts0=None, ext=None,
		_units=dict( h=3600, m=60, s=1,
			y=365.25*86400, mo=30.5*86400, w=7*86400, d=1*86400 ) ):
	delta = abs(
		(ts - (ts0 or datetime.now()))
		if not isinstance(ts, timedelta) else ts )

	res, s = list(), delta.total_seconds()
	for unit, unit_s in sorted(_units.items(), key=op.itemgetter(1), reverse=True):
		val = math.floor(s / float(unit_s))
		if not val: continue
		res.append('{:.0f}{}'.format(val, unit))
		if len(res) >= 2: break
		s -= val * unit_s

	if not res: return 'now'
	if ext: res.append(ext)
	return ' '.join(res)


@contextlib.contextmanager
def safe_replacement(path, mode=None):
	if mode is None:
		with contextlib.suppress(OSError):
			mode = stat.S_IMODE(os.lstat(path).st_mode)
	kws = dict(delete=False, dir=dirname(path), prefix=basename(path)+'.')
	with tempfile.NamedTemporaryFile(**kws) as tmp:
		try:
			if mode is not None: os.fchmod(tmp.fileno(), mode)
			yield tmp
			if not tmp.closed: tmp.flush()
			os.rename(tmp.name, path)
		finally:
			with contextlib.suppress(OSError): os.unlink(tmp.name)

@contextlib.contextmanager
def cleanup_on_err(file_obj):
	try: yield
	except:
		with contextlib.suppress(OSError): os.unlink(file_obj.name)
		file_obj.close()
		raise


def notify(title, body, critical=False, timeout=None, icon=None, fork=True):
	if fork and os.fork(): return
	try:
		import gi
		gi.require_version('Notify', '0.7')
		from gi.repository import Notify
		if not getattr(notify, '_init', False):
			Notify.init('logtail')
			notify._init = True
		note = Notify.Notification.new(summary=title, body=body, icon=icon)
		if critical: note.set_urgency(Notify.Urgency.CRITICAL)
		if timeout is not None: note.set_timeout(timeout)
		note.show()
	except: pass
	if fork: os._exit(0)


class MLStripper(HTMLParser):
	def __init__(self):
		super(MLStripper, self).__init__()
		self.fed = list()
	def handle_data(self, d): self.fed.append(d)
	def get_data(self): return ''.join(self.fed)

def strip_tags(html):
	s = MLStripper()
	s.feed(html)
	s.close()
	return s.get_data()


pid_file_dir = '.fg-alarm-pids'

def main(args=None):
	import argparse
	parser = argparse.ArgumentParser(
		description='Send notification at a certain future point in time.',
		epilog=r'Sending SIGQUIT (Ctrl+\ [Ctrl + backslash] from'
			' terminal) to the process will print time left until notification to stdout.')

	parser.add_argument('time', nargs='?',
		help='Absolute or relative time spec until notification.')
	parser.add_argument('summary', nargs='?',
		help='Optional notification summary (e.g. alarm reason).')
	parser.add_argument('body', nargs='?',
		help='Optional notification body. Always prefixed by timer info.'
			' Default is to print some machine info (uname) here.')

	parser.add_argument('-c', '--body-cmd', action='append', metavar='cmd/args',
		help='Command to run at the time of triggering and use output in a notification body.'
			' Not passed to shell, but split into arguments on spaces, if specified only once.'
			' Can be specified multiple times, in which case all'
				' values will be used as command/args as-is, without splitting.'
			' Notification will be delayed until command exits.')

	parser.add_argument('-f', '--fork', action='store_true',
		help='Fork-out immediately and wait for time to come in the background.')
	parser.add_argument('-a', '--at', action='store_true',
		help='Schedule notification as a job via "at" daemon/command.'
			' This should work with any posix-ish at(1p), and allow for reboot-persistent alarms.'
			' One issue though is that desktop environment might go stale/away between reboots.')

	parser.add_argument('-p', '--pid-file', action='store_true',
		help=(
			'Create pid file for the process in TMPDIR, so that it can be easily tracked.'
			' Also creates .msg file alongside .pid with'
				' the alarm message summary (to e.g. identify it).'
			' All such files will be created in one "{}" directory'
				' (created under TMPDIR, if does not exists), and have exclusive'
				' lock held on them while process runs, cleaning-up these files (but not dir) afterwards.'
			).format(pid_file_dir))
	parser.add_argument('-l', '--list', action='store_true',
		help='List delayed alarms, for which -p/--pid-file option was used, and exit.'
			' All other options except --kill are ignored with this one.')
	parser.add_argument('-k', '--kill',
		const=True, nargs='?', action='append', metavar='pid',
		help='Kill either specified (by pid) alarm process'
			' (making sure it is the right pid) or all of the ones that would be printed by --list,'
			' (i.e. for which --pid-file option was specified), and exit.'
			' Can be used multiple times. All other options except --list are ignored with this one.')

	parser.add_argument('-x', '--notification-non-critical', action='store_true',
		help='Default is to issue notifications with'
			' "critical" urgency set. This option disables that.')
	parser.add_argument('-t', '--notification-timeout',
		default=0, type=int, metavar='seconds',
		help='Timeout to use for notification, in seconds (default: %(default)s).'
			' "0" means "no timeout" (according to desktop notification spec).')
	parser.add_argument('-i', '--notification-icon',
		default='alarm', metavar='xdg-icon-or-path',
		help='Icon to use for notification (default: %(default)s). Empty - no icon.')

	parser.add_argument('-s', '--notification-sound', metavar='xdg-sound',
		help='Sound to play via canberra-gtk-play for notification (default: no sound).')
	parser.add_argument('-1', '--important', action='store_true',
		help='If any of the usual notification means fail,'
			' try email (to current $USER) and wall(1) notifications.')
	parser.add_argument('-e', '--email', action='store_true',
		help='Send email as per --important, but always do it,'
				' regardless of whether other notification means succeed or fail.'
			' --important still adds e.g. "wall" notification'
				' fallback if email fails, so not completely with this option.')

	parser.add_argument('-n', '--dry-run', action='store_true',
		help='Do not queue the actual alarm, stopping after time parsing stage.')
	parser.add_argument('-q', '--quiet', action='store_true', help='Quiet operation mode.')
	parser.add_argument('-d', '--debug', action='store_true', help='Verbose operation mode.')
	opts = parser.parse_args(sys.argv[1:] if args is None else args)

	global log
	logging.basicConfig(level=logging.DEBUG if opts.debug else logging.WARNING)
	log = get_logger('main')

	p = lambda fmt,*a,**k: print(fmt.format(*a,**k), flush=True)

	if opts.list or opts.kill:
		alarms, pid_dir = dict(), join(tempfile.gettempdir(), pid_file_dir)
		try: pids = os.listdir(pid_dir)
		except OSError as err:
			if err.errno != errno.ENOENT: raise
			pids = list()
		for fn in pids:
			if not fn.endswith('.pid'): continue
			pid_file = join(pid_dir, fn)
			with contextlib.suppress(OSError):
				with open(pid_file, 'a+') as src:
					try: fcntl.lockf(src.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
					except OSError: pass
					else: continue
					src.seek(0)
					pid = int(src.read().strip())
				with open('{}.msg'.format(pid_file[:-4])) as src:
					ts0, ts = map(float, src.readline().strip().split())
					msg = src.read().strip()
				alarms[pid] = ts0, ts, pid, msg
		if opts.list:
			for ts0, ts, pid, msg in sorted(alarms.values()):
				ts0, ts = (datetime.fromtimestamp(int(ts)) for ts in [ts0, ts])
				msg = msg.replace('\n', ' // ')
				if len(msg) > 80: msg = '{} ...[len={}]'.format(msg[:60], len(msg))
				p('{:>7d} :: {} (set: {}, in {}) :: {}', pid, ts, ts0, naturaltime_diff(ts), msg)
		if opts.kill:
			pids = set()
			for pid in opts.kill:
				if pid is True: pids.update(alarms)
				else: pids.add(int(pid))
			for pid in pids:
				if pid not in alarms:
					log.warning('Alarm process with specified pid does is not running: {}', pid)
					continue
				try: os.kill(pid, signal.SIGTERM)
				except OSError as err:
					log.warning('Failed to signal alarm-pid {}, skipping it: {}', pid, err)
		return

	elif not opts.time: parser.error('"time" argument must be specified.')

	icon = opts.notification_icon
	if icon and os.path.exists(icon): icon = os.path.realpath(icon)
	os.chdir('/') # to release mountpoints in case of long sleep

	pid_files, ts_env = None, os.environ.get('ALARM_AT_TS')
	if not ts_env:
		ts0, ts = datetime.now(), parse_timestamp(opts.time)
		ts_diff = abs(ts - ts0).total_seconds()
		if ts_diff < 2: ts, ts_diff = ts0, 0 # to account for "slightly in the past by now" timestamps
		if ts0 > ts:
			parser.error( 'Destination time cannot be in the past'
				' (parsed {!r} as {}, {})'.format(opts.time, ts, naturaltime_diff(ts, ts0, 'ago')) )
		if not opts.quiet:
			p('Parsed time spec {!r} as {} (delta: {})', opts.time, ts, naturaltime_diff(ts, ts0))
		if opts.dry_run: sys.exit()

		if opts.fork:
			r, w = os.pipe()
			if os.fork():
				pid = os.fdopen(r).readline().strip()
				print('Forked notification'
					' process to background (pid: {})'.format(pid or '???'))
				sys.exit()
			os.close(r)
			with contextlib.suppress(OSError): os.setsid()
			pid = os.fork()
			if pid:
				os.write(w, '{}\n'.format(pid).encode())
				sys.exit()
			os.close(w)

		if opts.at:
			os.environ['ALARM_AT_TS'] = '{} {}'.format(
				time.mktime(ts0.timetuple()), time.mktime(ts.timetuple()) )
			with open(os.devnull, 'w') as devnull:
				proc = subprocess.Popen(
					['at', ts.strftime('%H:%M %Y-%m-%d')],
					stdin=subprocess.PIPE, stdout=devnull,
					stderr=subprocess.STDOUT, close_fds=True )
				proc.stdin.write(' '.join( [sys.argv[0]]
					+ list("'{}'".format(arg.replace("'", r"'\''")) for arg in sys.argv[1:]) ))
				proc.stdin.close()
				err = proc.wait()
				if err: raise RuntimeError('"at" command failed (exit status: {})'.format(err))
			return

		if opts.pid_file:
			pid_file = join(tempfile.gettempdir(), pid_file_dir)
			with contextlib.suppress(OSError): os.makedirs(pid_file)
			pid_file = join(pid_file, 'alarm-{:08d}'.format(os.getpid()))
			pid_file, pid_file_msg = ('{}.{}'.format(pid_file, ext) for ext in ['pid', 'msg'])
			try:
				with contextlib.ExitStack() as stack:
					pid_file = open(pid_file, 'a+')
					stack.enter_context(cleanup_on_err(pid_file))
					pid_file.seek(0)
					pid_file.truncate()
					pid_file.write('{}\n'.format(os.getpid()))
					pid_file.flush()
					fcntl.lockf(pid_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
					pid_file_msg = open(pid_file_msg, 'w')
					stack.enter_context(cleanup_on_err(pid_file_msg))
					pid_file_msg.write('{:f} {:f}\n'.format(ts0.timestamp(), ts.timestamp()))
					pid_file_msg.write(opts.summary or 'notification [{}]'.format(ts0))
					pid_file_msg.flush()
			except OSError as err: log.error('Failed to create/lock pid/msg file(s): {}', err)
			else: pid_files = [pid_file, pid_file_msg]
			del pid_file, pid_file_msg
			def pid_file_cleanup():
				if pid_files:
					for pid_file in pid_files:
						pid_file.close()
						try: os.unlink(pid_file.name)
						except OSError as err: log.debug('Failed to remove pid/msg file: {}', err)
					pid_files.clear()
			for sig in 'INT', 'TERM':
				signal.signal(getattr(signal, 'SIG{}'.format(sig)), lambda sig,frm: sys.exit(0))
			atexit.register(pid_file_cleanup)

		if ts_diff and ts_diff > 0:
			ts_diff_print = set()
			signal.signal(signal.SIGQUIT, lambda sig,frm: ts_diff_print.add(True))
			while True:
				ts1 = datetime.now()
				ts_diff = ts - ts1
				if ts_diff.total_seconds() < 0: break
				if ts_diff_print:
					p( 'Time left: {} (total: {:.1f}s)',\
						naturaltime_diff(ts_diff), abs(ts_diff.total_seconds()) )
					ts_diff_print.clear()
				time.sleep(max(0, ts_diff.total_seconds()))

	else: # running from serialized state at the right time, e.g. "at" job
		ts0, ts = (datetime.fromtimestamp(float(ts)) for ts in ts_env.strip().split())

	ts1 = datetime.now()
	ts_diff = abs(ts1 - ts).total_seconds()
	if ts_diff > 60:
		log.info( 'Detected significant deviation from'
			' expected wake-up time: {}s (expected: {}, actual: {})', ts_diff, ts, ts1 )
	ts = ts1

	summary = ( opts.summary
		or 'notification from {}'.format(naturaltime_diff(ts, ts0, 'ago')) )

	body_cmd = None
	if opts.body_cmd:
		body_cmd = opts.body_cmd
		if len(body_cmd) == 1: body_cmd = body_cmd[0].split()
		try:
			body_cmd = 'Command output', '\n<tt>{}</tt>'.format(
				xml_escape(subprocess.check_output(body_cmd, close_fds=True)) )
		except Exception as err:
			body_cmd = ( '<span color="red"'
					' stretch="expanded" underline="double">'
				'Error while running command:</span>\n' )
			body_cmd += xml_escape(traceback.format_exc())

	label_len, uname = 14, list(map(xml_escape, os.uname()))
	body = [
		('Current time', datetime.now()),
		('Set', '{} (<b>{}</b>)'.format(ts0, naturaltime_diff(ts, ts0, 'ago'))) ]
	if opts.body: body.extend([('', ''), ('Info', '\n' + opts.body)])
	else:
		body.append(
			('Machine', '\n  {}\n  {}'.format(' '.join(uname[:3] + [uname[4]]), uname[3])) )
	if body_cmd:
		body.append(('', ''))
		if isinstance(body_cmd, tuple): body.append(body_cmd)
		elif isinstance(body_cmd, list): body.extend(body_cmd)
		else: body.append((None, body_cmd))
	body = '\n'.join(
		('<tt>{{:<{}s}}</tt>{{}}'.format(label_len).format(t + ':', v) if t else v)
		for t, v in body )

	any_fails = None
	try:
		proc = None
		if opts.notification_sound:
			try: proc = Popen(['canberra-gtk-play', '-i', opts.notification_sound])
			except Exception as err:
				log.exception('Failed to start "canberra-gtk-play" subprocess: {}', err)
				any_fails = True

		notify(
			'Alarm: {}'.format(summary), body, fork=False,
			critical=not opts.notification_non_critical,
			timeout=opts.notification_timeout * 1000, icon=icon or None )

		if proc:
			err = proc.wait()
			if err:
				log.error('"canberra-gtk-play" subprocess exited with error status: {}', err)
				any_fails = True
	except Exception as err:
		any_fails = True
		log.exception('Notification delivery failure: {}', err)

	if opts.email or (any_fails and opts.important):
		log.debug('Using fallback notification mechanisms due to --important flag')
		proc_summary = ' [fallback delivery]' if not opts.email else ''
		proc_summary, proc_body = '[Alarm] {}{}'.format(summary, proc_summary), strip_tags(body)
		proc = subprocess.Popen([ 'mail', '-s', proc_summary,
			os.environ['USER'] ], stdin=subprocess.PIPE, close_fds=True)
		proc.stdin.write(proc_body.encode())
		proc.stdin.close()
		err = proc.wait()
		if err: log.error('"mail" subprocess exited with error status: {}', err)
		if opts.important:
			proc = subprocess.Popen(['wall', '-t', '5'], stdin=subprocess.PIPE, close_fds=True)
			proc.stdin.write('\n'.join([proc_summary, '', proc_body, '']).encode())
			proc.stdin.close()
			err = proc.wait()
			if err: log.error('"wall" subprocess exited with error status: {}', err)

	if pid_files: pid_file_cleanup()

	return any_fails

if __name__ == '__main__': sys.exit(main())