~epidermis/epidermis/devel

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright © David D Lowe 2008-2011
# This program is free software under the terms of GPLv2. For more
# information, see the file named COPYING distributed with this project.

"""Contains the Shell class and its subclasses"""

from __future__ import division, absolute_import
from . import const
from .const import logger, MY_CACHE_HOME, MY_DATA_HOME, PIGMENT_TYPES
import pexpect
import os, sys
import time
from ConfigParser import SafeConfigParser
import ConfigParser
from gettext import gettext as _
import subprocess
import dbus
import warnings
import getpass
from . import handy

DEBUG_SHELL = True
if not const.DEBUG:
    DEBUG_SHELL = False

def bash_prepare(string): 
    """Prepare a string for bash by surrounding it with single quotes.
    
    Returns the formatted string.
    
    >>> shell.bash_prepare("two'words")
    "'two'\\''words'"
    >>> print shell.bash_prepare("two'words")
    'two'\''words'
    
    """ 
    return "'" + string.replace("'", r"'\''") + "'"

class Shell():
    """An object which can run commands
    This is an "abstract" class, only its subclasses should be used.
    
    """
    def __init__(self):
        self.ready = False
    
    def prepare(self):
        """Prepare the shell, asking user for authentication if necessary"""
        handy.abstract()
    
    def do(self, command, intrShell=False, cwd=None):
        """Run the command in this shell.
        
        Keyword arguments:
        command -- a list of command arguments
        intrShell -- if True, allow the shell to interpret special 
                     characters such as *
        cwd -- the current directory for the command
        
        Returns the exit code of the command.
        
        """
        handy.abstract()
    
    def check_do(self, command, intrShell=False, cwd=None):
        """Run the command in this shell. If it exits with a non-zero code, throw a
        CommandError exception.
        
        """
        exitcode = self.do(command, intrShell, cwd)
        if exitcode != 0:
            raise CommandError(command, exitcode)
        
    def set_out_file(self, outfile):
        """Set the file object where stdout output is put"""
        handy.abstract()
    
    def exit(self):
        """Exit the shell"""
        handy.abstract()

class RootShell(Shell):
    """A shell that runs commands as root. It will automatically use
    BashRootShell or PolicyKitShell in its backend depending on the
    user's settings"""
    
    def __init__(self):
        Shell.__init__(self)
        configParser = SafeConfigParser()
        configParser.read("/etc/epidermis.conf")
        try:
            val = configParser.get("epidermis","access_root")
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError),ee:
            print >> sys.stderr, "Error parsing /etc/epidermis.conf, using gksu"
            val = "gksu"
        if val == "policykit":
            self._shell = PolicyKitShell()
        elif val == "gksu":
            self._shell = BashRootShell()
        else:
            print >> sys.stderr, "Error parsing /etc/epidermis.conf, using gksu"
            self._shell = BashRootShell()
        
        def wrapper(func):
            def execute_and_update_ready(*args, **kwargs):
                func(*args, **kwargs)
                self.ready = self._shell.ready
            return execute_and_update_ready
        
        self.prepare = wrapper(self._shell.prepare)
        self.do = wrapper(self._shell.do)
        self.check_do = wrapper(self._shell.check_do)
        if hasattr(self._shell,"set_err_file"):
            self.set_err_file = wrapper(self._shell.set_err_file)
        self.set_out_file = wrapper(self._shell.set_out_file)
        self.exit = wrapper(self._shell.exit)
    
    def __str__(self):
        return "RootShell instance, with _shell: %s" % repr(self._shell)
    
    

class BashRootShell(Shell):
    """The Shell object which has access to a bash terminal run as root.
    It uses pexpect, and it runs sudo bash -i"""
    
    def __init__(self, prompt="epidermisbash:"):
        """Initialise this instance.
        
        Keyword arguments:
        prompt -- the prompt line of the command line (string)
        
        """
        Shell.__init__(self)
        self.prompt = prompt
        self.outfile = os.tmpfile()
        self.gksubash, self.gksued = None, None
        self.spawn = None
        self.ready = False
        self.gksubash = handy.get_data_file("gksubash.sh")
        self.gksued = handy.get_data_file("gksued.sh")
        if not os.path.exists(self.gksubash):
            raise(Exception("Cannot find gksubash.sh"))
        if not os.path.exists(self.gksued):
            raise(Exception("Cannot find gksued.sh"))
        
    def prepare(self):
        """Initialise the bash shell as root, asking user for 
        password graphically if necessary using gksu.
        
        This function raises a ShellNotReadyException if it fails to 
        prepare itself.
        
        # Unfortunately, due to bug #244930, gksu cannot run bash -i correctly
        # This workaround uses sudo, but still asks for the password
        # graphically, if required
        # See bug: https://bugs.launchpad.net/ubuntu/+source/gksu/+bug/244930
        
        """
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            import gtk
            display = gtk.gdk.display_get_default()
        
        self.spawn = pexpect.spawn("sh", [self.gksubash, self.gksued])
        self.spawn.logfile = self.outfile
        res = self.spawn.expect(["EPIDERMIS_PASS", "EPIDERMIS_SUDO_RUN_SUCCESSFULLY", \
            "EPIDERMIS_SUDO_FINISHED"])
        self.ready = True
        while res != 1: # sudo didn't run successfully
            if res == 2:
                # sudo exited unexpectedly
                self.ready = False
                raise(ShellNotReadyException())
                break
            # I need a password, use gksu gui or getpass to get one
            if display:
                pop = os.popen("gksu -p --message " + bash_prepare(_("Enter password to allow Epidermis to access root privileges")))
                password = None
                for line in pop.readlines():
                    if "sn_launcher_context_complete" in line:
                        continue
                    if len(line) == 0:
                        continue
                    password = line[:-1]
                    break
            else:
                try:
                    print _("Enter password (press Ctrl-D to cancel")
                    password = getpass.getpass(_("[sudo] password for %(username)s: ") % {"username":getpass.getuser()})
                except EOFError:
                    password = None
                    print ""
            if not password:
                retry = self._yes_no_box(_("Could not retrieve password, try again?"))
                if not retry:
                    self.ready = False
                    raise(ShellNotReadyException())
                    break
                else:
                    continue
            # try this password, disable outfile to avoid printing password
            oldoutfile = self.spawn.logfile
            self.spawn.logfile = os.tmpfile()
            self.spawn.sendline(password)
            self.spawn.logfile = oldoutfile
            # recheck to see if successful
            res = self.spawn.expect(["EPIDERMIS_PASS", "EPIDERMIS_SUDO_RUN_SUCCESSFULLY", \
                "EPIDERMIS_SUDO_FINISHED"])
    
        self.spawn.readline()
        
        
        if not self.ready:
            raise(ShellNotReadyException())
        else:
            if hasattr(self, "outfile"):
                self.spawn.logfile = self.outfile
            ##self.spawn.sendline(" PS1=" + self.prompt)
            ##self.spawn.readline()
            ##self.spawn.expect(self.prompt)
            self.do(" PS1=" + self.prompt)
            self.do(" HISTCONTROL=\"ignoreboth\"")
            return True
    
    def set_out_file(self, outfile):
        if not self.spawn is None:
            Sfile = outfile
        self.outfile = outfile
    
    
    def do(self, command, intrShell=False, cwd=None):
        if not self.ready:
            raise(ShellNotReadyException())
            
        cmd = ""
        if isinstance(command, list):
            for item in command:
                if intrShell:
                    cmd = cmd + item.replace(" ", r"\ ") + " "
                else:
                    cmd = cmd + bash_prepare(item) + " "
            cmd = cmd[:-1]
        elif isinstance(command, str):
            cmd = command
        else:
            raise(Exception("command must be list or string"))
        
        if DEBUG_SHELL:
            logger.debug("debug+++root do: " + cmd)
        if not cwd is None:
            self.spawn.sendline(" cd " + bash_prepare(cwd))
            self.spawn.readline()
            self.spawn.expect(self.prompt)
        self.spawn.sendline(" " + cmd) # space is necessary so that the command is not 
                                       # recorded in history
        self.spawn.readline() # reads the line in which the command was entered
        self.spawn.expect(self.prompt)
        self.spawn.sendline(" echo $?")
        self.spawn.readline()
        exitcode = self.spawn.readline().strip()
        self.spawn.expect(self.prompt)
        if DEBUG_SHELL:
            logger.debug("debug++root exitcode: " + exitcode)
        return int(exitcode)
            
    
    def _yes_no_box(self, message):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            import gtk
            display = gtk.gdk.display_get_default()
        if display:
            dd = gtk.MessageDialog(type=gtk.MESSAGE_ERROR, buttons=gtk.BUTTONS_YES_NO, \
                message_format=message)
            res = dd.run()
            dd.hide()
            dd.destroy()
            if res == gtk.RESPONSE_YES:
                return True
            else:
                return False
        else:
            while 1:
                print message, "[Y/n]",
                res = raw_input()
                if res.strip().lower() in ["", "y"]:
                    return True
                elif res.strip().lower() == "n":
                    return False
    
    def exit(self):
        if self.spawn.isalive():
            self.spawn.sendline(" exit")
            self.spawn.readline()
            self.spawn.expect("exit")
        
        if self.spawn.isalive():
            time.sleep(0.3)
            if self.spawn.isalive():
                self.spawn.terminate()
        self.ready = False

class PolicyKitShell(Shell):
    """An object which can run commands using a DBus service and 
    PolicyKit authentication
    
    """
    
    def __init__(self):
        self.ready = False
        self.outFile, self.errFile = None, None
        self.pshell = None
        
    def prepare(self):
        bus = dbus.SystemBus()
        self.pshell = bus.get_object("org.tuxfamily.epidermis.Shell",
            "/org/tuxfamily/epidermis/PShell")
        rt = self.pshell.prepare(dbus_interface="org.tuxfamily.epidermis.PShellInterface")
        self.ready = rt
        if not rt:
            raise(ShellNotReadyException())
        return self.ready
    
    def set_out_file(self, outfile):
        self.outFile = outfile
    
    def set_err_file(self, errfile):
        self.errFile = errfile
    
    def do(self, command, intrShell=False, cwd=None):
        if not isinstance(command, list):
            raise(Exception("Not list"))
        global DEBUG_SHELL
        if DEBUG_SHELL:
            cmdStr = ""
            for it in command:
                cmdStr = cmdStr + it + " "
            if len(cmdStr) > 0:
                cmdStr = cmdStr[:-1]
            logger.debug("debug+++pshell do: " + cmdStr)
        if cwd is None:
            cwd = "/tmp/"
        exitcode, out, err = self.pshell.do(command,intrShell,
            cwd, dbus_interface="org.tuxfamily.epidermis.PShellInterface")
        if not self.outFile is None:
            self.outFile.write(out)
        if not self.errFile is None:
            self.errFile.write(err)
        return exitcode
    
    def exit(self):
        self.pshell.leave(dbus_interface="org.tuxfamily.epidermis.PShellInterface")
        self.pshell = None
        self.ready = False


class SubprocessShell(Shell):
    """An object which can run commands using subprocess.call as current user"""
    def __init__(self):
        self.ready = True
        self.outFile = sys.stdout
        self.errFile = sys.stderr
    
    def prepare(self):
        """Prepare the shell"""
        self.ready = True
        self.outFile = sys.stdout
        return self.ready
    
    def do(self, command, intrShell=False, cwd=None):
        if not isinstance(command, list):
            raise(Exception("Not list"))
        global DEBUG_SHELL
        if DEBUG_SHELL:
            cmdStr = ""
            for it in command:
                cmdStr = cmdStr + it + " "
            if len(cmdStr) > 0:
                cmdStr = cmdStr[:-1]
            logger.debug("debug+++subprocess do: " + cmdStr)
        exitcode = subprocess.call(command, shell=intrShell, cwd=cwd, stdout=self.outFile, \
            stderr=self.errFile)
        return exitcode
        
    def set_out_file(self, outfile):
        self.outFile = outfile
    
    def set_err_file(self, errfile):
        self.errFile = errfile
    
    def exit(self):
        """Exits the shell, this shell is always ready so this method has no
        function"""
        self.ready = True

class ShellNotReadyException(Exception):
    def __str__(self):
        return "Shell has not been prepared"

class CommandError(Exception):
    def __init__(self, command, exitcode):
        self.command = command
        self.exitcode = int(exitcode)
    def __str__(self):
        return "Command %s returned non-zero status %d" % (repr(self.command), self.exitcode)

def testshell():
    """Test some of BashRootShell's and PolicyKitShell's
    functionality
    
    """
    bashShell = BashRootShell()
    bashShell.set_out_file(sys.stdout)
    bashShell.prepare()
    if bashShell.ready == False:
        print >> sys.stderr, "bash shell is not ready, user cancelled?"
        return
    print "Bash shell is ready"
    bashShell.do(["whoami"])
    bashShell.exit()
    
    pkShell = PolicyKitShell()
    pkShell.set_out_file(sys.stdout)
    pkShell.prepare()
    if pkShell.ready == False:
        print >> sys.stderr, "policykit shell is not ready, user cancelled?"
        return
    print "PolicyKit shell is ready"
    pkShell.do(["whoami"])
    pkShell.exit()
    
    rtShell = RootShell()
    rtShell.set_out_file(sys.stdout)
    rtShell.prepare()
    if rtShell.ready == False:
        print >> sys.stderr, "root shell is not ready"
        return
    print "Root shell is ready"
    rtShell.do(["whoami"])
    rtShell.exit()

def determine_path ():
    """Borrowed from wxglade.py"""
    try:
        root = __file__
        if os.path.islink (root):
            root = os.path.realpath (root)
        return os.path.dirname (os.path.abspath (root))
    except Exception, ee:
        print >> sys.stderr, "I'm sorry, but something is wrong."
        print >> sys.stderr, "There is no __file__ variable. Please contact the author."
        sys.exit (1)
        raise(ee)