path: root/rainbow/util.py

from __future__ import with_statement

import os

from stat import ST_MODE, ST_UID, ST_GID, S_ISREG

from os import strerror, stat, mkdir, chown, walk, listdir, fstat
from os.path import realpath, split, join, exists, isdir

from errno import errorcode

from ctypes import CDLL, POINTER, Structure, c_int, c_char_p, c_uint32, c_long
from ctypes.util import find_library

from cStringIO import StringIO as stringio
from subprocess import Popen, PIPE

### errno manipulationA

libc = CDLL(find_library('c'))
libc.__errno_location.restype = POINTER(c_int)

def get_errno():
    return libc.__errno_location().contents.value

class CError(Exception):
    """Base class for VServer syscall errors."""

    def __init__(self, errno=None):
        Exception.__init__(self)
        if errno is None:
            errno = get_errno()
        self.errno = errno
        self.errsym = errorcode[errno]
        self.errstring = strerror(errno)

    def __repr__(self):
        return '<%s errno=%s>'%(self.__class__.__name__, self.errno)

    def __str__(self):
        return '%s: %s'%(self.errsym, self.errstring)

def errno_errcheck(errval, exception):
    """Assign to the 'errcheck' attr of a ctypes._CFuncPtr.
       :errval: a return code indicating error, and
       :exception: an exception type like CError.
       >>> func.errcheck = errno_errcheck(errval, excpt)"""
    def _inner(result, func, args):
        if result == errval:
            errno = get_errno()
            raise exception(errno)
        return result
    return _inner

libc_errcheck = errno_errcheck(-1, CError)

### Logging, profiling, and error reporting

def make_reporter(verbosity, quiet, filelike):
    if not quiet:
        def report(level, msg, *args):
            if level <= verbosity:
                if len(args):
                    filelike.write(msg % args + '\n')
                else:
                    filelike.write(msg + '\n')
    else:
        def report(level, msg, *args): pass
    return report

def profile(profiler):
    def wrapper(fn):
        def inner(*args, **kwargs):
            profiler.runcall(fn, *args, **kwargs)
        return inner
    return wrapper

def trace(etype=None, value=None, tb=None):
    try:
        from IPython.ultraTB import AutoFormattedTB
        trace_any = AutoFormattedTB(mode='Verbose', color_scheme='NoColor', call_pdb=0)
        trace_exact = trace_any
    except:
        from traceback import print_exc, print_exception
        trace_any = print_exc
        trace_exact = print_exception
    if etype or value or tb:
        trace_exact(etype, value, tb)
    else:
        trace_any()

### File mode checking.

class Checker(object):
    """Sanity-check file-system permissions.
       Note: This class should NOT to be used for security-critical tests.
       In particular, this class does not know the special rules that adhere to
       uid 0.
    """
    def __init__(self, path, uid, gid, groups=None):
        self.path = path
        self.uid = uid
        self.gid = gid
        self.groups = set(groups or []).union(set([gid]))
        self.observation = stat(self.path)

    def __repr__(self):
        return "Checker(%r, %r, %r, groups=%r)" % (self.path, self.uid, self.gid, self.groups)

    def __str__(self):
        return ("Checker(%s, %s, %s, %s) -> mode: %o uid: %s, gid: %s" %
            (self.path, self.uid, self.gid, self.groups, self.observation[ST_MODE],
            self.observation[ST_UID], self.observation[ST_GID]))

    def positive(self, needed, mask):
        # I'm going to try to turn off bits in mode as I verify that they can
        # be satisfied.
        o = self.observation
        real_mode = o[ST_MODE]
        if self.uid == o[ST_UID]:
            needed &= ~((real_mode & 0700) >> 6)
        if o[ST_GID] in self.groups:
            needed &= ~((real_mode & 0070) >> 3)
        needed &= ~(real_mode & 0007)
        # Make sure needed is empty and real_mode satisfies stat.
        return (needed == 0) and ((mask & real_mode) == mask)

    def negative(self, avoid, mask):
        # I'm going to try to turn on bits for things we're permitted to do,
        # then check if any of them are disallowed.
        o = self.observation
        real_mode = o[ST_MODE]
        allowed = 0
        if self.uid == o[ST_UID]:
            allowed |= (real_mode & 0700) >> 6
        if o[ST_GID] in self.groups:
            allowed |= (real_mode & 0070) >> 3
        allowed |= (real_mode & 0007)
        # Did we avoid everything?
        return ((avoid & allowed) == 0) and (mask & real_mode == 0)

### Paths

def components(p):
    """Take a path and return a list of its components:
    e.g. /abc/de or /abc/de/ ---> ['abc', 'de']"""
    path = realpath(p)
    (head, tail) = split(path)
    if tail == '':
        (head, tail) = split(head)
    cpts = []
    while tail != '':
        cpts = [tail] + cpts
        (head, tail) = split(head)
    return cpts

def make_dirs(pth, uid, gid, mode):
    """A replacement for os.makedirs that sets ownership as it creates
    directories."""
    if exists(pth):
        return
    root = '/'
    for cpt in components(pth):
        root = join(root, cpt)
        if not exists(root):
            #print "Made %s for (%d, %d)" % (root, uid, gid)
            # XXX: A picky person would be worried about races here.
            mkdir(root, mode)
            chown(root, uid, gid)

def get_mounted_paths():
    mounts = []
    procfile = open('/proc/mounts')
    for line in procfile.readlines():
        mounts.append(line.split(' ')[1])
    procfile.close()
    return mounts

### Processes

def tokenize(data):
    """Take a string and split it on whitespace-strings not contained in
    matched quotes.
    """
    quotes = ''''"'''
    words = []
    state = '\0'
    current_word = stringio()
    c = ''
    for c in data:
        if state == '\0':
            if c in quotes:
                state = c
                continue
            if c.isspace():
                val = current_word.getvalue()
                if val != '':
                    words.append(val)
                    current_word.truncate(0)
                continue

        if state in quotes:
            if c == state:
                state = '\0'
                continue

        current_word.write(c)

    if state in quotes:
        current_word.write(c)

    words.append(current_word.getvalue())

    return words

def lout(cmd):
    if isinstance(cmd, basestring):
        cmd = tokenize(cmd)
    return Popen(cmd, stdout=PIPE).stdout.readlines()

def get_fds():
    for _, _, files in walk("/proc/self/fd"):
        for name in files:
            yield int(name)

### syscall(), clone(), daemon(), faccessat(), unshare()

CSIGNAL              = 0xff      # /* Signal mask to be sent at exit.  */
CLONE_VM             = 0x100     # /* Set if VM shared between processes.  */
CLONE_FS             = 0x200     # /* Set if fs info shared between processes.  */
CLONE_FILES          = 0x400     # /* Set if open files shared between processes.  */
CLONE_SIGHAND        = 0x800     # /* Set if signal handlers shared.  */
CLONE_PTRACE         = 0x2000    # /* Set if tracing continues on the child.  */
CLONE_VFORK          = 0x4000    # /* Set if the parent wants the child to wake it up on mm_release.  */
CLONE_PARENT         = 0x8000    # /* Set if we want to have the same parent as the cloner.  */
CLONE_THREAD         = 0x10000   # /* Set to add to same thread group.  */
CLONE_NEWNS          = 0x20000   # /* Set to create new namespace.  */
CLONE_SYSVSEM        = 0x40000   # /* Set to shared SVID SEM_UNDO semantics.  */
CLONE_SETTLS         = 0x80000   # /* Set TLS info.  */
CLONE_PARENT_SETTID  = 0x100000  # /* Store TID in userlevel buffer before MM copy.  */
CLONE_CHILD_CLEARTID = 0x200000  # /* Register exit futex and memory location to clear.  */
CLONE_DETACHED       = 0x400000  # /* Create clone detached.  */
CLONE_UNTRACED       = 0x800000  # /* Set if the tracing process can't force CLONE_PTRACE on this clone.  */
CLONE_CHILD_SETTID   = 0x1000000 # /* Store TID in userlevel buffer in the child.  */
CLONE_STOPPED        = 0x2000000 # /* Start in stopped state.  */
CLONE_NEWUTS         = 0x04000000# /* New utsname group? */
CLONE_NEWIPC         = 0x08000000# /* New ipcs */
CLONE_NEWUSER        = 0x10000000# /* New user namespace */
CLONE_NEWPID         = 0x20000000# /* New pid namespace */
CLONE_NEWNET         = 0x40000000# /* New network namespace */
CLONE_IO             = 0x80000000# /* Clone io context */

syscall = libc.syscall

# We want the simple one. <Bertl>
# int clone(int number, int flags, void *child_stack)
# int clone(int number, int flags, void *child_stack, int *parent_tidptr, int *child_tidptr)

SYS_clone = 120

def clone(flags):
    """Perform sys_clone() system call, which creates a new process in
    a manner similar to fork().  If flags is SIGCHLD, then this is
    identical to fork().
    XXX: i386 wrapper. <MS>"""
    return libc_errcheck(syscall(SYS_clone, flags, None), None, None)

# int unshare(int flags);
unshare = libc.unshare
unshare.__doc__ = "Libc's unshare call; see unshare (2)"
unshare.argtypes = [c_int]
unshare.restype = c_int
unshare.errcheck = libc_errcheck

# int daemon(int nochdir, int noclose);
daemon = libc.daemon
daemon.__doc__ = "Libc's daemon call; see daemon (3)."
daemon.argtypes = [c_int, c_int]
daemon.restype = c_int
daemon.errcheck = libc_errcheck

AT_FDCWD            = -100	# Special value used to indicate the at functions should use the current working directory.
AT_SYMLINK_NOFOLLOW = 0x100	# Do not follow symbolic links.
AT_REMOVEDIR        = 0x200	# Remove directory instead of unlinking file.
AT_SYMLINK_FOLLOW   = 0x400	# Follow symbolic links.
AT_EACCESS          = 0x200	# Test access permitted for effective IDs, not real IDs.

# int faccessat(int fd, const char* path, int mode, int flags);
faccessat = libc.faccessat
faccessat.__doc__ = "Libc's faccessat call; see faccessat (3)."
faccessat.argtypes = [c_int, c_char_p, c_int, c_int]
faccessat.restype = c_int
faccessat.errcheck = libc_errcheck

### reboot(), shutdown(), and poweroff()

_reboot = libc.reboot
_reboot.__doc__ = "Libc's reboot call; see reboot (2)."
_reboot.argtypes = [c_int]
_reboot.restype = c_int
# XXX: What is reboot()'s errval?
_reboot.errcheck = libc_errcheck

RB_AUTOBOOT    = 0x01234567
RB_HALT_SYSTEM = 0xcdef0123
RB_ENABLE_CAD  = 0x89abcdef
RB_DISABLE_CAD = 0
RB_POWER_OFF   = 0x4321fedc

def reboot():
    """Reboot the machine immediately.
    If no sync is done first, data will be lost."""
    return _reboot(RB_AUTOBOOT)

def halt():
    """Halt the machine immediately.
    If no sync is done first, data will be lost."""
    return _reboot(RB_HALT_SYSTEM)

def poweroff():
    """Poweroff the machine immediately.
    If no sync is done first, data will be lost."""
    return _reboot(RB_POWER_OFF)

### statfs(), fstatfs()

# XXX: The LSB has deprecated statfs() and recommends statvfs() instead.
# Unfortunately, python's os.statvfs() doesn't give any information about file
# system time. <MS>

# XXX: This is a very 32-bit Linux-only python module. <MS>

# struct statfs {
#     long    f_type;     /* type of filesystem (see below) */
#     long    f_bsize;    /* optimal transfer block size */
#     long    f_blocks;   /* total data blocks in file system */
#     long    f_bfree;    /* free blocks in fs */
#     long    f_bavail;   /* free blocks avail to non-superuser */
#     long    f_files;    /* total file nodes in file system */
#     long    f_ffree;    /* free file nodes in fs */
#     fsid_t  f_fsid;     /* file system id */
#     long    f_namelen;  /* maximum length of filenames */
# };

class c_fsid_t(Structure):
    _fields_ = [('val', c_int*3)]

class c_statfs(Structure):
    _fields_ = [('f_type',    c_long),
                ('f_bsize',   c_long),
                ('f_blocks',  c_long),
                ('f_bfree',   c_long),
                ('f_bavail',  c_long),
                ('f_files',   c_long),
                ('f_ffree',   c_long),
                ('f_fsid',    c_fsid_t),
                ('f_namelen', c_long),]

c_statfs_p = POINTER(c_statfs)

# The  function  statfs()  returns  information  about a mounted file system.
# path is the pathname of any file within the mounted filesystem. buf is a
# pointer to a statfs structure defined approximately as follows:

# include <sys/vfs.h>    /* or <sys/statfs.h> */

# int statfs(const char *path, struct statfs *buf);
statfs = libc.statfs
statfs.__doc__ = "Libc's statfs call; see statfs(2)"
statfs.argtypes = [c_char_p, c_statfs_p]
statfs.restype = c_int
statfs.errcheck = libc_errcheck

# int fstatfs(int fd, struct statfs *buf);
fstatfs = libc.fstatfs
fstatfs.__doc__ = "Libc's fstatfs call; see fstatfs(2)"
fstatfs.argtypes = [c_int, c_statfs_p]
fstatfs.restype = c_int
fstatfs.errcheck = libc_errcheck

# File system types:

ADFS_SUPER_MAGIC      = 0xadf5
AFFS_SUPER_MAGIC      = 0xADFF
BEFS_SUPER_MAGIC      = 0x42465331
BFS_MAGIC             = 0x1BADFACE
CIFS_MAGIC_NUMBER     = 0xFF534D42
CODA_SUPER_MAGIC      = 0x73757245
COH_SUPER_MAGIC       = 0x012FF7B7
CRAMFS_MAGIC          = 0x28cd3d45
DEVFS_SUPER_MAGIC     = 0x1373
EFS_SUPER_MAGIC       = 0x00414A53
EXT_SUPER_MAGIC       = 0x137D
EXT2_OLD_SUPER_MAGIC  = 0xEF51
EXT2_SUPER_MAGIC      = 0xEF53
EXT3_SUPER_MAGIC      = 0xEF53
HFS_SUPER_MAGIC       = 0x4244
HPFS_SUPER_MAGIC      = 0xF995E849
HUGETLBFS_MAGIC       = 0x958458f6
ISOFS_SUPER_MAGIC     = 0x9660
JFFS2_SUPER_MAGIC     = 0x72b6
JFS_SUPER_MAGIC       = 0x3153464a
MINIX_SUPER_MAGIC     = 0x137F # /* orig. minix */
MINIX_SUPER_MAGIC2    = 0x138F # /* 30 char minix */
MINIX2_SUPER_MAGIC    = 0x2468 # /* minix V2 */
MINIX2_SUPER_MAGIC2   = 0x2478 # /* minix V2, 30 char names */
MSDOS_SUPER_MAGIC     = 0x4d44
NCP_SUPER_MAGIC       = 0x564c
NFS_SUPER_MAGIC       = 0x6969
NTFS_SB_MAGIC         = 0x5346544e
OPENPROM_SUPER_MAGIC  = 0x9fa1
PROC_SUPER_MAGIC      = 0x9fa0
QNX4_SUPER_MAGIC      = 0x002f
REISERFS_SUPER_MAGIC  = 0x52654973
ROMFS_MAGIC           = 0x7275
SMB_SUPER_MAGIC       = 0x517B
SYSV2_SUPER_MAGIC     = 0x012FF7B6
SYSV4_SUPER_MAGIC     = 0x012FF7B5
TMPFS_MAGIC           = 0x01021994
UDF_SUPER_MAGIC       = 0x15013346
UFS_MAGIC             = 0x00011954
USBDEVICE_SUPER_MAGIC = 0x9fa2
VXFS_SUPER_MAGIC      = 0xa501FCF5
XENIX_SUPER_MAGIC     = 0x012FF7B4
XFS_SUPER_MAGIC       = 0x58465342
_XIAFS_SUPER_MAGIC    = 0x012FD16D

### mount(), umount(), and bindmount()

class MountError(CError):
    pass


MS_RDONLY = 1           # Mount read-only.
MS_NOSUID = 2           # Ignore suid and sgid bits.
MS_NODEV = 4            # Disallow access to device special files.
MS_NOEXEC = 8           # Disallow program execution.
MS_SYNCHRONOUS = 16     # Writes are synced at once.
MS_REMOUNT = 32         # Alter flags of a mounted FS.
MS_MANDLOCK = 64        # Allow mandatory locks on an FS.
S_WRITE = 128           # Write on file/directory/symlink.
S_APPEND = 256          # Append-only file.
S_IMMUTABLE = 512       # Immutable file.
MS_NOATIME = 1024       # Do not update access times.
MS_NODIRATIME = 2048    # Do not update directory access times.
MS_BIND = 4096          # Bind directory at different place.
MS_REC = 16384          # Recursive mount (use with MS_BIN).

MS_MGC_VAL = 0xc0ed0000 # Magic flag number to indicate "new" flags
MS_MGC_MSK = 0xffff0000 # Magic flag number mask

#extern int mount (__const char *__special_file, __const char *__dir,
#          __const char *__fstype, unsigned long int __rwflag,
#          __const void *__data) __THROW;

mount = libc.mount
mount.__doc__ = "Libc's mount call; see mount (2)."
mount.argtypes = [c_char_p, c_char_p, c_char_p, c_uint32, c_char_p]
mount.restype = c_int
mount.errcheck = libc_errcheck

#extern int umount (__const char *__special_file) __THROW;
umount = libc.umount
umount.__doc__ = "Libc's umount call; see umount (2)."
umount.argtypes = [c_char_p]
umount.restype = c_int
umount.errcheck = libc_errcheck

def bindmount(src, tgt, read_only, recursive):
    """Calls mount with appropriate arguments.
    """
    flags = MS_BIND
    if read_only:
        flags |= MS_RDONLY
    if recursive:
        flags |= MS_REC
    mount(src, tgt, '', flags, '')


def read_envdir(envdir, overrides=None):
    # We're going to read an envdir.

    # We'll add k->v pairs to our dictionary for every
    # regular file or symlink to a regular file in opts.envdir.

    # XXX: Several of these calls can block. Do we care?
    # XXX: How many bytes should we read?
    # XXX: Any issues with special characters?

    env = {}
    d = envdir
    if d and isdir(d):
        for p in listdir(d):
            f = None
            try:
                f = os.open(join(d, p), os.O_RDONLY)
                s = fstat(f)
                if S_ISREG(s[ST_MODE]):
                    env[p] = os.read(f, 1024)
            finally:
                if f is not None:
                    os.close(f)
    if overrides:
        for binding in overrides:
            k, v = binding.split('=',1)
            env[k] = v
    return env

class EnvMerge(object):
    def __init__(self):
        self.parent_env = None

    def read_parent_envfile(self):
        ppid = str(os.getppid())
        with open(join('/proc', ppid,'environ')) as f:
            self.parent_env = f.read().split('\0')

    def parent_envvar(self, key):
        if self.parent_env is None:
            self.read_parent_envfile()
        for kv in self.parent_env:
            if kv.startswith(key + "="):
                return kv[len(key +"="):]
        return None

    def prefer_our_envvar(self, key):
        if key in os.environ:
            return os.environ[key]
        return self.parent_envvar(key)