You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
424 lines
12 KiB
424 lines
12 KiB
#!/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:]))
|