#!/usr/bin/env python
#
# lsflock displays a list of files in a given directory which have been locked with the flock syscall, and the PID
# of the process holding the lock. The /proc/locks file is used to identify locks and related processes.
import os
import sys
import optparse

LOCKS_FILE = '/proc/locks'

def parse_command_line(args=None):
    """
    Parse options from the command line.

    :param args: Args passed to the script from the command line
    :return: Path to the directory given to check locks.
    """
    usage = 'usage: %prog [options] [directory]'
    parser = optparse.OptionParser(usage=usage)

    parser.add_option("-t", "--test", action="store_true", help=("Enable test mode. Simplifies output "
                                                                 "so that it's easier to parse"))
    parser.add_option("-r", "--recursive", action="store_true", help="Enable recursion")

    opts, args = parser.parse_args(args)

    # We only want a single directory given as an argument, but one must be provided.
    if len(args) == 0:
        parser.print_help()
        sys.exit(1)
    if len(args) > 1:
        sys.stderr.write("Error: Too many arguments.\n\n")
        parser.print_help()
        sys.exit(1)

    return opts, ''.join(args)


def get_locks(lockfile=LOCKS_FILE):
    """
    :param lockfile: Path to the file containing data about file locks. /proc/locks in Linux.
    :return: A dictionary of pid: inode for all locked file inodes.
    """
    locks = list()
    try:
        f = open(lockfile, "r")
    except IOError, error:
        # If for some reason this file can't be opened...
        print "Unable to open {0} for reading. {1}".format(lockfile, error.message)
        sys.exit(1)

    for line in f.readlines():
        columns = line.split()
        inode = columns[5].split(':')[2]
        pid = columns[4]
        locks.append({'inode': inode, 'pid': pid})

    return locks


def get_lock_pids_for_file(check_file, lockfile=LOCKS_FILE):
    """
    Parses the lock file (/proc/locks) to find all PIDs holding locks for the given file.

    :param check_file: Full path to the file you want to check for locks.
    :param lockfile: File to check for file locks. Should be /proc/locks.
    :return: List of PIDs of the process holding a lock on the given file. Empty list if none are found.
    """
    pids = list()
    try:
        f = open(lockfile, "r")
    except IOError, error:
        # If for some reason this file can't be opened...
        sys.stderr.write("Error: Unable to open {0} for reading. {1}".format(lockfile, error.message))
        sys.exit(1)

    try:
        inode = os.stat(check_file).st_ino
        for line in f:
            columns = line.split()
            lock_inode = columns[5].split(':')[2]
            if inode == int(lock_inode):
                pids.append(columns[4])
    except OSError:
        # It's possible that it won't be possible to stat the file due to permissions issues etc
        print sys.stderr.write("Error: Unable to check locks for {0}: Could not stat file.".format(check_file))

    return pids


def get_file_locks_by_pid(directory, recurse=False):
    """
    Recurses through a given directory and checks each file to see if a lock exists.

    :param directory: Path to the directory to check for locks.
    :param recurse: Set whether to search the directory recursively.
    :return: A dictionary mapping PIDs to the files they're holding locks for.
    """
    file_locks = {}
    if recurse:
        for root, dirs, files in os.walk(directory):
            for f in files:
                file_abspath = os.path.join(root, f)
                # Get all PIDs holding a lock for this file, if any exist, and add to the return dictionary.
                lock_pids = get_lock_pids_for_file(file_abspath)
                if lock_pids:
                    for pid in lock_pids:
                        file_locks[pid] = file_abspath
    else:
        files = os.listdir(directory)
        for file in files:
            file_abspath = os.path.join(directory, file)
            # Ignore directories if not searching recursively.
            if os.path.isdir(file_abspath):
                pass
            else:
                lock_pids = get_lock_pids_for_file(file_abspath)
                if lock_pids:
                    for pid in lock_pids:
                        file_locks[pid] = file_abspath

    return file_locks


def main():
    """
    Find locks for the given directory and print them to the command line.

    :return: None
    """
    opts, directory = parse_command_line()

    # Check locks for each directory given at the command line
    full_path = os.path.abspath(directory)

    # Make sure the value given is actually a directory.
    if os.path.isdir(full_path):
        file_locks = get_file_locks_by_pid(full_path, opts.recursive)
        if len(file_locks) > 0:
            if not opts.test:
                print "{0} locks found for files in {1}:\n".format(len(file_locks),full_path)
                print "{:>7}  {:>6}".format("PID", "FILE")
                for pid in file_locks.keys():
                    print "{:>8}  {:>8}".format(pid, file_locks[pid])
            else:
                # Don't do any fancy formatting in test mode.
                for pid in file_locks.keys():
                    print "{0} {1}".format(pid, file_locks[pid])
        else:
            if not opts.test:
                print "No locked files found in {0}".format(full_path)
    else:
        sys.stderr.write("Error: {0} is not a directory".format(full_path))
        sys.exit(1)


if __name__ == "__main__":
    # Call the main function when run interactively as a script
    try:
        main()
    except KeyboardInterrupt:
        print "Caught Keyboard Interrupt. Quitting."
        sys.exit(1)
