Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ c8a0948f

History | View | Annotate | Download (24.6 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 AddEtcHostsEntry(file_name, hostname, ip):
756
  """
757

758
  """
759
  f = open(file_name, 'a+')
760
  try:
761
    nl = True
762
    for line in f:
763
      fields = line.split()
764
      if len(fields) < 2 or fields[0].startswith('#'):
765
        continue
766
      if fields[0] == ip and hostname in fields[1:]:
767
        break
768
      nl = line.endswith('\n')
769
    else:
770
      if not nl:
771
        f.write("\n")
772
      f.write(ip)
773
      f.write(' ')
774
      f.write(hostname)
775
      f.write("\n")
776
      f.flush()
777
  finally:
778
    f.close()
779

    
780

    
781
def RemoveEtcHostsEntry(file_name, hostname):
782
  """
783

784
  """
785
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
786
  try:
787
    out = os.fdopen(fd, 'w')
788
    try:
789
      f = open(file_name, 'r')
790
      try:
791
        for line in f:
792
          fields = line.split()
793
          if len(fields) > 1 and not fields[0].startswith('#'):
794
            names = fields[1:]
795
            if hostname in names:
796
              while hostname in names:
797
                names.remove(hostname)
798
              if names:
799
                out.write(fields[0])
800
                out.write(' ')
801
                out.write(' '.join(names))
802
              continue
803

    
804
          out.write(line)
805

    
806
        out.flush()
807
        os.rename(tmpname, file_name)
808
      finally:
809
        f.close()
810
    finally:
811
      out.close()
812
  except:
813
    RemoveFile(tmpname)
814
    raise
815

    
816

    
817
def CreateBackup(file_name):
818
  """Creates a backup of a file.
819

820
  Returns: the path to the newly created backup file.
821

822
  """
823
  if not os.path.isfile(file_name):
824
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
825
                                file_name)
826

    
827
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
828
  dir_name = os.path.dirname(file_name)
829

    
830
  fsrc = open(file_name, 'rb')
831
  try:
832
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
833
    fdst = os.fdopen(fd, 'wb')
834
    try:
835
      shutil.copyfileobj(fsrc, fdst)
836
    finally:
837
      fdst.close()
838
  finally:
839
    fsrc.close()
840

    
841
  return backup_name
842

    
843

    
844
def ShellQuote(value):
845
  """Quotes shell argument according to POSIX.
846

847
  """
848
  if _re_shell_unquoted.match(value):
849
    return value
850
  else:
851
    return "'%s'" % value.replace("'", "'\\''")
852

    
853

    
854
def ShellQuoteArgs(args):
855
  """Quotes all given shell arguments and concatenates using spaces.
856

857
  """
858
  return ' '.join([ShellQuote(i) for i in args])
859

    
860

    
861

    
862
def TcpPing(source, target, port, timeout=10, live_port_needed=False):
863
  """Simple ping implementation using TCP connect(2).
864

865
  Try to do a TCP connect(2) from the specified source IP to the specified
866
  target IP and the specified target port. If live_port_needed is set to true,
867
  requires the remote end to accept the connection. The timeout is specified
868
  in seconds and defaults to 10 seconds
869

870
  """
871
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
872

    
873
  sucess = False
874

    
875
  try:
876
    sock.bind((source, 0))
877
  except socket.error, (errcode, errstring):
878
    if errcode == errno.EADDRNOTAVAIL:
879
      success = False
880

    
881
  sock.settimeout(timeout)
882

    
883
  try:
884
    sock.connect((target, port))
885
    sock.close()
886
    success = True
887
  except socket.timeout:
888
    success = False
889
  except socket.error, (errcode, errstring):
890
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
891

    
892
  return success
893

    
894

    
895
def ListVisibleFiles(path):
896
  """Returns a list of all visible files in a directory.
897

898
  """
899
  return [i for i in os.listdir(path) if not i.startswith(".")]
900

    
901

    
902
def GetHomeDir(user, default=None):
903
  """Try to get the homedir of the given user.
904

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

909
  """
910
  try:
911
    if isinstance(user, basestring):
912
      result = pwd.getpwnam(user)
913
    elif isinstance(user, (int, long)):
914
      result = pwd.getpwuid(user)
915
    else:
916
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
917
                                   type(user))
918
  except KeyError:
919
    return default
920
  return result.pw_dir
921

    
922

    
923
def NewUUID():
924
  """Returns a random UUID.
925

926
  """
927
  f = open("/proc/sys/kernel/random/uuid", "r")
928
  try:
929
    return f.read(128).rstrip("\n")
930
  finally:
931
    f.close()
932

    
933

    
934
def WriteFile(file_name, fn=None, data=None,
935
              mode=None, uid=-1, gid=-1,
936
              atime=None, mtime=None):
937
  """(Over)write a file atomically.
938

939
  The file_name and either fn (a function taking one argument, the
940
  file descriptor, and which should write the data to it) or data (the
941
  contents of the file) must be passed. The other arguments are
942
  optional and allow setting the file mode, owner and group, and the
943
  mtime/atime of the file.
944

945
  If the function doesn't raise an exception, it has succeeded and the
946
  target file has the new contents. If the file has raised an
947
  exception, an existing target file should be unmodified and the
948
  temporary file should be removed.
949

950
  """
951
  if not os.path.isabs(file_name):
952
    raise errors.ProgrammerError("Path passed to WriteFile is not"
953
                                 " absolute: '%s'" % file_name)
954

    
955
  if [fn, data].count(None) != 1:
956
    raise errors.ProgrammerError("fn or data required")
957

    
958
  if [atime, mtime].count(None) == 1:
959
    raise errors.ProgrammerError("Both atime and mtime must be either"
960
                                 " set or None")
961

    
962

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

    
984

    
985
def all(seq, pred=bool):
986
  "Returns True if pred(x) is True for every element in the iterable"
987
  for elem in itertools.ifilterfalse(pred, seq):
988
    return False
989
  return True
990

    
991

    
992
def any(seq, pred=bool):
993
  "Returns True if pred(x) is True for at least one element in the iterable"
994
  for elem in itertools.ifilter(pred, seq):
995
    return True
996
  return False