#!/usr/bin/env python

# Copyright 2009-2016 Peter Poeml
# 
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
# 
#        http://www.apache.org/licenses/LICENSE-2.0
# 
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
#
#
#
# withlock -- https://github.com/poeml/withlock
#
#
# A locking wrapper that make sure that a program isn't run more than once.
# It creates locks that are valid only while the wrapper is running, and thus
# will never require a cleanup, e.g. after a reboot.

# Parts of the locking strategy, and parts of the usage semantics, of this
# script were inspired from with-lock-ex.c, which was written by Ian Jackson and
# placed in the public domain by him.
#
#
# Usage is simple. Instead of your command 
#   CMD ARGS...
# you simply use
#   withlock LOCKFILE CMD ARGS...
#
# See --help output for more options.

# Modernize features for Python 2 and 3 compatibility
from __future__ import absolute_import
from __future__ import print_function

__version__ = '0.5'

import os
import os.path
import errno
import sys
import time
import stat
import fcntl
import signal
import atexit
from optparse import OptionParser

got_lock = False

class SignalInterrupt(Exception):
    """Exception raised on SIGTERM and SIGHUP."""


def catchterm(*args):
    raise SignalInterrupt


for name in 'SIGBREAK', 'SIGHUP', 'SIGTERM':
    num = getattr(signal, name, None)
    if num: signal.signal(num, catchterm)


def cleanup(lockfile, verbose):
    global got_lock
    if got_lock:
        if verbose:
            sys.stderr.write('removing lockfile: %r\n' % lockfile)

        try:
            os.unlink(lockfile)
        except:
            pass


def main():

    usage = 'usage: %prog [options] LOCKFILE CMD ARGS...'
    version = '%prog ' + __version__

    parser = OptionParser(usage=usage, version=version)
    parser.disable_interspersed_args()

    parser.add_option('-w', '--wait',
                      dest='wait',
                      help="wait for maximum SECONDS until the lock is acquired",
                      metavar="SECONDS")

    parser.add_option("-q", "--quiet",
                      action="store_true", dest="quiet", default=False,
                      help="if lock can't be acquired immediately, silently "
                           "quit without error")

    parser.add_option("-v", "--verbose",
                      action="store_true", dest="verbose", default=False,
                      help="print debug messages to stderr")



    (options, args) = parser.parse_args()

    usage = usage.replace('%prog', os.path.basename(sys.argv[0]))


    if len(args) < 2:
        sys.exit(usage)

    lockfile = args[0]
    cmd = args[1:]
    waited = 0
    if options.wait:
        options.wait = int(options.wait)

    if options.verbose:
        sys.stderr.write('lockfile: %r\n' % lockfile)
        sys.stderr.write('cmd: %r\n' % cmd)

    # check that symlink attacks in the lockfile directory are not possible
    lockdir = os.path.dirname(os.path.realpath(lockfile)) or os.curdir
    lockdir_mode = stat.S_IMODE(os.stat(lockdir).st_mode)
    if options.verbose:
        sys.stderr.write('lockdir: %r\n' % lockdir)
        sys.stderr.write('lockdir mode: %s\n' % oct(lockdir_mode))
    if (lockdir_mode & stat.S_IWGRP) or (lockdir_mode & stat.S_IWOTH):
        sys.stderr.write('withlock: the destination directory for %r is %r, \n'
                         '          which is writable by other users. That allows for symlink attacks. \n'
                         '          Choose another directory.\n' \
                                % (lockfile, lockdir))
        sys.exit(3)


    prev_umask = os.umask(0o066)

    lock = open(lockfile, 'w')

    global got_lock
    while 1 + 1 == 2:

        try:
            fcntl.lockf(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            got_lock = True
            break
        except IOError as e:
            if e.errno in [ errno.EAGAIN, 
                            errno.EACCES, 
                            errno.EWOULDBLOCK, 
                            errno.EBUSY ]:
                if options.wait:
                    if waited >= options.wait:
                        if options.quiet:
                            if options.verbose:
                                sys.stderr.write('quitting silently\n')
                            sys.exit(0)
                        else:
                            sys.stderr.write('could not acquire lock\n')
                            sys.exit(1)
                    if options.verbose and waited == 0:
                        sys.stderr.write('waiting up to %ss for lock\n' \
                                            % options.wait)
                    time.sleep(1)
                    waited += 1
                    continue
                elif options.quiet:
                    if options.verbose:
                        sys.stderr.write('quitting silently\n')
                    sys.exit(0)
                else:
                    sys.exit('could not acquire lock')
                
        try:
            os.stat(lockfile)
            break
        except OSError as e: 
            if e.errno == errno.ENOENT:
                sys.exit('====== could not stat %s, which we have just locked' \
                            % lockfile)


    atexit.register(cleanup, lockfile, options.verbose)
    os.umask(prev_umask)

    import subprocess
    rc = subprocess.call(' '.join(cmd), shell=True)
    sys.stdout.flush()
    sys.stderr.flush()

    if options.verbose:
        sys.stderr.write('command terminated with exit code %s\n' % rc)

    sys.exit(rc)


if __name__ == '__main__':

    try:
        main()

    except SignalInterrupt:
        print('killed!', file=sys.stderr)

    except KeyboardInterrupt:
        print('interrupted!', file=sys.stderr)