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