#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2019 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Script to lock/unlock machines."""


__author__ = "asharif@google.com (Ahmad Sharif)"

import argparse
import datetime
import fcntl
import getpass
import glob
import json
import os
import socket
import sys
import time

from cros_utils import logger


LOCK_SUFFIX = "_check_lock_liveness"

# The locks file directory REQUIRES that 'group' only has read/write
# privileges and 'world' has no privileges.  So the mask must be
# '0o27': 0o777 - 0o27 = 0o750.
LOCK_MASK = 0o27


def FileCheckName(name):
    return name + LOCK_SUFFIX


def OpenLiveCheck(file_name):
    with FileCreationMask(LOCK_MASK):
        fd = open(file_name, "a")
    try:
        fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError as e:
        logger.GetLogger().LogError(e)
        raise
    return fd


class FileCreationMask(object):
    """Class for the file creation mask."""

    def __init__(self, mask):
        self._mask = mask
        self._old_mask = None

    def __enter__(self):
        self._old_mask = os.umask(self._mask)

    def __exit__(self, typ, value, traceback):
        os.umask(self._old_mask)


class LockDescription(object):
    """The description of the lock."""

    def __init__(self, desc=None):
        try:
            self.owner = desc["owner"]
            self.exclusive = desc["exclusive"]
            self.counter = desc["counter"]
            self.time = desc["time"]
            self.reason = desc["reason"]
            self.auto = desc["auto"]
        except (KeyError, TypeError):
            self.owner = ""
            self.exclusive = False
            self.counter = 0
            self.time = 0
            self.reason = ""
            self.auto = False

    def IsLocked(self):
        return self.counter or self.exclusive

    def __str__(self):
        return " ".join(
            [
                "Owner: %s" % self.owner,
                "Exclusive: %s" % self.exclusive,
                "Counter: %s" % self.counter,
                "Time: %s" % self.time,
                "Reason: %s" % self.reason,
                "Auto: %s" % self.auto,
            ]
        )


class FileLock(object):
    """File lock operation class."""

    FILE_OPS = []

    def __init__(self, lock_filename):
        self._filepath = lock_filename
        lock_dir = os.path.dirname(lock_filename)
        assert os.path.isdir(lock_dir), (
            "Locks dir: %s doesn't exist!" % lock_dir
        )
        self._file = None
        self._description = None

        self.exclusive = None
        self.auto = None
        self.reason = None
        self.time = None
        self.owner = None

    def getDescription(self):
        return self._description

    def getFilePath(self):
        return self._filepath

    def setDescription(self, desc):
        self._description = desc

    @classmethod
    def AsString(cls, file_locks):
        stringify_fmt = "%-30s %-15s %-4s %-4s %-15s %-40s %-4s"
        header = stringify_fmt % (
            "machine",
            "owner",
            "excl",
            "ctr",
            "elapsed",
            "reason",
            "auto",
        )
        lock_strings = []
        for file_lock in file_locks:

            elapsed_time = datetime.timedelta(
                seconds=int(time.time() - file_lock.getDescription().time)
            )
            elapsed_time = "%s ago" % elapsed_time
            lock_strings.append(
                stringify_fmt
                % (
                    os.path.basename(file_lock.getFilePath),
                    file_lock.getDescription().owner,
                    file_lock.getDescription().exclusive,
                    file_lock.getDescription().counter,
                    elapsed_time,
                    file_lock.getDescription().reason,
                    file_lock.getDescription().auto,
                )
            )
        table = "\n".join(lock_strings)
        return "\n".join([header, table])

    @classmethod
    def ListLock(cls, pattern, locks_dir):
        if not locks_dir:
            locks_dir = Machine.LOCKS_DIR
        full_pattern = os.path.join(locks_dir, pattern)
        file_locks = []
        for lock_filename in glob.glob(full_pattern):
            if LOCK_SUFFIX in lock_filename:
                continue
            file_lock = FileLock(lock_filename)
            with file_lock as lock:
                if lock.IsLocked():
                    file_locks.append(file_lock)
        logger.GetLogger().LogOutput("\n%s" % cls.AsString(file_locks))

    def __enter__(self):
        with FileCreationMask(LOCK_MASK):
            try:
                self._file = open(self._filepath, "a+")
                self._file.seek(0, os.SEEK_SET)

                if fcntl.flock(self._file.fileno(), fcntl.LOCK_EX) == -1:
                    raise IOError("flock(%s, LOCK_EX) failed!" % self._filepath)

                try:
                    desc = json.load(self._file)
                except (EOFError, ValueError):
                    desc = None
                self._description = LockDescription(desc)

                if self._description.exclusive and self._description.auto:
                    locked_byself = False
                    for fd in self.FILE_OPS:
                        if fd.name == FileCheckName(self._filepath):
                            locked_byself = True
                            break
                    if not locked_byself:
                        try:
                            fp = OpenLiveCheck(FileCheckName(self._filepath))
                        except IOError:
                            pass
                        else:
                            self._description = LockDescription()
                            fcntl.lockf(fp, fcntl.LOCK_UN)
                            fp.close()
                return self._description
            # Check this differently?
            except IOError as ex:
                logger.GetLogger().LogError(ex)
                return None

    def __exit__(self, typ, value, traceback):
        self._file.truncate(0)
        self._file.write(json.dumps(self._description.__dict__, skipkeys=True))
        self._file.close()

    def __str__(self):
        return self.AsString([self])


class Lock(object):
    """Lock class"""

    def __init__(self, lock_file, auto=True):
        self._to_lock = os.path.basename(lock_file)
        self._lock_file = lock_file
        self._logger = logger.GetLogger()
        self._auto = auto

    def NonBlockingLock(self, exclusive, reason=""):
        with FileLock(self._lock_file) as lock:
            if lock.exclusive:
                self._logger.LogError(
                    "Exclusive lock already acquired by %s. Reason: %s"
                    % (lock.owner, lock.reason)
                )
                return False

            if exclusive:
                if lock.counter:
                    self._logger.LogError("Shared lock already acquired")
                    return False
                lock_file_check = FileCheckName(self._lock_file)
                fd = OpenLiveCheck(lock_file_check)
                FileLock.FILE_OPS.append(fd)

                lock.exclusive = True
                lock.reason = reason
                lock.owner = getpass.getuser()
                lock.time = time.time()
                lock.auto = self._auto
            else:
                lock.counter += 1
        self._logger.LogOutput("Successfully locked: %s" % self._to_lock)
        return True

    def Unlock(self, exclusive, force=False):
        with FileLock(self._lock_file) as lock:
            if not lock.IsLocked():
                self._logger.LogWarning("Can't unlock unlocked machine!")
                return True

            if lock.exclusive != exclusive:
                self._logger.LogError(
                    "shared locks must be unlocked with --shared"
                )
                return False

            if lock.exclusive:
                if lock.owner != getpass.getuser() and not force:
                    self._logger.LogError(
                        "%s can't unlock lock owned by: %s"
                        % (getpass.getuser(), lock.owner)
                    )
                    return False
                if lock.auto != self._auto:
                    self._logger.LogError(
                        "Can't unlock lock with different -a" " parameter."
                    )
                    return False
                lock.exclusive = False
                lock.reason = ""
                lock.owner = ""

                if self._auto:
                    del_list = [
                        i
                        for i in FileLock.FILE_OPS
                        if i.name == FileCheckName(self._lock_file)
                    ]
                    for i in del_list:
                        FileLock.FILE_OPS.remove(i)
                    for f in del_list:
                        fcntl.lockf(f, fcntl.LOCK_UN)
                        f.close()
                    del del_list
                    os.remove(FileCheckName(self._lock_file))

            else:
                lock.counter -= 1
        return True


class Machine(object):
    """Machine class"""

    LOCKS_DIR = "/google/data/rw/users/mo/mobiletc-prebuild/locks"

    def __init__(self, name, locks_dir=LOCKS_DIR, auto=True):
        self._name = name
        self._auto = auto
        try:
            self._full_name = socket.gethostbyaddr(name)[0]
        except socket.error:
            self._full_name = self._name
        self._full_name = os.path.join(locks_dir, self._full_name)

    def Lock(self, exclusive=False, reason=""):
        lock = Lock(self._full_name, self._auto)
        return lock.NonBlockingLock(exclusive, reason)

    def TryLock(self, timeout=300, exclusive=False, reason=""):
        locked = False
        sleep = timeout / 10
        while True:
            locked = self.Lock(exclusive, reason)
            if locked or timeout < 0:
                break
            print(
                "Lock not acquired for {0}, wait {1} seconds ...".format(
                    self._name, sleep
                )
            )
            time.sleep(sleep)
            timeout -= sleep
        return locked

    def Unlock(self, exclusive=False, ignore_ownership=False):
        lock = Lock(self._full_name, self._auto)
        return lock.Unlock(exclusive, ignore_ownership)


def Main(argv):
    """The main function."""

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-r", "--reason", dest="reason", default="", help="The lock reason."
    )
    parser.add_argument(
        "-u",
        "--unlock",
        dest="unlock",
        action="store_true",
        default=False,
        help="Use this to unlock.",
    )
    parser.add_argument(
        "-l",
        "--list_locks",
        dest="list_locks",
        action="store_true",
        default=False,
        help="Use this to list locks.",
    )
    parser.add_argument(
        "-f",
        "--ignore_ownership",
        dest="ignore_ownership",
        action="store_true",
        default=False,
        help="Use this to force unlock on a lock you don't own.",
    )
    parser.add_argument(
        "-s",
        "--shared",
        dest="shared",
        action="store_true",
        default=False,
        help="Use this for a shared (non-exclusive) lock.",
    )
    parser.add_argument(
        "-d",
        "--dir",
        dest="locks_dir",
        action="store",
        default=Machine.LOCKS_DIR,
        help="Use this to set different locks_dir",
    )
    parser.add_argument("args", nargs="*", help="Machine arg.")

    options = parser.parse_args(argv)

    options.locks_dir = os.path.abspath(options.locks_dir)
    exclusive = not options.shared

    if not options.list_locks and len(options.args) != 2:
        logger.GetLogger().LogError(
            "Either --list_locks or a machine arg is needed."
        )
        return 1

    if len(options.args) > 1:
        machine = Machine(options.args[1], options.locks_dir, auto=False)
    else:
        machine = None

    if options.list_locks:
        FileLock.ListLock("*", options.locks_dir)
        retval = True
    elif options.unlock:
        retval = machine.Unlock(exclusive, options.ignore_ownership)
    else:
        retval = machine.Lock(exclusive, options.reason)

    if retval:
        return 0
    else:
        return 1


if __name__ == "__main__":
    sys.exit(Main(sys.argv[1:]))