Statistics
| Branch: | Tag: | Revision:

root / lib / utils / io.py @ 5ae4945a

History | View | Annotate | Download (28.2 kB)

1
#
2
#
3

    
4
# Copyright (C) 2006, 2007, 2010, 2011, 2012 Google Inc.
5
#
6
# This program is free software; you can redistribute it and/or modify
7
# it under the terms of the GNU General Public License as published by
8
# the Free Software Foundation; either version 2 of the License, or
9
# (at your option) any later version.
10
#
11
# This program is distributed in the hope that it will be useful, but
12
# WITHOUT ANY WARRANTY; without even the implied warranty of
13
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14
# General Public License for more details.
15
#
16
# You should have received a copy of the GNU General Public License
17
# along with this program; if not, write to the Free Software
18
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
19
# 02110-1301, USA.
20

    
21
"""Utility functions for I/O.
22

23
"""
24

    
25
import os
26
import logging
27
import shutil
28
import tempfile
29
import errno
30
import time
31
import stat
32

    
33
from ganeti import errors
34
from ganeti import constants
35
from ganeti.utils import filelock
36

    
37

    
38
#: Path generating random UUID
39
_RANDOM_UUID_FILE = "/proc/sys/kernel/random/uuid"
40

    
41
#: Directory used by fsck(8) to store recovered data, usually at a file
42
#: system's root directory
43
_LOST_AND_FOUND = "lost+found"
44

    
45
# Possible values for keep_perms in WriteFile()
46
KP_NEVER = 0
47
KP_ALWAYS = 1
48
KP_IF_EXISTS = 2
49

    
50
KEEP_PERMS_VALUES = [
51
  KP_NEVER,
52
  KP_ALWAYS,
53
  KP_IF_EXISTS,
54
  ]
55

    
56

    
57
def ErrnoOrStr(err):
58
  """Format an EnvironmentError exception.
59

60
  If the L{err} argument has an errno attribute, it will be looked up
61
  and converted into a textual C{E...} description. Otherwise the
62
  string representation of the error will be returned.
63

64
  @type err: L{EnvironmentError}
65
  @param err: the exception to format
66

67
  """
68
  if hasattr(err, "errno"):
69
    detail = errno.errorcode[err.errno]
70
  else:
71
    detail = str(err)
72
  return detail
73

    
74

    
75
class FileStatHelper:
76
  """Helper to store file handle's C{fstat}.
77

78
  Useful in combination with L{ReadFile}'s C{preread} parameter.
79

80
  """
81
  def __init__(self):
82
    """Initializes this class.
83

84
    """
85
    self.st = None
86

    
87
  def __call__(self, fh):
88
    """Calls C{fstat} on file handle.
89

90
    """
91
    self.st = os.fstat(fh.fileno())
92

    
93

    
94
def ReadFile(file_name, size=-1, preread=None):
95
  """Reads a file.
96

97
  @type size: int
98
  @param size: Read at most size bytes (if negative, entire file)
99
  @type preread: callable receiving file handle as single parameter
100
  @param preread: Function called before file is read
101
  @rtype: str
102
  @return: the (possibly partial) content of the file
103

104
  """
105
  f = open(file_name, "r")
106
  try:
107
    if preread:
108
      preread(f)
109

    
110
    return f.read(size)
111
  finally:
112
    f.close()
113

    
114

    
115
def WriteFile(file_name, fn=None, data=None,
116
              mode=None, uid=-1, gid=-1,
117
              atime=None, mtime=None, close=True,
118
              dry_run=False, backup=False,
119
              prewrite=None, postwrite=None, keep_perms=KP_NEVER):
120
  """(Over)write a file atomically.
121

122
  The file_name and either fn (a function taking one argument, the
123
  file descriptor, and which should write the data to it) or data (the
124
  contents of the file) must be passed. The other arguments are
125
  optional and allow setting the file mode, owner and group, and the
126
  mtime/atime of the file.
127

128
  If the function doesn't raise an exception, it has succeeded and the
129
  target file has the new contents. If the function has raised an
130
  exception, an existing target file should be unmodified and the
131
  temporary file should be removed.
132

133
  @type file_name: str
134
  @param file_name: the target filename
135
  @type fn: callable
136
  @param fn: content writing function, called with
137
      file descriptor as parameter
138
  @type data: str
139
  @param data: contents of the file
140
  @type mode: int
141
  @param mode: file mode
142
  @type uid: int
143
  @param uid: the owner of the file
144
  @type gid: int
145
  @param gid: the group of the file
146
  @type atime: int
147
  @param atime: a custom access time to be set on the file
148
  @type mtime: int
149
  @param mtime: a custom modification time to be set on the file
150
  @type close: boolean
151
  @param close: whether to close file after writing it
152
  @type prewrite: callable
153
  @param prewrite: function to be called before writing content
154
  @type postwrite: callable
155
  @param postwrite: function to be called after writing content
156
  @type keep_perms: members of L{KEEP_PERMS_VALUES}
157
  @param keep_perms: if L{KP_NEVER} (default), owner, group, and mode are
158
      taken from the other parameters; if L{KP_ALWAYS}, owner, group, and
159
      mode are copied from the existing file; if L{KP_IF_EXISTS}, owner,
160
      group, and mode are taken from the file, and if the file doesn't
161
      exist, they are taken from the other parameters. It is an error to
162
      pass L{KP_ALWAYS} when the file doesn't exist or when C{uid}, C{gid},
163
      or C{mode} are set to non-default values.
164

165
  @rtype: None or int
166
  @return: None if the 'close' parameter evaluates to True,
167
      otherwise the file descriptor
168

169
  @raise errors.ProgrammerError: if any of the arguments are not valid
170

171
  """
172
  if not os.path.isabs(file_name):
173
    raise errors.ProgrammerError("Path passed to WriteFile is not"
174
                                 " absolute: '%s'" % file_name)
175

    
176
  if [fn, data].count(None) != 1:
177
    raise errors.ProgrammerError("fn or data required")
178

    
179
  if [atime, mtime].count(None) == 1:
180
    raise errors.ProgrammerError("Both atime and mtime must be either"
181
                                 " set or None")
182

    
183
  if not keep_perms in KEEP_PERMS_VALUES:
184
    raise errors.ProgrammerError("Invalid value for keep_perms: %s" %
185
                                 keep_perms)
186
  if keep_perms == KP_ALWAYS and (uid != -1 or gid != -1 or mode is not None):
187
    raise errors.ProgrammerError("When keep_perms==KP_ALWAYS, 'uid', 'gid',"
188
                                 " and 'mode' cannot be set")
189

    
190
  if backup and not dry_run and os.path.isfile(file_name):
191
    CreateBackup(file_name)
192

    
193
  if keep_perms == KP_ALWAYS or keep_perms == KP_IF_EXISTS:
194
    # os.stat() raises an exception if the file doesn't exist
195
    try:
196
      file_stat = os.stat(file_name)
197
      mode = stat.S_IMODE(file_stat.st_mode)
198
      uid = file_stat.st_uid
199
      gid = file_stat.st_gid
200
    except OSError:
201
      if keep_perms == KP_ALWAYS:
202
        raise
203
      # else: if keeep_perms == KP_IF_EXISTS it's ok if the file doesn't exist
204

    
205
  # Whether temporary file needs to be removed (e.g. if any error occurs)
206
  do_remove = True
207

    
208
  # Function result
209
  result = None
210

    
211
  (dir_name, base_name) = os.path.split(file_name)
212
  (fd, new_name) = tempfile.mkstemp(suffix=".new", prefix=base_name,
213
                                    dir=dir_name)
214
  try:
215
    try:
216
      if uid != -1 or gid != -1:
217
        os.chown(new_name, uid, gid)
218
      if mode:
219
        os.chmod(new_name, mode)
220
      if callable(prewrite):
221
        prewrite(fd)
222
      if data is not None:
223
        if isinstance(data, unicode):
224
          data = data.encode()
225
        assert isinstance(data, str)
226
        to_write = len(data)
227
        offset = 0
228
        while offset < to_write:
229
          written = os.write(fd, buffer(data, offset))
230
          assert written >= 0
231
          assert written <= to_write - offset
232
          offset += written
233
        assert offset == to_write
234
      else:
235
        fn(fd)
236
      if callable(postwrite):
237
        postwrite(fd)
238
      os.fsync(fd)
239
      if atime is not None and mtime is not None:
240
        os.utime(new_name, (atime, mtime))
241
    finally:
242
      # Close file unless the file descriptor should be returned
243
      if close:
244
        os.close(fd)
245
      else:
246
        result = fd
247

    
248
    # Rename file to destination name
249
    if not dry_run:
250
      os.rename(new_name, file_name)
251
      # Successful, no need to remove anymore
252
      do_remove = False
253
  finally:
254
    if do_remove:
255
      RemoveFile(new_name)
256

    
257
  return result
258

    
259

    
260
def GetFileID(path=None, fd=None):
261
  """Returns the file 'id', i.e. the dev/inode and mtime information.
262

263
  Either the path to the file or the fd must be given.
264

265
  @param path: the file path
266
  @param fd: a file descriptor
267
  @return: a tuple of (device number, inode number, mtime)
268

269
  """
270
  if [path, fd].count(None) != 1:
271
    raise errors.ProgrammerError("One and only one of fd/path must be given")
272

    
273
  if fd is None:
274
    st = os.stat(path)
275
  else:
276
    st = os.fstat(fd)
277

    
278
  return (st.st_dev, st.st_ino, st.st_mtime)
279

    
280

    
281
def VerifyFileID(fi_disk, fi_ours):
282
  """Verifies that two file IDs are matching.
283

284
  Differences in the inode/device are not accepted, but and older
285
  timestamp for fi_disk is accepted.
286

287
  @param fi_disk: tuple (dev, inode, mtime) representing the actual
288
      file data
289
  @param fi_ours: tuple (dev, inode, mtime) representing the last
290
      written file data
291
  @rtype: boolean
292

293
  """
294
  (d1, i1, m1) = fi_disk
295
  (d2, i2, m2) = fi_ours
296

    
297
  return (d1, i1) == (d2, i2) and m1 <= m2
298

    
299

    
300
def SafeWriteFile(file_name, file_id, **kwargs):
301
  """Wraper over L{WriteFile} that locks the target file.
302

303
  By keeping the target file locked during WriteFile, we ensure that
304
  cooperating writers will safely serialise access to the file.
305

306
  @type file_name: str
307
  @param file_name: the target filename
308
  @type file_id: tuple
309
  @param file_id: a result from L{GetFileID}
310

311
  """
312
  fd = os.open(file_name, os.O_RDONLY | os.O_CREAT)
313
  try:
314
    filelock.LockFile(fd)
315
    if file_id is not None:
316
      disk_id = GetFileID(fd=fd)
317
      if not VerifyFileID(disk_id, file_id):
318
        raise errors.LockError("Cannot overwrite file %s, it has been modified"
319
                               " since last written" % file_name)
320
    return WriteFile(file_name, **kwargs)
321
  finally:
322
    os.close(fd)
323

    
324

    
325
def ReadOneLineFile(file_name, strict=False):
326
  """Return the first non-empty line from a file.
327

328
  @type strict: boolean
329
  @param strict: if True, abort if the file has more than one
330
      non-empty line
331

332
  """
333
  file_lines = ReadFile(file_name).splitlines()
334
  full_lines = filter(bool, file_lines)
335
  if not file_lines or not full_lines:
336
    raise errors.GenericError("No data in one-liner file %s" % file_name)
337
  elif strict and len(full_lines) > 1:
338
    raise errors.GenericError("Too many lines in one-liner file %s" %
339
                              file_name)
340
  return full_lines[0]
341

    
342

    
343
def RemoveFile(filename):
344
  """Remove a file ignoring some errors.
345

346
  Remove a file, ignoring non-existing ones or directories. Other
347
  errors are passed.
348

349
  @type filename: str
350
  @param filename: the file to be removed
351

352
  """
353
  try:
354
    os.unlink(filename)
355
  except OSError, err:
356
    if err.errno not in (errno.ENOENT, errno.EISDIR):
357
      raise
358

    
359

    
360
def RemoveDir(dirname):
361
  """Remove an empty directory.
362

363
  Remove a directory, ignoring non-existing ones.
364
  Other errors are passed. This includes the case,
365
  where the directory is not empty, so it can't be removed.
366

367
  @type dirname: str
368
  @param dirname: the empty directory to be removed
369

370
  """
371
  try:
372
    os.rmdir(dirname)
373
  except OSError, err:
374
    if err.errno != errno.ENOENT:
375
      raise
376

    
377

    
378
def RenameFile(old, new, mkdir=False, mkdir_mode=0750, dir_uid=None,
379
               dir_gid=None):
380
  """Renames a file.
381

382
  This just creates the very least directory if it does not exist and C{mkdir}
383
  is set to true.
384

385
  @type old: string
386
  @param old: Original path
387
  @type new: string
388
  @param new: New path
389
  @type mkdir: bool
390
  @param mkdir: Whether to create target directory if it doesn't exist
391
  @type mkdir_mode: int
392
  @param mkdir_mode: Mode for newly created directories
393
  @type dir_uid: int
394
  @param dir_uid: The uid for the (if fresh created) dir
395
  @type dir_gid: int
396
  @param dir_gid: The gid for the (if fresh created) dir
397

398
  """
399
  try:
400
    return os.rename(old, new)
401
  except OSError, err:
402
    # In at least one use case of this function, the job queue, directory
403
    # creation is very rare. Checking for the directory before renaming is not
404
    # as efficient.
405
    if mkdir and err.errno == errno.ENOENT:
406
      # Create directory and try again
407
      dir_path = os.path.dirname(new)
408
      MakeDirWithPerm(dir_path, mkdir_mode, dir_uid, dir_gid)
409

    
410
      return os.rename(old, new)
411

    
412
    raise
413

    
414

    
415
def EnforcePermission(path, mode, uid=None, gid=None, must_exist=True,
416
                      _chmod_fn=os.chmod, _chown_fn=os.chown, _stat_fn=os.stat):
417
  """Enforces that given path has given permissions.
418

419
  @param path: The path to the file
420
  @param mode: The mode of the file
421
  @param uid: The uid of the owner of this file
422
  @param gid: The gid of the owner of this file
423
  @param must_exist: Specifies if non-existance of path will be an error
424
  @param _chmod_fn: chmod function to use (unittest only)
425
  @param _chown_fn: chown function to use (unittest only)
426

427
  """
428
  logging.debug("Checking %s", path)
429

    
430
  # chown takes -1 if you want to keep one part of the ownership, however
431
  # None is Python standard for that. So we remap them here.
432
  if uid is None:
433
    uid = -1
434
  if gid is None:
435
    gid = -1
436

    
437
  try:
438
    st = _stat_fn(path)
439

    
440
    fmode = stat.S_IMODE(st[stat.ST_MODE])
441
    if fmode != mode:
442
      logging.debug("Changing mode of %s from %#o to %#o", path, fmode, mode)
443
      _chmod_fn(path, mode)
444

    
445
    if max(uid, gid) > -1:
446
      fuid = st[stat.ST_UID]
447
      fgid = st[stat.ST_GID]
448
      if fuid != uid or fgid != gid:
449
        logging.debug("Changing owner of %s from UID %s/GID %s to"
450
                      " UID %s/GID %s", path, fuid, fgid, uid, gid)
451
        _chown_fn(path, uid, gid)
452
  except EnvironmentError, err:
453
    if err.errno == errno.ENOENT:
454
      if must_exist:
455
        raise errors.GenericError("Path %s should exist, but does not" % path)
456
    else:
457
      raise errors.GenericError("Error while changing permissions on %s: %s" %
458
                                (path, err))
459

    
460

    
461
def MakeDirWithPerm(path, mode, uid, gid, _lstat_fn=os.lstat,
462
                    _mkdir_fn=os.mkdir, _perm_fn=EnforcePermission):
463
  """Enforces that given path is a dir and has given mode, uid and gid set.
464

465
  @param path: The path to the file
466
  @param mode: The mode of the file
467
  @param uid: The uid of the owner of this file
468
  @param gid: The gid of the owner of this file
469
  @param _lstat_fn: Stat function to use (unittest only)
470
  @param _mkdir_fn: mkdir function to use (unittest only)
471
  @param _perm_fn: permission setter function to use (unittest only)
472

473
  """
474
  logging.debug("Checking directory %s", path)
475
  try:
476
    # We don't want to follow symlinks
477
    st = _lstat_fn(path)
478
  except EnvironmentError, err:
479
    if err.errno != errno.ENOENT:
480
      raise errors.GenericError("stat(2) on %s failed: %s" % (path, err))
481
    _mkdir_fn(path)
482
  else:
483
    if not stat.S_ISDIR(st[stat.ST_MODE]):
484
      raise errors.GenericError(("Path %s is expected to be a directory, but "
485
                                 "isn't") % path)
486

    
487
  _perm_fn(path, mode, uid=uid, gid=gid)
488

    
489

    
490
def Makedirs(path, mode=0750):
491
  """Super-mkdir; create a leaf directory and all intermediate ones.
492

493
  This is a wrapper around C{os.makedirs} adding error handling not implemented
494
  before Python 2.5.
495

496
  """
497
  try:
498
    os.makedirs(path, mode)
499
  except OSError, err:
500
    # Ignore EEXIST. This is only handled in os.makedirs as included in
501
    # Python 2.5 and above.
502
    if err.errno != errno.EEXIST or not os.path.exists(path):
503
      raise
504

    
505

    
506
def TimestampForFilename():
507
  """Returns the current time formatted for filenames.
508

509
  The format doesn't contain colons as some shells and applications treat them
510
  as separators. Uses the local timezone.
511

512
  """
513
  return time.strftime("%Y-%m-%d_%H_%M_%S")
514

    
515

    
516
def CreateBackup(file_name):
517
  """Creates a backup of a file.
518

519
  @type file_name: str
520
  @param file_name: file to be backed up
521
  @rtype: str
522
  @return: the path to the newly created backup
523
  @raise errors.ProgrammerError: for invalid file names
524

525
  """
526
  if not os.path.isfile(file_name):
527
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
528
                                 file_name)
529

    
530
  prefix = ("%s.backup-%s." %
531
            (os.path.basename(file_name), TimestampForFilename()))
532
  dir_name = os.path.dirname(file_name)
533

    
534
  fsrc = open(file_name, "rb")
535
  try:
536
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
537
    fdst = os.fdopen(fd, "wb")
538
    try:
539
      logging.debug("Backing up %s at %s", file_name, backup_name)
540
      shutil.copyfileobj(fsrc, fdst)
541
    finally:
542
      fdst.close()
543
  finally:
544
    fsrc.close()
545

    
546
  return backup_name
547

    
548

    
549
def ListVisibleFiles(path, _is_mountpoint=os.path.ismount):
550
  """Returns a list of visible files in a directory.
551

552
  @type path: str
553
  @param path: the directory to enumerate
554
  @rtype: list
555
  @return: the list of all files not starting with a dot
556
  @raise ProgrammerError: if L{path} is not an absolue and normalized path
557

558
  """
559
  if not IsNormAbsPath(path):
560
    raise errors.ProgrammerError("Path passed to ListVisibleFiles is not"
561
                                 " absolute/normalized: '%s'" % path)
562

    
563
  mountpoint = _is_mountpoint(path)
564

    
565
  def fn(name):
566
    """File name filter.
567

568
    Ignores files starting with a dot (".") as by Unix convention they're
569
    considered hidden. The "lost+found" directory found at the root of some
570
    filesystems is also hidden.
571

572
    """
573
    return not (name.startswith(".") or
574
                (mountpoint and name == _LOST_AND_FOUND and
575
                 os.path.isdir(os.path.join(path, name))))
576

    
577
  return filter(fn, os.listdir(path))
578

    
579

    
580
def EnsureDirs(dirs):
581
  """Make required directories, if they don't exist.
582

583
  @param dirs: list of tuples (dir_name, dir_mode)
584
  @type dirs: list of (string, integer)
585

586
  """
587
  for dir_name, dir_mode in dirs:
588
    try:
589
      os.mkdir(dir_name, dir_mode)
590
    except EnvironmentError, err:
591
      if err.errno != errno.EEXIST:
592
        raise errors.GenericError("Cannot create needed directory"
593
                                  " '%s': %s" % (dir_name, err))
594
    try:
595
      os.chmod(dir_name, dir_mode)
596
    except EnvironmentError, err:
597
      raise errors.GenericError("Cannot change directory permissions on"
598
                                " '%s': %s" % (dir_name, err))
599
    if not os.path.isdir(dir_name):
600
      raise errors.GenericError("%s is not a directory" % dir_name)
601

    
602

    
603
def FindFile(name, search_path, test=os.path.exists):
604
  """Look for a filesystem object in a given path.
605

606
  This is an abstract method to search for filesystem object (files,
607
  dirs) under a given search path.
608

609
  @type name: str
610
  @param name: the name to look for
611
  @type search_path: str
612
  @param search_path: location to start at
613
  @type test: callable
614
  @param test: a function taking one argument that should return True
615
      if the a given object is valid; the default value is
616
      os.path.exists, causing only existing files to be returned
617
  @rtype: str or None
618
  @return: full path to the object if found, None otherwise
619

620
  """
621
  # validate the filename mask
622
  if constants.EXT_PLUGIN_MASK.match(name) is None:
623
    logging.critical("Invalid value passed for external script name: '%s'",
624
                     name)
625
    return None
626

    
627
  for dir_name in search_path:
628
    # FIXME: investigate switch to PathJoin
629
    item_name = os.path.sep.join([dir_name, name])
630
    # check the user test and that we're indeed resolving to the given
631
    # basename
632
    if test(item_name) and os.path.basename(item_name) == name:
633
      return item_name
634
  return None
635

    
636

    
637
def IsNormAbsPath(path):
638
  """Check whether a path is absolute and also normalized
639

640
  This avoids things like /dir/../../other/path to be valid.
641

642
  """
643
  return os.path.normpath(path) == path and os.path.isabs(path)
644

    
645

    
646
def IsBelowDir(root, other_path):
647
  """Check whether a path is below a root dir.
648

649
  This works around the nasty byte-byte comparisation of commonprefix.
650

651
  """
652
  if not (os.path.isabs(root) and os.path.isabs(other_path)):
653
    raise ValueError("Provided paths '%s' and '%s' are not absolute" %
654
                     (root, other_path))
655
  prepared_root = "%s%s" % (os.path.normpath(root), os.sep)
656
  return os.path.commonprefix([prepared_root,
657
                               os.path.normpath(other_path)]) == prepared_root
658

    
659

    
660
def PathJoin(*args):
661
  """Safe-join a list of path components.
662

663
  Requirements:
664
      - the first argument must be an absolute path
665
      - no component in the path must have backtracking (e.g. /../),
666
        since we check for normalization at the end
667

668
  @param args: the path components to be joined
669
  @raise ValueError: for invalid paths
670

671
  """
672
  # ensure we're having at least one path passed in
673
  assert args
674
  # ensure the first component is an absolute and normalized path name
675
  root = args[0]
676
  if not IsNormAbsPath(root):
677
    raise ValueError("Invalid parameter to PathJoin: '%s'" % str(args[0]))
678
  result = os.path.join(*args)
679
  # ensure that the whole path is normalized
680
  if not IsNormAbsPath(result):
681
    raise ValueError("Invalid parameters to PathJoin: '%s'" % str(args))
682
  # check that we're still under the original prefix
683
  if not IsBelowDir(root, result):
684
    raise ValueError("Error: path joining resulted in different prefix"
685
                     " (%s != %s)" % (result, root))
686
  return result
687

    
688

    
689
def TailFile(fname, lines=20):
690
  """Return the last lines from a file.
691

692
  @note: this function will only read and parse the last 4KB of
693
      the file; if the lines are very long, it could be that less
694
      than the requested number of lines are returned
695

696
  @param fname: the file name
697
  @type lines: int
698
  @param lines: the (maximum) number of lines to return
699

700
  """
701
  fd = open(fname, "r")
702
  try:
703
    fd.seek(0, 2)
704
    pos = fd.tell()
705
    pos = max(0, pos - 4096)
706
    fd.seek(pos, 0)
707
    raw_data = fd.read()
708
  finally:
709
    fd.close()
710

    
711
  rows = raw_data.splitlines()
712
  return rows[-lines:]
713

    
714

    
715
def BytesToMebibyte(value):
716
  """Converts bytes to mebibytes.
717

718
  @type value: int
719
  @param value: Value in bytes
720
  @rtype: int
721
  @return: Value in mebibytes
722

723
  """
724
  return int(round(value / (1024.0 * 1024.0), 0))
725

    
726

    
727
def CalculateDirectorySize(path):
728
  """Calculates the size of a directory recursively.
729

730
  @type path: string
731
  @param path: Path to directory
732
  @rtype: int
733
  @return: Size in mebibytes
734

735
  """
736
  size = 0
737

    
738
  for (curpath, _, files) in os.walk(path):
739
    for filename in files:
740
      st = os.lstat(PathJoin(curpath, filename))
741
      size += st.st_size
742

    
743
  return BytesToMebibyte(size)
744

    
745

    
746
def GetFilesystemStats(path):
747
  """Returns the total and free space on a filesystem.
748

749
  @type path: string
750
  @param path: Path on filesystem to be examined
751
  @rtype: int
752
  @return: tuple of (Total space, Free space) in mebibytes
753

754
  """
755
  st = os.statvfs(path)
756

    
757
  fsize = BytesToMebibyte(st.f_bavail * st.f_frsize)
758
  tsize = BytesToMebibyte(st.f_blocks * st.f_frsize)
759
  return (tsize, fsize)
760

    
761

    
762
def ReadPidFile(pidfile):
763
  """Read a pid from a file.
764

765
  @type  pidfile: string
766
  @param pidfile: path to the file containing the pid
767
  @rtype: int
768
  @return: The process id, if the file exists and contains a valid PID,
769
           otherwise 0
770

771
  """
772
  try:
773
    raw_data = ReadOneLineFile(pidfile)
774
  except EnvironmentError, err:
775
    if err.errno != errno.ENOENT:
776
      logging.exception("Can't read pid file")
777
    return 0
778

    
779
  return _ParsePidFileContents(raw_data)
780

    
781

    
782
def _ParsePidFileContents(data):
783
  """Tries to extract a process ID from a PID file's content.
784

785
  @type data: string
786
  @rtype: int
787
  @return: Zero if nothing could be read, PID otherwise
788

789
  """
790
  try:
791
    pid = int(data)
792
  except (TypeError, ValueError):
793
    logging.info("Can't parse pid file contents", exc_info=True)
794
    return 0
795
  else:
796
    return pid
797

    
798

    
799
def ReadLockedPidFile(path):
800
  """Reads a locked PID file.
801

802
  This can be used together with L{utils.process.StartDaemon}.
803

804
  @type path: string
805
  @param path: Path to PID file
806
  @return: PID as integer or, if file was unlocked or couldn't be opened, None
807

808
  """
809
  try:
810
    fd = os.open(path, os.O_RDONLY)
811
  except EnvironmentError, err:
812
    if err.errno == errno.ENOENT:
813
      # PID file doesn't exist
814
      return None
815
    raise
816

    
817
  try:
818
    try:
819
      # Try to acquire lock
820
      filelock.LockFile(fd)
821
    except errors.LockError:
822
      # Couldn't lock, daemon is running
823
      return int(os.read(fd, 100))
824
  finally:
825
    os.close(fd)
826

    
827
  return None
828

    
829

    
830
def AddAuthorizedKey(file_obj, key):
831
  """Adds an SSH public key to an authorized_keys file.
832

833
  @type file_obj: str or file handle
834
  @param file_obj: path to authorized_keys file
835
  @type key: str
836
  @param key: string containing key
837

838
  """
839
  key_fields = key.split()
840

    
841
  if isinstance(file_obj, basestring):
842
    f = open(file_obj, "a+")
843
  else:
844
    f = file_obj
845

    
846
  try:
847
    nl = True
848
    for line in f:
849
      # Ignore whitespace changes
850
      if line.split() == key_fields:
851
        break
852
      nl = line.endswith("\n")
853
    else:
854
      if not nl:
855
        f.write("\n")
856
      f.write(key.rstrip("\r\n"))
857
      f.write("\n")
858
      f.flush()
859
  finally:
860
    f.close()
861

    
862

    
863
def RemoveAuthorizedKey(file_name, key):
864
  """Removes an SSH public key from an authorized_keys file.
865

866
  @type file_name: str
867
  @param file_name: path to authorized_keys file
868
  @type key: str
869
  @param key: string containing key
870

871
  """
872
  key_fields = key.split()
873

    
874
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
875
  try:
876
    out = os.fdopen(fd, "w")
877
    try:
878
      f = open(file_name, "r")
879
      try:
880
        for line in f:
881
          # Ignore whitespace changes while comparing lines
882
          if line.split() != key_fields:
883
            out.write(line)
884

    
885
        out.flush()
886
        os.rename(tmpname, file_name)
887
      finally:
888
        f.close()
889
    finally:
890
      out.close()
891
  except:
892
    RemoveFile(tmpname)
893
    raise
894

    
895

    
896
def DaemonPidFileName(name):
897
  """Compute a ganeti pid file absolute path
898

899
  @type name: str
900
  @param name: the daemon name
901
  @rtype: str
902
  @return: the full path to the pidfile corresponding to the given
903
      daemon name
904

905
  """
906
  return PathJoin(constants.RUN_GANETI_DIR, "%s.pid" % name)
907

    
908

    
909
def WritePidFile(pidfile):
910
  """Write the current process pidfile.
911

912
  @type pidfile: string
913
  @param pidfile: the path to the file to be written
914
  @raise errors.LockError: if the pid file already exists and
915
      points to a live process
916
  @rtype: int
917
  @return: the file descriptor of the lock file; do not close this unless
918
      you want to unlock the pid file
919

920
  """
921
  # We don't rename nor truncate the file to not drop locks under
922
  # existing processes
923
  fd_pidfile = os.open(pidfile, os.O_RDWR | os.O_CREAT, 0600)
924

    
925
  # Lock the PID file (and fail if not possible to do so). Any code
926
  # wanting to send a signal to the daemon should try to lock the PID
927
  # file before reading it. If acquiring the lock succeeds, the daemon is
928
  # no longer running and the signal should not be sent.
929
  try:
930
    filelock.LockFile(fd_pidfile)
931
  except errors.LockError:
932
    msg = ["PID file '%s' is already locked by another process" % pidfile]
933
    # Try to read PID file
934
    pid = _ParsePidFileContents(os.read(fd_pidfile, 100))
935
    if pid > 0:
936
      msg.append(", PID read from file is %s" % pid)
937
    raise errors.PidFileLockError("".join(msg))
938

    
939
  os.write(fd_pidfile, "%d\n" % os.getpid())
940

    
941
  return fd_pidfile
942

    
943

    
944
def ReadWatcherPauseFile(filename, now=None, remove_after=3600):
945
  """Reads the watcher pause file.
946

947
  @type filename: string
948
  @param filename: Path to watcher pause file
949
  @type now: None, float or int
950
  @param now: Current time as Unix timestamp
951
  @type remove_after: int
952
  @param remove_after: Remove watcher pause file after specified amount of
953
    seconds past the pause end time
954

955
  """
956
  if now is None:
957
    now = time.time()
958

    
959
  try:
960
    value = ReadFile(filename)
961
  except IOError, err:
962
    if err.errno != errno.ENOENT:
963
      raise
964
    value = None
965

    
966
  if value is not None:
967
    try:
968
      value = int(value)
969
    except ValueError:
970
      logging.warning(("Watcher pause file (%s) contains invalid value,"
971
                       " removing it"), filename)
972
      RemoveFile(filename)
973
      value = None
974

    
975
    if value is not None:
976
      # Remove file if it's outdated
977
      if now > (value + remove_after):
978
        RemoveFile(filename)
979
        value = None
980

    
981
      elif now > value:
982
        value = None
983

    
984
  return value
985

    
986

    
987
def NewUUID():
988
  """Returns a random UUID.
989

990
  @note: This is a Linux-specific method as it uses the /proc
991
      filesystem.
992
  @rtype: str
993

994
  """
995
  return ReadFile(_RANDOM_UUID_FILE, size=128).rstrip("\n")
996

    
997

    
998
class TemporaryFileManager(object):
999
  """Stores the list of files to be deleted and removes them on demand.
1000

1001
  """
1002

    
1003
  def __init__(self):
1004
    self._files = []
1005

    
1006
  def __del__(self):
1007
    self.Cleanup()
1008

    
1009
  def Add(self, filename):
1010
    """Add file to list of files to be deleted.
1011

1012
    @type filename: string
1013
    @param filename: path to filename to be added
1014

1015
    """
1016
    self._files.append(filename)
1017

    
1018
  def Remove(self, filename):
1019
    """Remove file from list of files to be deleted.
1020

1021
    @type filename: string
1022
    @param filename: path to filename to be deleted
1023

1024
    """
1025
    self._files.remove(filename)
1026

    
1027
  def Cleanup(self):
1028
    """Delete all files marked for deletion
1029

1030
    """
1031
    while self._files:
1032
      RemoveFile(self._files.pop())