import os
from os import R_OK, W_OK, X_OK, fork, symlink, unlink, O_CREAT, O_EXCL, chown, chmod
from os import setgroups, setgid, setuid, chdir, umask, execvpe, waitpid, WEXITSTATUS
from os import getpid, getuid, _exit, rename, readlink
from os.path import join, basename, realpath, lexists, exists, dirname
from subprocess import check_call, Popen, PIPE
from stat import S_IFDIR
from tempfile import mkdtemp, mkstemp
from grp import getgrnam, getgrgid
from pwd import getpwuid
from glob import glob
import resource

from rainbow.util import Checker, mount, make_dirs, get_fds, read_envdir
from rainbow.util import unshare, CLONE_NEWNET

def reserve_elt(pool_dir, elt, max_elt, incr, elt_name):
    fd = None
    while elt < max_elt:
        try:
            fd = os.open(join(pool_dir, str(elt)), O_EXCL | O_CREAT, 0600)
            os.close(fd)
            break
        except OSError:
            elt += incr
    if fd is None:
        raise RuntimeError("No " + elt_name + " available.")
    return elt

def reserve_tag(log, spool, tag, tag_map, tag_type, tag_type_plural, min, max, step):
    # Pick an element for yourself.
    pool_dir = join(spool, tag_type+'_pool')
    elt = reserve_elt(pool_dir, min, max, step, tag_type_plural)
    log(1, 'reserved %s (%d) for tag %s', tag_type, elt, tag)

    # Then try to atomically symlink the elt you picked to 'tag' in 'tag_map' dir.
    # If you succeed, then you have the right elt.
    # If you fail, someone else has the right elt so release yours and use theirs.
    tag_path = join(spool, tag_map, tag)
    elt_path = join(pool_dir, str(elt))
    try:
        symlink(str(elt), tag_path)
    except OSError:
        unlink(elt_path)
        elt = int(basename(realpath(tag_path)))
    return elt

def add_uid_to_group(log, spool, owner_uid, uid, gid):
    map_dir = join(spool, 'gid_to_members', str(gid))
    make_dirs(map_dir, 0, 0, 0755)
    owner_path = join(spool, 'gid_to_owner', str(gid))
    if not lexists(owner_path):
        symlink(str(owner_uid), join(spool, 'gid_to_owner', str(gid)))
    # Only join groups that we own.
    assert readlink(owner_path) == str(owner_uid)
    open(join(map_dir, str(uid)), 'w').close()
    open(join(map_dir, str(getpwuid(owner_uid).pw_name)), 'w').close()
    log(1, "added owner (%d) and uid (%d) to gid (%d)", owner_uid, uid, gid)

def reserve_uid(log, spool, owner_uid):
    gid = reserve_elt(join(spool, 'gid_pool'), 10000, 65534, 1, 'gids')
    uid = reserve_elt(join(spool, 'uid_pool'), 10000, 65534, 1, 'uids')
    symlink(str(gid), join(spool, 'uid_to_gid', str(uid)))
    add_uid_to_group(log, spool, owner_uid, uid, gid)
    return (uid, gid)

def reserve_group(log, spool, owner_uid, uid, group):
    gid = reserve_tag(log, spool, group, 'bundle_id_to_gid', 'gid', 'gids', 10000, 65534, 1)
    add_uid_to_group(log, spool, owner_uid, uid, gid)
    return gid

def grab_home(_, spool, uid, __, owner_gid):
    home = join(spool, 'uid_to_home_dir', str(uid))
    make_dirs(home, uid, owner_gid, 0770)
    chown(home, uid, owner_gid)
    # Per discussion with Bert Freudenberg, set the setgid bit on $HOME
    # (i.e. $HOME) so that Sugar can better write inside it. <MS>
    chmod(home, 02770)
    return home

def configure_home(_, spool, home, owner_uid, __, ___, gid, data_group_to_gid):
    for group, gid in data_group_to_gid:
        path = join(spool, 'gid_to_data_dir', str(gid))
        make_dirs(path, owner_uid, gid, 0770)
        chown(path, owner_uid, gid)
        chmod(path, 02770)
        if not lexists(join(home, group)):
            symlink(path, join(home, group))

def mount_fsen(log, _):
    log(1, 'Mounting tmpfsen on /var/tmp and... drat; /tmp kills the X socket. :( <MS>')
    #mount('tmpfs', '/tmp', 'tmpfs', 0, '')
    mount('tmpfs', '/var/tmp', 'tmpfs', 0, '')

def run_assistant(log, assistant, env, owner_uid, owner_gid, uid, groups, safe_fds):
    envdir = None
    try:
        envdir = mkdtemp()
        chown(envdir, owner_uid, owner_gid)
        pid = fork()
    except:
        if envdir: check_call(['/bin/rm', '-rf', envdir])
        raise
    else:
        if not pid:
            log(1, 'Dropping privilege to run assistant.')
            setgroups(groups)
            setgid(owner_gid)
            setuid(owner_uid)
            log(1, 'Closing fds.')
            for fd in get_fds():
                if fd not in safe_fds:
                    try: os.close(fd) # propagate failure from EIO or EBADF.
                    except: pass
            log(1, 'Running assistant.')
            assistant_argv = [assistant, '-v', '-v', '-v', '-u', str(uid), '-e', envdir]
            log(1, '%r %r', assistant_argv, env)
            execvpe(assistant_argv[0], assistant_argv, env)
            _exit(55)
        else:
            try:
                pid, status = waitpid(pid, 0)
                log(1, 'Assistant returned %d.', status)
                assert not WEXITSTATUS(status)
                return read_envdir(envdir)
            finally:
                log(1, 'pid %d  uid %d', getpid(), getuid())
                if envdir: check_call(['/bin/rm', '-rf', envdir])

def launch(log, _, uid, gid, groups, argv, env, cwd, pset, safe_fds):
    # Set appropriate group membership(s), depending on requested permissions
    log(1, 'dropping privilege to (%d, %d, %r)', uid, gid, groups)
    setgroups(groups)
    setgid(gid)
    setuid(uid)

    # Limit various resources
    # Must be done *after* setting uid/gid
    # This should come from the permissions.info file, but this is OK
    #   for now.
    try:
        def set_limit(lim_name, unix_name):
            p = pset.permission_params('lim_' + lim_name)
            if p != None:
                x = int(float(p[0]))
                y = int(float(x*1.10))
                log(1, 'Setting RLIMIT_%s to %d,%d', unix_name, x, y)
                resource.setrlimit(getattr(resource, 'RLIMIT_'+unix_name), (x,y))

        set_limit('nofile', 'NOFILE')
        set_limit('fsize', 'FSIZE')
        set_limit('mem', 'AS')
        set_limit('nproc', 'NPROC')

    except:
        pass # Why would we be throwing exceptions when setting rlimits? <MS>

    log(1, 'chdir to %s' % cwd)
    chdir(cwd)

    log(1, 'umask(0)')
    umask(0)

    log(1, 'about to execve\nargv: %s\nenv: %s', argv, env)
    log(1, 'closing all fds but %s', safe_fds)
    for fd in get_fds():
        if fd not in safe_fds:
            try: os.close(fd) # propagate failure from EIO or EBADF.
            except: pass

    execvpe(argv[0], argv, env)

def check_data_groups(data_groups):
    # XXX: How do I figure out what MAX_PATH_LEN is in python? <MS>
    # XXX: How long are GECOS fields permitted to be? <MS>
    for data_id in data_groups:
        assert data_id and '\0' not in data_id and len(data_id) < 128

def check_argv(argv):
    assert argv

def check_cwd(uid, gid, cwd):
    ck = Checker(cwd, uid, gid)
    assert ck.positive(R_OK | X_OK, S_IFDIR)

def check_spool(spool, owner_uid, owner_gid):
    make_dirs(spool, 0, 0, 0755)
    spool_dirs = ('uid_pool', 'gid_pool', 'uid_to_gid', 'bundle_id_to_gid',
                  'gid_to_data_dir', 'uid_to_home_dir', 'xephyr_display_pool',
                  'uid_to_xephyr_cookie', 'uid_to_xephyr_display',
                  'uid_to_xephyr_auth', 'gid_to_members', 'gid_to_owner')
    for frag in spool_dirs:
        make_dirs(join(spool, frag), 0, 0, 0755)
        ck = Checker(join(spool, frag), owner_uid, owner_gid)
        assert ck.positive(R_OK | X_OK, S_IFDIR)

def check_owner(_, __):
    return True

def check_home_dirs(uid, gid, home, data_group_to_gid):
    for frag, _ in data_group_to_gid:
        ck = Checker(join(home, frag), uid, gid, [x for _, x in data_group_to_gid])
        assert ck.positive(R_OK | W_OK | X_OK, S_IFDIR)

def check_home(uid, gid, home):
    ck = Checker(home, uid, gid)
    assert ck.positive(R_OK | W_OK | X_OK, S_IFDIR)

def maybe_add_gid(log, owner_uid, gid):
    # rainbow should only let you drop privilege.
    owner = getpwuid(owner_uid).pw_name
    members = getgrgid(gid).gr_mem
    log(1, "maybe_add_gid owner: %s members: %s result: %s", owner, members, owner in members)
    return owner in members

def configure_groups(log, owner_uid, groups, gid, data_group_to_gid, recorded_groups, pset):
    groups.insert(0, gid)
    groups += recorded_groups

    for _, data_gid in data_group_to_gid:
        if maybe_add_gid(log, owner_uid, data_gid):
            groups.insert(0, data_gid)

    for cap in ("audio", "video", "serial"):
        try:
            if pset.has_permission(cap):
                cap_gid = getgrnam(cap).gr_gid
                if maybe_add_gid(log, owner_uid, cap_gid):
                    groups.insert(0, cap_gid)
        except Exception, e: log(1, "Skipping permission (%s) because of (%s).", cap, e)
    return list(set(groups))

def configure_xephyr(_, spool, owner_gid, uid, env, safe_fds):
    # XXX: MUST CHECK RETURN VALUES on subprocesses!!!!!
    # XXX: I shouldn't be running these subprocesses as uid 0.
    # XXX: Must get env, fds right!!!!
    cookie_path = join(spool, 'uid_to_xephyr_cookie', str(uid))
    if lexists(cookie_path):
        cookie = readlink(cookie_path)
    else:
        cookie = Popen(["mcookie"], stdout=PIPE).communicate()[0]
        symlink(cookie, join(spool, 'uid_to_xephyr_cookie', str(uid)))

    display_path = join(spool, 'uid_to_xephyr_display', str(uid))
    if lexists(display_path):
        display = int(readlink(display_path))
    else:
        display = reserve_elt(join(spool, 'xephyr_display_pool'), 100, 10000, 2, 'displays')
        symlink(str(display), display_path)

    auth_path = join(spool, 'uid_to_xephyr_auth', str(uid))
    if not exists(auth_path):
        fd, name = mkstemp(prefix='tmp', dir=join(spool, 'uid_to_xephyr_auth'))
        os.close(fd)
        Popen(["xauth", "-f", name], stdin=PIPE).communicate("add :%d . %s\n" % (display, cookie))
        rename(name, auth_path)
        chmod(auth_path, 0640)
        chown(auth_path, 0, owner_gid)

    # NB: Current versions of Xephyr will exit if the display we specified is
    # already in use. This permits us to run Xephyr unconditionally even when
    # we're resuming an old uid. <MS>
    Popen(["Xephyr", "-screen", "800x600x24", "-auth", auth_path, "-reset", "-terminate", ":%d" % display])

    newenv = {'DISPLAY' : ':%d' % display, 'XAUTHORITY' : auth_path}
    return newenv

def configure_network(log, pset):
    log(1, "networking shared with parent: %s", pset.has_permission("network"))
    if not pset.has_permission("network"):
        unshare(CLONE_NEWNET)

def check_uid(_, spool, owner_uid, uid):
    assert 10000 <= uid and uid <= 65534
    assert getpwuid(owner_uid).pw_name in getgrgid(uid).gr_mem

def inject(log, spool, env, argv, cwd, pset, safe_fds, owner_uid, owner_gid,
           uid, groups, data_groups, assistant, xephyr):
    # Note: exceptions are intended to bubble up to the caller and should
    # terminate execution.
    check_data_groups(data_groups)
    check_argv(argv)
    check_owner(owner_uid, owner_gid)

    check_spool(spool, owner_uid, owner_gid)

    if not uid:
        uid, gid = reserve_uid(log, spool, owner_uid)
        home = grab_home(log, spool, uid, gid, owner_gid)
    else:
        check_uid(log, spool, owner_uid, uid)
        pw = getpwuid(uid)
        gid, home = pw.pw_gid, pw.pw_dir
        log(1, "resuming uid (%d) for owner (%d) with gid (%d) and home (%s)", uid, owner_uid, gid, home)

    # XXX: Need to verify ownership and membership before joining data groups.
    recorded_groups = [int(basename(dirname(p))) for p in glob(join(spool, 'gid_to_members', '*', str(uid)))]
    data_group_to_gid = [(group, reserve_group(log, spool, owner_uid, uid, group)) for group in data_groups]
    configure_home(log, spool, home, owner_uid, owner_gid, uid, gid, data_group_to_gid)

    if cwd is None:
        cwd = home
    check_cwd(uid, gid, cwd)
    check_home_dirs(uid, gid, home, data_group_to_gid)
    check_home_dirs(owner_uid, owner_gid, home, data_group_to_gid)
    check_home(uid, gid, home)

    groups = configure_groups(log, owner_uid, groups, gid, data_group_to_gid, recorded_groups, pset)
    if xephyr:
        env.update(configure_xephyr(log, spool, owner_gid, uid, env, safe_fds))
    if assistant:
        env.update(run_assistant(log, assistant, env, owner_uid, owner_gid, uid, groups, safe_fds))

    mount_fsen(log, home)
    configure_network(log, pset)

    launch(log, home, uid, gid, groups, argv, env, cwd, pset, safe_fds)