Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ f7414041

History | View | Annotate | Download (25.2 kB)

1
#
2
#
3

    
4
# Copyright (C) 2006, 2007 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

    
22
"""Ganeti small utilities
23

24
"""
25

    
26

    
27
import sys
28
import os
29
import sha
30
import time
31
import subprocess
32
import re
33
import socket
34
import tempfile
35
import shutil
36
import errno
37
import pwd
38
import itertools
39

    
40
from ganeti import logger
41
from ganeti import errors
42

    
43

    
44
_locksheld = []
45
_re_shell_unquoted = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')
46

    
47
class RunResult(object):
48
  """Simple class for holding the result of running external programs.
49

50
  Instance variables:
51
    exit_code: the exit code of the program, or None (if the program
52
               didn't exit())
53
    signal: numeric signal that caused the program to finish, or None
54
            (if the program wasn't terminated by a signal)
55
    stdout: the standard output of the program
56
    stderr: the standard error of the program
57
    failed: a Boolean value which is True in case the program was
58
            terminated by a signal or exited with a non-zero exit code
59
    fail_reason: a string detailing the termination reason
60

61
  """
62
  __slots__ = ["exit_code", "signal", "stdout", "stderr",
63
               "failed", "fail_reason", "cmd"]
64

    
65

    
66
  def __init__(self, exit_code, signal, stdout, stderr, cmd):
67
    self.cmd = cmd
68
    self.exit_code = exit_code
69
    self.signal = signal
70
    self.stdout = stdout
71
    self.stderr = stderr
72
    self.failed = (signal is not None or exit_code != 0)
73

    
74
    if self.signal is not None:
75
      self.fail_reason = "terminated by signal %s" % self.signal
76
    elif self.exit_code is not None:
77
      self.fail_reason = "exited with exit code %s" % self.exit_code
78
    else:
79
      self.fail_reason = "unable to determine termination reason"
80

    
81
  def _GetOutput(self):
82
    """Returns the combined stdout and stderr for easier usage.
83

84
    """
85
    return self.stdout + self.stderr
86

    
87
  output = property(_GetOutput, None, None, "Return full output")
88

    
89

    
90
def _GetLockFile(subsystem):
91
  """Compute the file name for a given lock name."""
92
  return "/var/lock/ganeti_lock_%s" % subsystem
93

    
94

    
95
def Lock(name, max_retries=None, debug=False):
96
  """Lock a given subsystem.
97

98
  In case the lock is already held by an alive process, the function
99
  will sleep indefintely and poll with a one second interval.
100

101
  When the optional integer argument 'max_retries' is passed with a
102
  non-zero value, the function will sleep only for this number of
103
  times, and then it will will raise a LockError if the lock can't be
104
  acquired. Passing in a negative number will cause only one try to
105
  get the lock. Passing a positive number will make the function retry
106
  for approximately that number of seconds.
107

108
  """
109
  lockfile = _GetLockFile(name)
110

    
111
  if name in _locksheld:
112
    raise errors.LockError('Lock "%s" already held!' % (name,))
113

    
114
  errcount = 0
115

    
116
  retries = 0
117
  while True:
118
    try:
119
      fd = os.open(lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_SYNC)
120
      break
121
    except OSError, creat_err:
122
      if creat_err.errno != errno.EEXIST:
123
        raise errors.LockError("Can't create the lock file. Error '%s'." %
124
                               str(creat_err))
125

    
126
      try:
127
        pf = open(lockfile, 'r')
128
      except IOError, open_err:
129
        errcount += 1
130
        if errcount >= 5:
131
          raise errors.LockError("Lock file exists but cannot be opened."
132
                                 " Error: '%s'." % str(open_err))
133
        time.sleep(1)
134
        continue
135

    
136
      try:
137
        pid = int(pf.read())
138
      except ValueError:
139
        raise errors.LockError("Invalid pid string in %s" %
140
                               (lockfile,))
141

    
142
      if not IsProcessAlive(pid):
143
        raise errors.LockError("Stale lockfile %s for pid %d?" %
144
                               (lockfile, pid))
145

    
146
      if max_retries and max_retries <= retries:
147
        raise errors.LockError("Can't acquire lock during the specified"
148
                               " time, aborting.")
149
      if retries == 5 and (debug or sys.stdin.isatty()):
150
        logger.ToStderr("Waiting for '%s' lock from pid %d..." % (name, pid))
151

    
152
      time.sleep(1)
153
      retries += 1
154
      continue
155

    
156
  os.write(fd, '%d\n' % (os.getpid(),))
157
  os.close(fd)
158

    
159
  _locksheld.append(name)
160

    
161

    
162
def Unlock(name):
163
  """Unlock a given subsystem.
164

165
  """
166
  lockfile = _GetLockFile(name)
167

    
168
  try:
169
    fd = os.open(lockfile, os.O_RDONLY)
170
  except OSError:
171
    raise errors.LockError('Lock "%s" not held.' % (name,))
172

    
173
  f = os.fdopen(fd, 'r')
174
  pid_str = f.read()
175

    
176
  try:
177
    pid = int(pid_str)
178
  except ValueError:
179
    raise errors.LockError('Unable to determine PID of locking process.')
180

    
181
  if pid != os.getpid():
182
    raise errors.LockError('Lock not held by me (%d != %d)' %
183
                           (os.getpid(), pid,))
184

    
185
  os.unlink(lockfile)
186
  _locksheld.remove(name)
187

    
188

    
189
def LockCleanup():
190
  """Remove all locks.
191

192
  """
193
  for lock in _locksheld:
194
    Unlock(lock)
195

    
196

    
197
def RunCmd(cmd):
198
  """Execute a (shell) command.
199

200
  The command should not read from its standard input, as it will be
201
  closed.
202

203
  Args:
204
    cmd: command to run. (str)
205

206
  Returns: `RunResult` instance
207

208
  """
209
  if isinstance(cmd, list):
210
    cmd = [str(val) for val in cmd]
211
    strcmd = " ".join(cmd)
212
    shell = False
213
  else:
214
    strcmd = cmd
215
    shell = True
216
  env = os.environ.copy()
217
  env["LC_ALL"] = "C"
218
  child = subprocess.Popen(cmd, shell=shell,
219
                           stderr=subprocess.PIPE,
220
                           stdout=subprocess.PIPE,
221
                           stdin=subprocess.PIPE,
222
                           close_fds=True, env=env)
223

    
224
  child.stdin.close()
225
  out = child.stdout.read()
226
  err = child.stderr.read()
227

    
228
  status = child.wait()
229
  if status >= 0:
230
    exitcode = status
231
    signal = None
232
  else:
233
    exitcode = None
234
    signal = -status
235

    
236
  return RunResult(exitcode, signal, out, err, strcmd)
237

    
238

    
239
def RunCmdUnlocked(cmd):
240
  """Execute a shell command without the 'cmd' lock.
241

242
  This variant of `RunCmd()` drops the 'cmd' lock before running the
243
  command and re-aquires it afterwards, thus it can be used to call
244
  other ganeti commands.
245

246
  The argument and return values are the same as for the `RunCmd()`
247
  function.
248

249
  Args:
250
    cmd - command to run. (str)
251

252
  Returns:
253
    `RunResult`
254

255
  """
256
  Unlock('cmd')
257
  ret = RunCmd(cmd)
258
  Lock('cmd')
259

    
260
  return ret
261

    
262

    
263
def RemoveFile(filename):
264
  """Remove a file ignoring some errors.
265

266
  Remove a file, ignoring non-existing ones or directories. Other
267
  errors are passed.
268

269
  """
270
  try:
271
    os.unlink(filename)
272
  except OSError, err:
273
    if err.errno not in (errno.ENOENT, errno.EISDIR):
274
      raise
275

    
276

    
277
def _FingerprintFile(filename):
278
  """Compute the fingerprint of a file.
279

280
  If the file does not exist, a None will be returned
281
  instead.
282

283
  Args:
284
    filename - Filename (str)
285

286
  """
287
  if not (os.path.exists(filename) and os.path.isfile(filename)):
288
    return None
289

    
290
  f = open(filename)
291

    
292
  fp = sha.sha()
293
  while True:
294
    data = f.read(4096)
295
    if not data:
296
      break
297

    
298
    fp.update(data)
299

    
300
  return fp.hexdigest()
301

    
302

    
303
def FingerprintFiles(files):
304
  """Compute fingerprints for a list of files.
305

306
  Args:
307
    files - array of filenames.  ( [str, ...] )
308

309
  Return value:
310
    dictionary of filename: fingerprint for the files that exist
311

312
  """
313
  ret = {}
314

    
315
  for filename in files:
316
    cksum = _FingerprintFile(filename)
317
    if cksum:
318
      ret[filename] = cksum
319

    
320
  return ret
321

    
322

    
323
def CheckDict(target, template, logname=None):
324
  """Ensure a dictionary has a required set of keys.
325

326
  For the given dictionaries `target` and `template`, ensure target
327
  has all the keys from template. Missing keys are added with values
328
  from template.
329

330
  Args:
331
    target   - the dictionary to check
332
    template - template dictionary
333
    logname  - a caller-chosen string to identify the debug log
334
               entry; if None, no logging will be done
335

336
  Returns value:
337
    None
338

339
  """
340
  missing = []
341
  for k in template:
342
    if k not in target:
343
      missing.append(k)
344
      target[k] = template[k]
345

    
346
  if missing and logname:
347
    logger.Debug('%s missing keys %s' %
348
                 (logname, ', '.join(missing)))
349

    
350

    
351
def IsProcessAlive(pid):
352
  """Check if a given pid exists on the system.
353

354
  Returns: true or false, depending on if the pid exists or not
355

356
  Remarks: zombie processes treated as not alive
357

358
  """
359
  try:
360
    f = open("/proc/%d/status" % pid)
361
  except IOError, err:
362
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
363
      return False
364

    
365
  alive = True
366
  try:
367
    data = f.readlines()
368
    if len(data) > 1:
369
      state = data[1].split()
370
      if len(state) > 1 and state[1] == "Z":
371
        alive = False
372
  finally:
373
    f.close()
374

    
375
  return alive
376

    
377

    
378
def MatchNameComponent(key, name_list):
379
  """Try to match a name against a list.
380

381
  This function will try to match a name like test1 against a list
382
  like ['test1.example.com', 'test2.example.com', ...]. Against this
383
  list, 'test1' as well as 'test1.example' will match, but not
384
  'test1.ex'. A multiple match will be considered as no match at all
385
  (e.g. 'test1' against ['test1.example.com', 'test1.example.org']).
386

387
  Args:
388
    key: the name to be searched
389
    name_list: the list of strings against which to search the key
390

391
  Returns:
392
    None if there is no match *or* if there are multiple matches
393
    otherwise the element from the list which matches
394

395
  """
396
  mo = re.compile("^%s(\..*)?$" % re.escape(key))
397
  names_filtered = [name for name in name_list if mo.match(name) is not None]
398
  if len(names_filtered) != 1:
399
    return None
400
  return names_filtered[0]
401

    
402

    
403
class HostInfo:
404
  """Class implementing resolver and hostname functionality
405

406
  """
407
  def __init__(self, name=None):
408
    """Initialize the host name object.
409

410
    If the name argument is not passed, it will use this system's
411
    name.
412

413
    """
414
    if name is None:
415
      name = self.SysName()
416

    
417
    self.query = name
418
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
419
    self.ip = self.ipaddrs[0]
420

    
421
  def ShortName(self):
422
    """Returns the hostname without domain.
423

424
    """
425
    return self.name.split('.')[0]
426

    
427
  @staticmethod
428
  def SysName():
429
    """Return the current system's name.
430

431
    This is simply a wrapper over socket.gethostname()
432

433
    """
434
    return socket.gethostname()
435

    
436
  @staticmethod
437
  def LookupHostname(hostname):
438
    """Look up hostname
439

440
    Args:
441
      hostname: hostname to look up
442

443
    Returns:
444
      a tuple (name, aliases, ipaddrs) as returned by socket.gethostbyname_ex
445
      in case of errors in resolving, we raise a ResolverError
446

447
    """
448
    try:
449
      result = socket.gethostbyname_ex(hostname)
450
    except socket.gaierror, err:
451
      # hostname not found in DNS
452
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
453

    
454
    return result
455

    
456

    
457
def ListVolumeGroups():
458
  """List volume groups and their size
459

460
  Returns:
461
     Dictionary with keys volume name and values the size of the volume
462

463
  """
464
  command = "vgs --noheadings --units m --nosuffix -o name,size"
465
  result = RunCmd(command)
466
  retval = {}
467
  if result.failed:
468
    return retval
469

    
470
  for line in result.stdout.splitlines():
471
    try:
472
      name, size = line.split()
473
      size = int(float(size))
474
    except (IndexError, ValueError), err:
475
      logger.Error("Invalid output from vgs (%s): %s" % (err, line))
476
      continue
477

    
478
    retval[name] = size
479

    
480
  return retval
481

    
482

    
483
def BridgeExists(bridge):
484
  """Check whether the given bridge exists in the system
485

486
  Returns:
487
     True if it does, false otherwise.
488

489
  """
490
  return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
491

    
492

    
493
def NiceSort(name_list):
494
  """Sort a list of strings based on digit and non-digit groupings.
495

496
  Given a list of names ['a1', 'a10', 'a11', 'a2'] this function will
497
  sort the list in the logical order ['a1', 'a2', 'a10', 'a11'].
498

499
  The sort algorithm breaks each name in groups of either only-digits
500
  or no-digits. Only the first eight such groups are considered, and
501
  after that we just use what's left of the string.
502

503
  Return value
504
    - a copy of the list sorted according to our algorithm
505

506
  """
507
  _SORTER_BASE = "(\D+|\d+)"
508
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
509
                                                  _SORTER_BASE, _SORTER_BASE,
510
                                                  _SORTER_BASE, _SORTER_BASE,
511
                                                  _SORTER_BASE, _SORTER_BASE)
512
  _SORTER_RE = re.compile(_SORTER_FULL)
513
  _SORTER_NODIGIT = re.compile("^\D*$")
514
  def _TryInt(val):
515
    """Attempts to convert a variable to integer."""
516
    if val is None or _SORTER_NODIGIT.match(val):
517
      return val
518
    rval = int(val)
519
    return rval
520

    
521
  to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
522
             for name in name_list]
523
  to_sort.sort()
524
  return [tup[1] for tup in to_sort]
525

    
526

    
527
def CheckDaemonAlive(pid_file, process_string):
528
  """Check wether the specified daemon is alive.
529

530
  Args:
531
   - pid_file: file to read the daemon pid from, the file is
532
               expected to contain only a single line containing
533
               only the PID
534
   - process_string: a substring that we expect to find in
535
                     the command line of the daemon process
536

537
  Returns:
538
   - True if the daemon is judged to be alive (that is:
539
      - the PID file exists, is readable and contains a number
540
      - a process of the specified PID is running
541
      - that process contains the specified string in its
542
        command line
543
      - the process is not in state Z (zombie))
544
   - False otherwise
545

546
  """
547
  try:
548
    pid_file = file(pid_file, 'r')
549
    try:
550
      pid = int(pid_file.readline())
551
    finally:
552
      pid_file.close()
553

    
554
    cmdline_file_path = "/proc/%s/cmdline" % (pid)
555
    cmdline_file = open(cmdline_file_path, 'r')
556
    try:
557
      cmdline = cmdline_file.readline()
558
    finally:
559
      cmdline_file.close()
560

    
561
    if not process_string in cmdline:
562
      return False
563

    
564
    stat_file_path =  "/proc/%s/stat" % (pid)
565
    stat_file = open(stat_file_path, 'r')
566
    try:
567
      process_state = stat_file.readline().split()[2]
568
    finally:
569
      stat_file.close()
570

    
571
    if process_state == 'Z':
572
      return False
573

    
574
  except (IndexError, IOError, ValueError):
575
    return False
576

    
577
  return True
578

    
579

    
580
def TryConvert(fn, val):
581
  """Try to convert a value ignoring errors.
582

583
  This function tries to apply function `fn` to `val`. If no
584
  ValueError or TypeError exceptions are raised, it will return the
585
  result, else it will return the original value. Any other exceptions
586
  are propagated to the caller.
587

588
  """
589
  try:
590
    nv = fn(val)
591
  except (ValueError, TypeError), err:
592
    nv = val
593
  return nv
594

    
595

    
596
def IsValidIP(ip):
597
  """Verifies the syntax of an IP address.
598

599
  This function checks if the ip address passes is valid or not based
600
  on syntax (not ip range, class calculations or anything).
601

602
  """
603
  unit = "(0|[1-9]\d{0,2})"
604
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
605

    
606

    
607
def IsValidShellParam(word):
608
  """Verifies is the given word is safe from the shell's p.o.v.
609

610
  This means that we can pass this to a command via the shell and be
611
  sure that it doesn't alter the command line and is passed as such to
612
  the actual command.
613

614
  Note that we are overly restrictive here, in order to be on the safe
615
  side.
616

617
  """
618
  return bool(re.match("^[-a-zA-Z0-9._+/:%@]+$", word))
619

    
620

    
621
def BuildShellCmd(template, *args):
622
  """Build a safe shell command line from the given arguments.
623

624
  This function will check all arguments in the args list so that they
625
  are valid shell parameters (i.e. they don't contain shell
626
  metacharaters). If everything is ok, it will return the result of
627
  template % args.
628

629
  """
630
  for word in args:
631
    if not IsValidShellParam(word):
632
      raise errors.ProgrammerError("Shell argument '%s' contains"
633
                                   " invalid characters" % word)
634
  return template % args
635

    
636

    
637
def FormatUnit(value):
638
  """Formats an incoming number of MiB with the appropriate unit.
639

640
  Value needs to be passed as a numeric type. Return value is always a string.
641

642
  """
643
  if value < 1024:
644
    return "%dM" % round(value, 0)
645

    
646
  elif value < (1024 * 1024):
647
    return "%0.1fG" % round(float(value) / 1024, 1)
648

    
649
  else:
650
    return "%0.1fT" % round(float(value) / 1024 / 1024, 1)
651

    
652

    
653
def ParseUnit(input_string):
654
  """Tries to extract number and scale from the given string.
655

656
  Input must be in the format NUMBER+ [DOT NUMBER+] SPACE* [UNIT]. If no unit
657
  is specified, it defaults to MiB. Return value is always an int in MiB.
658

659
  """
660
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', input_string)
661
  if not m:
662
    raise errors.UnitParseError("Invalid format")
663

    
664
  value = float(m.groups()[0])
665

    
666
  unit = m.groups()[1]
667
  if unit:
668
    lcunit = unit.lower()
669
  else:
670
    lcunit = 'm'
671

    
672
  if lcunit in ('m', 'mb', 'mib'):
673
    # Value already in MiB
674
    pass
675

    
676
  elif lcunit in ('g', 'gb', 'gib'):
677
    value *= 1024
678

    
679
  elif lcunit in ('t', 'tb', 'tib'):
680
    value *= 1024 * 1024
681

    
682
  else:
683
    raise errors.UnitParseError("Unknown unit: %s" % unit)
684

    
685
  # Make sure we round up
686
  if int(value) < value:
687
    value += 1
688

    
689
  # Round up to the next multiple of 4
690
  value = int(value)
691
  if value % 4:
692
    value += 4 - value % 4
693

    
694
  return value
695

    
696

    
697
def AddAuthorizedKey(file_name, key):
698
  """Adds an SSH public key to an authorized_keys file.
699

700
  Args:
701
    file_name: Path to authorized_keys file
702
    key: String containing key
703
  """
704
  key_fields = key.split()
705

    
706
  f = open(file_name, 'a+')
707
  try:
708
    nl = True
709
    for line in f:
710
      # Ignore whitespace changes
711
      if line.split() == key_fields:
712
        break
713
      nl = line.endswith('\n')
714
    else:
715
      if not nl:
716
        f.write("\n")
717
      f.write(key.rstrip('\r\n'))
718
      f.write("\n")
719
      f.flush()
720
  finally:
721
    f.close()
722

    
723

    
724
def RemoveAuthorizedKey(file_name, key):
725
  """Removes an SSH public key from an authorized_keys file.
726

727
  Args:
728
    file_name: Path to authorized_keys file
729
    key: String containing key
730
  """
731
  key_fields = key.split()
732

    
733
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
734
  try:
735
    out = os.fdopen(fd, 'w')
736
    try:
737
      f = open(file_name, 'r')
738
      try:
739
        for line in f:
740
          # Ignore whitespace changes while comparing lines
741
          if line.split() != key_fields:
742
            out.write(line)
743

    
744
        out.flush()
745
        os.rename(tmpname, file_name)
746
      finally:
747
        f.close()
748
    finally:
749
      out.close()
750
  except:
751
    RemoveFile(tmpname)
752
    raise
753

    
754

    
755
def SetEtcHostsEntry(file_name, ip, hostname, aliases):
756
  """Sets the name of an IP address and hostname in /etc/hosts.
757

758
  """
759
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
760
  try:
761
    out = os.fdopen(fd, 'w')
762
    try:
763
      f = open(file_name, 'r')
764
      try:
765
        written = False
766
        for line in f:
767
          fields = line.split()
768
          if fields and not fields[0].startswith('#') and ip == fields[0]:
769
            continue
770
          out.write(line)
771

    
772
        out.write("%s\t%s" % (ip, hostname))
773
        if aliases:
774
          out.write(" %s" % ' '.join(aliases))
775
        out.write('\n')
776

    
777
        out.flush()
778
        os.fsync(out)
779
        os.rename(tmpname, file_name)
780
      finally:
781
        f.close()
782
    finally:
783
      out.close()
784
  except:
785
    RemoveFile(tmpname)
786
    raise
787

    
788

    
789
def RemoveEtcHostsEntry(file_name, hostname):
790
  """Removes a hostname from /etc/hosts.
791

792
  IP addresses without names are removed from the file.
793
  """
794
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
795
  try:
796
    out = os.fdopen(fd, 'w')
797
    try:
798
      f = open(file_name, 'r')
799
      try:
800
        for line in f:
801
          fields = line.split()
802
          if len(fields) > 1 and not fields[0].startswith('#'):
803
            names = fields[1:]
804
            if hostname in names:
805
              while hostname in names:
806
                names.remove(hostname)
807
              if names:
808
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
809
              continue
810

    
811
          out.write(line)
812

    
813
        out.flush()
814
        os.fsync(out)
815
        os.rename(tmpname, file_name)
816
      finally:
817
        f.close()
818
    finally:
819
      out.close()
820
  except:
821
    RemoveFile(tmpname)
822
    raise
823

    
824

    
825
def CreateBackup(file_name):
826
  """Creates a backup of a file.
827

828
  Returns: the path to the newly created backup file.
829

830
  """
831
  if not os.path.isfile(file_name):
832
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
833
                                file_name)
834

    
835
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
836
  dir_name = os.path.dirname(file_name)
837

    
838
  fsrc = open(file_name, 'rb')
839
  try:
840
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
841
    fdst = os.fdopen(fd, 'wb')
842
    try:
843
      shutil.copyfileobj(fsrc, fdst)
844
    finally:
845
      fdst.close()
846
  finally:
847
    fsrc.close()
848

    
849
  return backup_name
850

    
851

    
852
def ShellQuote(value):
853
  """Quotes shell argument according to POSIX.
854

855
  """
856
  if _re_shell_unquoted.match(value):
857
    return value
858
  else:
859
    return "'%s'" % value.replace("'", "'\\''")
860

    
861

    
862
def ShellQuoteArgs(args):
863
  """Quotes all given shell arguments and concatenates using spaces.
864

865
  """
866
  return ' '.join([ShellQuote(i) for i in args])
867

    
868

    
869

    
870
def TcpPing(source, target, port, timeout=10, live_port_needed=False):
871
  """Simple ping implementation using TCP connect(2).
872

873
  Try to do a TCP connect(2) from the specified source IP to the specified
874
  target IP and the specified target port. If live_port_needed is set to true,
875
  requires the remote end to accept the connection. The timeout is specified
876
  in seconds and defaults to 10 seconds
877

878
  """
879
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
880

    
881
  sucess = False
882

    
883
  try:
884
    sock.bind((source, 0))
885
  except socket.error, (errcode, errstring):
886
    if errcode == errno.EADDRNOTAVAIL:
887
      success = False
888

    
889
  sock.settimeout(timeout)
890

    
891
  try:
892
    sock.connect((target, port))
893
    sock.close()
894
    success = True
895
  except socket.timeout:
896
    success = False
897
  except socket.error, (errcode, errstring):
898
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
899

    
900
  return success
901

    
902

    
903
def ListVisibleFiles(path):
904
  """Returns a list of all visible files in a directory.
905

906
  """
907
  return [i for i in os.listdir(path) if not i.startswith(".")]
908

    
909

    
910
def GetHomeDir(user, default=None):
911
  """Try to get the homedir of the given user.
912

913
  The user can be passed either as a string (denoting the name) or as
914
  an integer (denoting the user id). If the user is not found, the
915
  'default' argument is returned, which defaults to None.
916

917
  """
918
  try:
919
    if isinstance(user, basestring):
920
      result = pwd.getpwnam(user)
921
    elif isinstance(user, (int, long)):
922
      result = pwd.getpwuid(user)
923
    else:
924
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
925
                                   type(user))
926
  except KeyError:
927
    return default
928
  return result.pw_dir
929

    
930

    
931
def NewUUID():
932
  """Returns a random UUID.
933

934
  """
935
  f = open("/proc/sys/kernel/random/uuid", "r")
936
  try:
937
    return f.read(128).rstrip("\n")
938
  finally:
939
    f.close()
940

    
941

    
942
def WriteFile(file_name, fn=None, data=None,
943
              mode=None, uid=-1, gid=-1,
944
              atime=None, mtime=None):
945
  """(Over)write a file atomically.
946

947
  The file_name and either fn (a function taking one argument, the
948
  file descriptor, and which should write the data to it) or data (the
949
  contents of the file) must be passed. The other arguments are
950
  optional and allow setting the file mode, owner and group, and the
951
  mtime/atime of the file.
952

953
  If the function doesn't raise an exception, it has succeeded and the
954
  target file has the new contents. If the file has raised an
955
  exception, an existing target file should be unmodified and the
956
  temporary file should be removed.
957

958
  """
959
  if not os.path.isabs(file_name):
960
    raise errors.ProgrammerError("Path passed to WriteFile is not"
961
                                 " absolute: '%s'" % file_name)
962

    
963
  if [fn, data].count(None) != 1:
964
    raise errors.ProgrammerError("fn or data required")
965

    
966
  if [atime, mtime].count(None) == 1:
967
    raise errors.ProgrammerError("Both atime and mtime must be either"
968
                                 " set or None")
969

    
970

    
971
  dir_name, base_name = os.path.split(file_name)
972
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
973
  # here we need to make sure we remove the temp file, if any error
974
  # leaves it in place
975
  try:
976
    if uid != -1 or gid != -1:
977
      os.chown(new_name, uid, gid)
978
    if mode:
979
      os.chmod(new_name, mode)
980
    if data is not None:
981
      os.write(fd, data)
982
    else:
983
      fn(fd)
984
    os.fsync(fd)
985
    if atime is not None and mtime is not None:
986
      os.utime(new_name, (atime, mtime))
987
    os.rename(new_name, file_name)
988
  finally:
989
    os.close(fd)
990
    RemoveFile(new_name)
991

    
992

    
993
def all(seq, pred=bool):
994
  "Returns True if pred(x) is True for every element in the iterable"
995
  for elem in itertools.ifilterfalse(pred, seq):
996
    return False
997
  return True
998

    
999

    
1000
def any(seq, pred=bool):
1001
  "Returns True if pred(x) is True for at least one element in the iterable"
1002
  for elem in itertools.ifilter(pred, seq):
1003
    return True
1004
  return False
1005

    
1006

    
1007
def UniqueSequence(seq):
1008
  """Returns a list with unique elements.
1009

1010
  Element order is preserved.
1011
  """
1012
  seen = set()
1013
  return [i for i in seq if i not in seen and not seen.add(i)]