#
# Copyright (c) 2008--2016 Red Hat, Inc.
#
# This software is licensed to you under the GNU General Public License,
# version 2 (GPLv2). There is NO WARRANTY for this software, express or
# implied, including the implied warranties of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. You should have received a copy of GPLv2
# along with this software; if not, see
# http://www.gnu.org/licenses/old-licenses/gpl-2.0.txt.
#
# Red Hat trademarks are not licensed under GPLv2. No permission is
# granted to use or replicate Red Hat trademarks that are incorporated
# in this software or its documentation.
#

import os
import sys
import stat
import fcntl
from errno import EWOULDBLOCK, EEXIST
from rhn.i18n import bstr
import fcntl

class LockfileLockedException(Exception):
    """thrown ONLY when pid file is locked."""
    pass

class Lockfile:

    """class that provides simple access to a PID-style lockfile.

    methods: __init__(lockfile), acquire(), and release()
    NOTE: currently acquires upon init
    The *.pid file will be acquired, or an LockfileLockedException is raised.
    """

    def __init__(self, lockfile, pid=None):
        """create (if need be), and acquire lock on lockfile

        lockfile example: '/var/run/up2date.pid'
        """

        # cleanup the path and assign it.
        self.lockfile = os.path.abspath(
                          os.path.expanduser(
                            os.path.expandvars(lockfile)))

        self.pid = pid
        if not self.pid:
            self.pid = os.getpid()

        # create the directory structure
        dirname = os.path.dirname(self.lockfile)
        if not os.path.exists(dirname):
            try:
                os.makedirs(dirname)
            except OSError:
                e = sys.exc_info()[1]
                if hasattr(e, 'errno') and e.errno == EEXIST:
                    # race condition... dirname exists now.
                    pass
                else:
                    raise

        # open the file -- non-destructive read-write, unless it needs
        # to be created. O_NOFOLLOW refuses to follow a symlink planted at
        # the final path component and O_CLOEXEC keeps the descriptor from
        # leaking across exec(); this brings the lockfile open up to the
        # CLOS-4287 hardening the sibling root-owned os.open sites already
        # carry (up2date_client/up2dateLog.py, rhnreg.py, up2dateAuth.py,
        # config.py). O_EXCL/O_TRUNC are intentionally NOT added: a PID
        # lockfile legitimately pre-exists (it may be held by a live peer)
        # and is only truncated under flock in acquire().
        self.f = os.open(self.lockfile,
                         os.O_RDWR | os.O_CREAT | os.O_SYNC
                         | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600)

        # Defense-in-depth: the opened descriptor must be a regular file
        # owned by this process. O_NOFOLLOW only rejects a symlink at the
        # final component; this fstat additionally rejects a pre-planted
        # FIFO/device/foreign-owned file at the lock path -- closing the
        # reusable-library attack steelmanned in the verdict for callers
        # that pass a path under a non-root-owned directory. Fail closed
        # (refuse to operate) rather than acting on a suspicious inode.
        st = os.fstat(self.f)
        if not stat.S_ISREG(st.st_mode) or st.st_uid != os.geteuid():
            os.close(self.f)
            raise OSError(
                "refusing to use lockfile %s: not a regular file owned by "
                "uid %d (mode 0o%o, uid %d)"
                % (self.lockfile, os.geteuid(), st.st_mode, st.st_uid))
        self.acquire()

    def acquire(self):
        """acquire the lock; else raise LockfileLockedException."""

        try:
            fcntl.flock(self.f, fcntl.LOCK_EX|fcntl.LOCK_NB)
        except IOError:
            if sys.exc_info()[1].errno == EWOULDBLOCK:
                raise LockfileLockedException(
                  "cannot acquire lock on %s." % self.lockfile, None, sys.exc_info()[2])
            else:
                raise
        # unlock upon exit
        fcntl.fcntl(self.f, fcntl.F_SETFD, 1)
        # truncate and write the pid
        os.ftruncate(self.f, 0)
        os.write(self.f, bstr(str(self.pid) + '\n'))

    def release(self):
        # Remove the lock file, but only if the path still names the exact
        # inode we hold open. unlink-by-path is a TOCTOU sink: if the path
        # was swapped (or symlinked) between open and release, a bare
        # os.unlink(self.lockfile) would delete the attacker's substitute.
        # Compare the open fd's (dev, ino) against an lstat (no symlink
        # follow) of the path and only unlink on a match; otherwise leave
        # the foreign entry alone but still release our lock and fd.
        try:
            fd_st = os.fstat(self.f)
            path_st = os.lstat(self.lockfile)
            same_inode = (fd_st.st_dev == path_st.st_dev
                          and fd_st.st_ino == path_st.st_ino)
        except OSError:
            same_inode = False
        if same_inode:
            os.unlink(self.lockfile)
        fcntl.flock(self.f, fcntl.LOCK_UN)
        os.close(self.f)


def main():
    """test code"""

    try:
        L = Lockfile('./test.pid')
    except LockfileLockedException:
        sys.stderr.write("%s\n" % sys.exc_info()[1])
        sys.exit(-1)
    else:
        print("lock acquired ")
        print("...sleeping for 10 seconds")
        import time
        time.sleep(10)
        L.release()
        print("lock released ")

if __name__ == '__main__':
    # test code
    sys.exit(main() or 0)