#!/usr/bin/python

import sys
import os
import pwd

def abort(message):
    sys.stderr.write("setuser: %s\n" % message)
    sys.exit(1)


def main():
    '''
    A simple alternative to sudo that executes a command as a user by setting
    the user ID and user parameters to those described by the system and then
    using execvp(3) to execute the command without the necessity of a TTY
    '''

    if os.geteuid() == 0:
        # change user only if current user is root
        username = sys.argv[1]
        try:
            user = pwd.getpwnam(username)
        except KeyError:
            abort("user %s not found" % username)
        os.initgroups(username, user.pw_gid)
        os.setgid(user.pw_gid)
        os.setuid(user.pw_uid)
        os.environ['USER'] = username
        os.environ['HOME'] = user.pw_dir
        os.environ['UID'] = str(user.pw_uid)
    elif os.getenv("SINGULARITY_NAME"):
        # When running in Singularity, remove all the environment 
        # variables associated with python
        for key in ['PYTHONPATH', 'PYTHONUSERBASE', 
            'PYTHONUSERSITE', 'PYTHONNOUSERSITE']:
            if key in os.environ:
                del os.environ[key]

    try:
        # Append $DOCKER_CMD (or, if empty, $SHELL) for Singularity
        if len(sys.argv) <= 2 or len(sys.argv[2].strip()) == 0:
            args = os.getenv("DOCKER_CMD", os.getenv('SHELL'))
        elif sys.argv[2].strip()[:1] =='-':
            args = os.getenv("DOCKER_CMD") + " " + " ".join(sys.argv[2:])
        else:
            args = " ".join(sys.argv[2:])

        # Always source /etc/.profile and $DOCKER_HOME/.profile to set up path
        args = "source /etc/profile && source $DOCKER_HOME/.profile; " + args

        os.execvp('/bin/bash', ['--', '--noprofile', '-c', args])
    except OSError as e:
        abort("cannot execute %s: %s" % (sys.argv[2], str(e)))

if __name__ == '__main__':

    if len(sys.argv) < 2:
        sys.stderr.write("Usage: /sbin/setuser USERNAME [COMMAND args..]\n")
        sys.exit(1)

    main()
