import errno
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
from time import sleep
import resource

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

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_cmd_as_owner(log, cmd, env, owner_uid, owner_gid, groups, safe_fds):
    pid = fork()
    if not pid:
        log(1, 'Dropping privilege to run cmd as owner.')
        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 cmd.')
        log(1, '%r %r', cmd, env)
        execvpe(cmd[0], cmd, env)
        _exit(59)
    else:
        return pid

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)
        assistant_argv = [assistant, '-v', '-v', '-v', '-u', str(uid), '-e', envdir]
        pid = run_cmd_as_owner(log, assistant_argv, env, owner_uid, owner_gid, groups, safe_fds)
        pid, status = waitpid(pid, 0)
        log(1, 'Assistant returned %d.', status)
        log(1, 'pid %d  uid %d', getpid(), getuid())
        assert not WEXITSTATUS(status)
        return read_envdir(envdir)
    finally:
        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 execvpe\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', 'x11_display_pool',
                  'uid_to_x11_cookie', 'uid_to_x11_display',
                  'uid_to_x11_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_x11(log, spool, owner_uid, 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_x11_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_x11_cookie', str(uid)))

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

    auth_path = join(spool, 'uid_to_x11_auth', str(uid))
    if not exists(auth_path):
        fd, name = mkstemp(prefix='tmp', dir=join(spool, 'uid_to_x11_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)

    x11_server_argv = (['Xvnc', '-DisconnectClients', '-NeverShared',
                        '-localhost', '-SecurityTypes', 'None', '-auth',
                        auth_path, '-geometry', '1024x768', '-depth', '24',
                        ':%d' % display])
                        # -extension XFIXES, on ubuntu?

    run_cmd_as_owner(log, x11_server_argv, {}, owner_uid, owner_gid, [owner_gid], safe_fds)
    sleep(1.0) # XXX: Argh; race condition.

    owner_pw = getpwuid(owner_uid)
    client_env = { 'HOME' : owner_pw.pw_dir, 'DISPLAY' : env['DISPLAY'] }

    x11_client_argv = (['vncviewer', 'localhost:%d' % (5900 + display)])
    run_cmd_as_owner(log, x11_client_argv, client_env, owner_uid, owner_gid, [owner_gid], safe_fds)
    # XXX: need to check that the client successfully connected to the server.

    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"):
        try:
            unshare(CLONE_NEWNET)
        except CError, e:
            if e.errno == errno.EINVAL:
                log(1, "Warning: CLONE_NEWNET not supported")
                return

            raise

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, x11):
    # 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 x11:
        env.update(configure_x11(log, spool, owner_uid, 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)