Statistics
| Branch: | Tag: | Revision:

root / lib / utils / io.py @ 0232b768

History | View | Annotate | Download (28.9 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 import pathutils
36
from ganeti.utils import filelock
37

    
38

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

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

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

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

    
57

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

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

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

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

    
75

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

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

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

85
    """
86
    self.st = None
87

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

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

    
94

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

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

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

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

    
115

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

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

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

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

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

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

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

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

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

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

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

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

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

    
209
  # Function result
210
  result = None
211

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

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

    
258
  return result
259

    
260

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

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

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

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

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

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

    
281

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

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

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

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

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

    
300

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

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

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

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

    
325

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

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

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

    
343

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

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

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

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

    
360

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

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

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

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

    
378

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

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

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

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

    
411
      return os.rename(old, new)
412

    
413
    raise
414

    
415

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

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

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

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

    
438
  try:
439
    st = _stat_fn(path)
440

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

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

    
461

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

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

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

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

    
490

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

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

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

    
506

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

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

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

    
516

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

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

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

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

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

    
547
  return backup_name
548

    
549

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

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

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

    
564
  mountpoint = _is_mountpoint(path)
565

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

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

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

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

    
580

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

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

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

    
603

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

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

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

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

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

    
637

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

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

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

    
646

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

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

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

    
660

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

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

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

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

    
689

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

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

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

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

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

    
715

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

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

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

    
727

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

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

736
  """
737
  size = 0
738

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

    
744
  return BytesToMebibyte(size)
745

    
746

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

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

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

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

    
762

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

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

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

    
780
  return _ParsePidFileContents(raw_data)
781

    
782

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

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

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

    
799

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

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

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

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

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

    
828
  return None
829

    
830

    
831
_SSH_KEYS_WITH_TWO_PARTS = frozenset(["ssh-dss", "ssh-rsa"])
832

    
833

    
834
def _SplitSshKey(key):
835
  """Splits a line for SSH's C{authorized_keys} file.
836

837
  If the line has no options (e.g. no C{command="..."}), only the significant
838
  parts, the key type and its hash, are used. Otherwise the whole line is used
839
  (split at whitespace).
840

841
  @type key: string
842
  @param key: Key line
843
  @rtype: tuple
844

845
  """
846
  parts = key.split()
847

    
848
  if parts and parts[0] in _SSH_KEYS_WITH_TWO_PARTS:
849
    # If the key has no options in front of it, we only want the significant
850
    # fields
851
    return (False, parts[:2])
852
  else:
853
    # Can't properly split the line, so use everything
854
    return (True, parts)
855

    
856

    
857
def AddAuthorizedKey(file_obj, key):
858
  """Adds an SSH public key to an authorized_keys file.
859

860
  @type file_obj: str or file handle
861
  @param file_obj: path to authorized_keys file
862
  @type key: str
863
  @param key: string containing key
864

865
  """
866
  key_fields = _SplitSshKey(key)
867

    
868
  if isinstance(file_obj, basestring):
869
    f = open(file_obj, "a+")
870
  else:
871
    f = file_obj
872

    
873
  try:
874
    nl = True
875
    for line in f:
876
      # Ignore whitespace changes
877
      if _SplitSshKey(line) == key_fields:
878
        break
879
      nl = line.endswith("\n")
880
    else:
881
      if not nl:
882
        f.write("\n")
883
      f.write(key.rstrip("\r\n"))
884
      f.write("\n")
885
      f.flush()
886
  finally:
887
    f.close()
888

    
889

    
890
def RemoveAuthorizedKey(file_name, key):
891
  """Removes an SSH public key from an authorized_keys file.
892

893
  @type file_name: str
894
  @param file_name: path to authorized_keys file
895
  @type key: str
896
  @param key: string containing key
897

898
  """
899
  key_fields = _SplitSshKey(key)
900

    
901
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
902
  try:
903
    out = os.fdopen(fd, "w")
904
    try:
905
      f = open(file_name, "r")
906
      try:
907
        for line in f:
908
          # Ignore whitespace changes while comparing lines
909
          if _SplitSshKey(line) != key_fields:
910
            out.write(line)
911

    
912
        out.flush()
913
        os.rename(tmpname, file_name)
914
      finally:
915
        f.close()
916
    finally:
917
      out.close()
918
  except:
919
    RemoveFile(tmpname)
920
    raise
921

    
922

    
923
def DaemonPidFileName(name):
924
  """Compute a ganeti pid file absolute path
925

926
  @type name: str
927
  @param name: the daemon name
928
  @rtype: str
929
  @return: the full path to the pidfile corresponding to the given
930
      daemon name
931

932
  """
933
  return PathJoin(pathutils.RUN_DIR, "%s.pid" % name)
934

    
935

    
936
def WritePidFile(pidfile):
937
  """Write the current process pidfile.
938

939
  @type pidfile: string
940
  @param pidfile: the path to the file to be written
941
  @raise errors.LockError: if the pid file already exists and
942
      points to a live process
943
  @rtype: int
944
  @return: the file descriptor of the lock file; do not close this unless
945
      you want to unlock the pid file
946

947
  """
948
  # We don't rename nor truncate the file to not drop locks under
949
  # existing processes
950
  fd_pidfile = os.open(pidfile, os.O_RDWR | os.O_CREAT, 0600)
951

    
952
  # Lock the PID file (and fail if not possible to do so). Any code
953
  # wanting to send a signal to the daemon should try to lock the PID
954
  # file before reading it. If acquiring the lock succeeds, the daemon is
955
  # no longer running and the signal should not be sent.
956
  try:
957
    filelock.LockFile(fd_pidfile)
958
  except errors.LockError:
959
    msg = ["PID file '%s' is already locked by another process" % pidfile]
960
    # Try to read PID file
961
    pid = _ParsePidFileContents(os.read(fd_pidfile, 100))
962
    if pid > 0:
963
      msg.append(", PID read from file is %s" % pid)
964
    raise errors.PidFileLockError("".join(msg))
965

    
966
  os.write(fd_pidfile, "%d\n" % os.getpid())
967

    
968
  return fd_pidfile
969

    
970

    
971
def ReadWatcherPauseFile(filename, now=None, remove_after=3600):
972
  """Reads the watcher pause file.
973

974
  @type filename: string
975
  @param filename: Path to watcher pause file
976
  @type now: None, float or int
977
  @param now: Current time as Unix timestamp
978
  @type remove_after: int
979
  @param remove_after: Remove watcher pause file after specified amount of
980
    seconds past the pause end time
981

982
  """
983
  if now is None:
984
    now = time.time()
985

    
986
  try:
987
    value = ReadFile(filename)
988
  except IOError, err:
989
    if err.errno != errno.ENOENT:
990
      raise
991
    value = None
992

    
993
  if value is not None:
994
    try:
995
      value = int(value)
996
    except ValueError:
997
      logging.warning(("Watcher pause file (%s) contains invalid value,"
998
                       " removing it"), filename)
999
      RemoveFile(filename)
1000
      value = None
1001

    
1002
    if value is not None:
1003
      # Remove file if it's outdated
1004
      if now > (value + remove_after):
1005
        RemoveFile(filename)
1006
        value = None
1007

    
1008
      elif now > value:
1009
        value = None
1010

    
1011
  return value
1012

    
1013

    
1014
def NewUUID():
1015
  """Returns a random UUID.
1016

1017
  @note: This is a Linux-specific method as it uses the /proc
1018
      filesystem.
1019
  @rtype: str
1020

1021
  """
1022
  return ReadFile(_RANDOM_UUID_FILE, size=128).rstrip("\n")
1023

    
1024

    
1025
class TemporaryFileManager(object):
1026
  """Stores the list of files to be deleted and removes them on demand.
1027

1028
  """
1029

    
1030
  def __init__(self):
1031
    self._files = []
1032

    
1033
  def __del__(self):
1034
    self.Cleanup()
1035

    
1036
  def Add(self, filename):
1037
    """Add file to list of files to be deleted.
1038

1039
    @type filename: string
1040
    @param filename: path to filename to be added
1041

1042
    """
1043
    self._files.append(filename)
1044

    
1045
  def Remove(self, filename):
1046
    """Remove file from list of files to be deleted.
1047

1048
    @type filename: string
1049
    @param filename: path to filename to be deleted
1050

1051
    """
1052
    self._files.remove(filename)
1053

    
1054
  def Cleanup(self):
1055
    """Delete all files marked for deletion
1056

1057
    """
1058
    while self._files:
1059
      RemoveFile(self._files.pop())