| #!/usr/bin/env python3 |
| # -*- coding: utf-8 -*- |
| # Copyright 2019 The Chromium OS Authors. All rights reserved. |
| # 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.""" |
| |
| from __future__ import division |
| from __future__ import print_function |
| |
| __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:])) |
