#!/usr/bin/python ############################################################ # # Common utilities for the ONL python tools. # ############################################################ import logging import subprocess from collections import Iterable import sys import os import fcntl import glob from string import Template logger = None # Cheap colored terminal logging. Fixme. def color_logging(): if sys.stderr.isatty(): logging.addLevelName( logging.WARNING, "\033[1;31m%s\033[1;0m" % logging.getLevelName(logging.WARNING)) logging.addLevelName( logging.ERROR, "\033[1;41m%s\033[1;0m" % logging.getLevelName(logging.ERROR)) def init_logging(name, lvl=logging.DEBUG): global logger logging.basicConfig() logger = logging.getLogger(name) logger.setLevel(lvl) color_logging() return logger ############################################################ # # Log and execute system commands # def execute(args, sudo=False, chroot=None, ex=None): if type(args) is str: # Must be executed through the shell shell=True else: shell = False if chroot and os.geteuid() != 0: # Must be executed under sudo sudo = True if chroot: if type(args) is str: args = "chroot %s %s" % (chroot, args) elif type(args) in (list,tuple): args = ['chroot', chroot] + list(args) if sudo: if type(args) is str: args = "sudo %s" % (args) elif type(args) in (list, tuple): args = [ 'sudo' ] + list(args) logger.debug("Executing:%s", args) try: subprocess.check_call(args, shell=shell) return 0 except subprocess.CalledProcessError, e: if ex: raise ex return e.returncode # Flatten lists if string lists def sflatten(coll): for i in coll: if isinstance(i, Iterable) and not isinstance(i, basestring): for subc in sflatten(i): if subc: yield subc else: yield i def gen_salt(): # return an eight character salt value salt_map = './' + string.digits + string.ascii_uppercase + \ string.ascii_lowercase rand = random.SystemRandom() salt = '' for i in range(0, 8): salt += salt_map[rand.randint(0, 63)] return salt ############################################################ # # Delete a user # def userdel(username): # Can't use the userdel command because of potential uid 0 in-user problems while running ourselves for line in fileinput.input('/etc/passwd', inplace=True): if not line.startswith('%s:' % username): print line, for line in fileinput.input('/etc/shadow', inplace=True): if not line.startswith('%s:' % username): print line, ############################################################ # # Add a user # def useradd(username, uid, password, shell, deleteFirst=True): args = [ 'useradd', '--non-unique', '--shell', shell, '--home-dir', '/root', '--uid', '0', '--gid', '0', '--group', 'root' ] if deleteFirst: userdel(username) if password: epassword=crypt.crypt(password, '$1$%s$' % gen_salt()); args = args + ['-p', epassword] args.append(username) cc(args); if password is None: cc(('passwd', '-d', username)) logger.info("user %s password %s", username, password) class Lock(object): """File Locking class.""" def __init__(self, filename): self.filename = filename self.handle = open(filename, 'w') def take(self): # logger.debug("taking lock %s" % self.filename) fcntl.flock(self.handle, fcntl.LOCK_EX) # logger.debug("took lock %s" % self.filename) def give(self): fcntl.flock(self.handle, fcntl.LOCK_UN) # logger.debug("released lock %s" % self.filename) def __enter__(self): self.take() def __exit__(self ,type, value, traceback): self.give() def __del__(self): self.handle.close() def filepath(absdir, relpath, eklass): """Return the absolute path for the given absdir/repath file.""" p = None if os.path.isabs(relpath): p = relpath else: p = os.path.join(absdir, relpath) # Globs that result in a single file are allowed: g = glob.glob(p) if len(g) is 0: raise eklass("'%s' did not match any files." % p) elif len(g) > 1: raise eklass("'%s' matched too many files %s" % (p, g)) else: p = g[0] return p def validate_src_dst_file_tuples(absdir, data, dstsubs, eklass): files = [] if type(data) is dict: for (s,d) in data.iteritems(): files.append((s,d)) elif type(data) is list: for e in data: if type(e) is dict: for (s,d) in e.iteritems(): files.append((s,d)) elif type(e) in [ list, tuple ]: if len(e) != 2: raise eklass("Too many filenames: '%s'" % (e)) else: files.append((e[0], e[1])) else: raise eklass("File entry '%s' is invalid." % (e)) # # Validate/process source files. # Process dst paths. # flist = [] for f in files: src = filepath(absdir, f[0], eklass) if not os.path.exists(src): raise eklass("Source file or directory '%s' does not exist." % src) t = Template(f[1]) dst = t.substitute(dstsubs) flist.append((src, dst)) return flist