• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# Copyright 2019 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7"""Script to lock/unlock machines."""
8
9from __future__ import division
10from __future__ import print_function
11
12__author__ = 'asharif@google.com (Ahmad Sharif)'
13
14import argparse
15import datetime
16import fcntl
17import getpass
18import glob
19import json
20import os
21import socket
22import sys
23import time
24
25from cros_utils import logger
26
27LOCK_SUFFIX = '_check_lock_liveness'
28
29# The locks file directory REQUIRES that 'group' only has read/write
30# privileges and 'world' has no privileges.  So the mask must be
31# '0o27': 0o777 - 0o27 = 0o750.
32LOCK_MASK = 0o27
33
34
35def FileCheckName(name):
36  return name + LOCK_SUFFIX
37
38
39def OpenLiveCheck(file_name):
40  with FileCreationMask(LOCK_MASK):
41    fd = open(file_name, 'a')
42  try:
43    fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
44  except IOError as e:
45    logger.GetLogger().LogError(e)
46    raise
47  return fd
48
49
50class FileCreationMask(object):
51  """Class for the file creation mask."""
52
53  def __init__(self, mask):
54    self._mask = mask
55    self._old_mask = None
56
57  def __enter__(self):
58    self._old_mask = os.umask(self._mask)
59
60  def __exit__(self, typ, value, traceback):
61    os.umask(self._old_mask)
62
63
64class LockDescription(object):
65  """The description of the lock."""
66
67  def __init__(self, desc=None):
68    try:
69      self.owner = desc['owner']
70      self.exclusive = desc['exclusive']
71      self.counter = desc['counter']
72      self.time = desc['time']
73      self.reason = desc['reason']
74      self.auto = desc['auto']
75    except (KeyError, TypeError):
76      self.owner = ''
77      self.exclusive = False
78      self.counter = 0
79      self.time = 0
80      self.reason = ''
81      self.auto = False
82
83  def IsLocked(self):
84    return self.counter or self.exclusive
85
86  def __str__(self):
87    return ' '.join([
88        'Owner: %s' % self.owner,
89        'Exclusive: %s' % self.exclusive,
90        'Counter: %s' % self.counter,
91        'Time: %s' % self.time,
92        'Reason: %s' % self.reason,
93        'Auto: %s' % self.auto,
94    ])
95
96
97class FileLock(object):
98  """File lock operation class."""
99  FILE_OPS = []
100
101  def __init__(self, lock_filename):
102    self._filepath = lock_filename
103    lock_dir = os.path.dirname(lock_filename)
104    assert os.path.isdir(lock_dir), ("Locks dir: %s doesn't exist!" % lock_dir)
105    self._file = None
106    self._description = None
107
108    self.exclusive = None
109    self.auto = None
110    self.reason = None
111    self.time = None
112    self.owner = None
113
114  def getDescription(self):
115    return self._description
116
117  def getFilePath(self):
118    return self._filepath
119
120  def setDescription(self, desc):
121    self._description = desc
122
123  @classmethod
124  def AsString(cls, file_locks):
125    stringify_fmt = '%-30s %-15s %-4s %-4s %-15s %-40s %-4s'
126    header = stringify_fmt % ('machine', 'owner', 'excl', 'ctr', 'elapsed',
127                              'reason', 'auto')
128    lock_strings = []
129    for file_lock in file_locks:
130
131      elapsed_time = datetime.timedelta(
132          seconds=int(time.time() - file_lock.getDescription().time))
133      elapsed_time = '%s ago' % elapsed_time
134      lock_strings.append(
135          stringify_fmt % (os.path.basename(file_lock.getFilePath),
136                           file_lock.getDescription().owner,
137                           file_lock.getDescription().exclusive,
138                           file_lock.getDescription().counter, elapsed_time,
139                           file_lock.getDescription().reason,
140                           file_lock.getDescription().auto))
141    table = '\n'.join(lock_strings)
142    return '\n'.join([header, table])
143
144  @classmethod
145  def ListLock(cls, pattern, locks_dir):
146    if not locks_dir:
147      locks_dir = Machine.LOCKS_DIR
148    full_pattern = os.path.join(locks_dir, pattern)
149    file_locks = []
150    for lock_filename in glob.glob(full_pattern):
151      if LOCK_SUFFIX in lock_filename:
152        continue
153      file_lock = FileLock(lock_filename)
154      with file_lock as lock:
155        if lock.IsLocked():
156          file_locks.append(file_lock)
157    logger.GetLogger().LogOutput('\n%s' % cls.AsString(file_locks))
158
159  def __enter__(self):
160    with FileCreationMask(LOCK_MASK):
161      try:
162        self._file = open(self._filepath, 'a+')
163        self._file.seek(0, os.SEEK_SET)
164
165        if fcntl.flock(self._file.fileno(), fcntl.LOCK_EX) == -1:
166          raise IOError('flock(%s, LOCK_EX) failed!' % self._filepath)
167
168        try:
169          desc = json.load(self._file)
170        except (EOFError, ValueError):
171          desc = None
172        self._description = LockDescription(desc)
173
174        if self._description.exclusive and self._description.auto:
175          locked_byself = False
176          for fd in self.FILE_OPS:
177            if fd.name == FileCheckName(self._filepath):
178              locked_byself = True
179              break
180          if not locked_byself:
181            try:
182              fp = OpenLiveCheck(FileCheckName(self._filepath))
183            except IOError:
184              pass
185            else:
186              self._description = LockDescription()
187              fcntl.lockf(fp, fcntl.LOCK_UN)
188              fp.close()
189        return self._description
190      # Check this differently?
191      except IOError as ex:
192        logger.GetLogger().LogError(ex)
193        return None
194
195  def __exit__(self, typ, value, traceback):
196    self._file.truncate(0)
197    self._file.write(json.dumps(self._description.__dict__, skipkeys=True))
198    self._file.close()
199
200  def __str__(self):
201    return self.AsString([self])
202
203
204class Lock(object):
205  """Lock class"""
206
207  def __init__(self, lock_file, auto=True):
208    self._to_lock = os.path.basename(lock_file)
209    self._lock_file = lock_file
210    self._logger = logger.GetLogger()
211    self._auto = auto
212
213  def NonBlockingLock(self, exclusive, reason=''):
214    with FileLock(self._lock_file) as lock:
215      if lock.exclusive:
216        self._logger.LogError(
217            'Exclusive lock already acquired by %s. Reason: %s' % (lock.owner,
218                                                                   lock.reason))
219        return False
220
221      if exclusive:
222        if lock.counter:
223          self._logger.LogError('Shared lock already acquired')
224          return False
225        lock_file_check = FileCheckName(self._lock_file)
226        fd = OpenLiveCheck(lock_file_check)
227        FileLock.FILE_OPS.append(fd)
228
229        lock.exclusive = True
230        lock.reason = reason
231        lock.owner = getpass.getuser()
232        lock.time = time.time()
233        lock.auto = self._auto
234      else:
235        lock.counter += 1
236    self._logger.LogOutput('Successfully locked: %s' % self._to_lock)
237    return True
238
239  def Unlock(self, exclusive, force=False):
240    with FileLock(self._lock_file) as lock:
241      if not lock.IsLocked():
242        self._logger.LogWarning("Can't unlock unlocked machine!")
243        return True
244
245      if lock.exclusive != exclusive:
246        self._logger.LogError('shared locks must be unlocked with --shared')
247        return False
248
249      if lock.exclusive:
250        if lock.owner != getpass.getuser() and not force:
251          self._logger.LogError("%s can't unlock lock owned by: %s" %
252                                (getpass.getuser(), lock.owner))
253          return False
254        if lock.auto != self._auto:
255          self._logger.LogError("Can't unlock lock with different -a"
256                                ' parameter.')
257          return False
258        lock.exclusive = False
259        lock.reason = ''
260        lock.owner = ''
261
262        if self._auto:
263          del_list = [
264              i for i in FileLock.FILE_OPS
265              if i.name == FileCheckName(self._lock_file)
266          ]
267          for i in del_list:
268            FileLock.FILE_OPS.remove(i)
269          for f in del_list:
270            fcntl.lockf(f, fcntl.LOCK_UN)
271            f.close()
272          del del_list
273          os.remove(FileCheckName(self._lock_file))
274
275      else:
276        lock.counter -= 1
277    return True
278
279
280class Machine(object):
281  """Machine class"""
282
283  LOCKS_DIR = '/google/data/rw/users/mo/mobiletc-prebuild/locks'
284
285  def __init__(self, name, locks_dir=LOCKS_DIR, auto=True):
286    self._name = name
287    self._auto = auto
288    try:
289      self._full_name = socket.gethostbyaddr(name)[0]
290    except socket.error:
291      self._full_name = self._name
292    self._full_name = os.path.join(locks_dir, self._full_name)
293
294  def Lock(self, exclusive=False, reason=''):
295    lock = Lock(self._full_name, self._auto)
296    return lock.NonBlockingLock(exclusive, reason)
297
298  def TryLock(self, timeout=300, exclusive=False, reason=''):
299    locked = False
300    sleep = timeout / 10
301    while True:
302      locked = self.Lock(exclusive, reason)
303      if locked or timeout < 0:
304        break
305      print('Lock not acquired for {0}, wait {1} seconds ...'.format(
306          self._name, sleep))
307      time.sleep(sleep)
308      timeout -= sleep
309    return locked
310
311  def Unlock(self, exclusive=False, ignore_ownership=False):
312    lock = Lock(self._full_name, self._auto)
313    return lock.Unlock(exclusive, ignore_ownership)
314
315
316def Main(argv):
317  """The main function."""
318
319  parser = argparse.ArgumentParser()
320  parser.add_argument(
321      '-r', '--reason', dest='reason', default='', help='The lock reason.')
322  parser.add_argument(
323      '-u',
324      '--unlock',
325      dest='unlock',
326      action='store_true',
327      default=False,
328      help='Use this to unlock.')
329  parser.add_argument(
330      '-l',
331      '--list_locks',
332      dest='list_locks',
333      action='store_true',
334      default=False,
335      help='Use this to list locks.')
336  parser.add_argument(
337      '-f',
338      '--ignore_ownership',
339      dest='ignore_ownership',
340      action='store_true',
341      default=False,
342      help="Use this to force unlock on a lock you don't own.")
343  parser.add_argument(
344      '-s',
345      '--shared',
346      dest='shared',
347      action='store_true',
348      default=False,
349      help='Use this for a shared (non-exclusive) lock.')
350  parser.add_argument(
351      '-d',
352      '--dir',
353      dest='locks_dir',
354      action='store',
355      default=Machine.LOCKS_DIR,
356      help='Use this to set different locks_dir')
357  parser.add_argument('args', nargs='*', help='Machine arg.')
358
359  options = parser.parse_args(argv)
360
361  options.locks_dir = os.path.abspath(options.locks_dir)
362  exclusive = not options.shared
363
364  if not options.list_locks and len(options.args) != 2:
365    logger.GetLogger().LogError(
366        'Either --list_locks or a machine arg is needed.')
367    return 1
368
369  if len(options.args) > 1:
370    machine = Machine(options.args[1], options.locks_dir, auto=False)
371  else:
372    machine = None
373
374  if options.list_locks:
375    FileLock.ListLock('*', options.locks_dir)
376    retval = True
377  elif options.unlock:
378    retval = machine.Unlock(exclusive, options.ignore_ownership)
379  else:
380    retval = machine.Lock(exclusive, options.reason)
381
382  if retval:
383    return 0
384  else:
385    return 1
386
387
388if __name__ == '__main__':
389  sys.exit(Main(sys.argv[1:]))
390