"""
shell.py

Copyright 2007 Andres Riancho

This file is part of w3af, http://w3af.org/ .

w3af is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation version 2 of the License.

w3af is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with w3af; if not, write to the Free Software
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
"""
import shlex
import gtk
import gobject

import w3af.core.data.kb.knowledge_base as kb
import w3af.core.controllers.output_manager as om

from w3af.core.ui.gui import prompt
from w3af.core.data.kb.shell import Shell
from w3af.core.data.kb.kb_observer import KBObserver
from w3af.core.controllers.exceptions import OSDetectionException


class Shells(gtk.TreeView):
    """The list of shells produced from vulnerabilities.

    :author: Facundo Batista <facundobatista =at= taniquetil.com.ar>
    """
    def __init__(self, w3af):
        self.w3af = w3af

        # create the ListStore, with the shell name and id
        self.liststore = gtk.ListStore(str, str)
        self.listholder = {}

        # create the TreeView using liststore
        super(Shells, self).__init__(self.liststore)

        # create a TreeViewColumn for the text
        tvcolumn = gtk.TreeViewColumn('Shells')
        cell = gtk.CellRendererText()
        tvcolumn.pack_start(cell, True)
        tvcolumn.add_attribute(cell, 'text', 0)
        self.append_column(tvcolumn)

        self.connect('row-activated', self.use_shell)
        kb.kb.add_observer(ShellObserver(self))
        self.show()

    def update_gui(self, shell_show_str, shell_id):
        self.liststore.append([shell_show_str, shell_id])

    def use_shell(self, treeview, path, view_column):
        """Raises a prompt dialog to use the shell."""
        shellid = self.liststore[path][1]
        shell = self.listholder[shellid]
        
        title = "Shell - " + (shell.get_remote_system() or 'Unknown')
        
        def parse_input_run_command(user_input):
            try:
                parsed_command = shlex.split(user_input)
            except ValueError, ve:
                om.out.console('%s' % ve)
            else:
                command = parsed_command[0]
                params = parsed_command[1:]
            
                cmd_stdout = shell.generic_user_input(command, params)
                om.out.console(cmd_stdout)

            om.manager.process_all_messages()

        prompt_text = '%s@%s' % (shell.get_remote_user(),
                                 shell.get_remote_system_name())
        prompt.PromptDialog(title, prompt_text, parse_input_run_command)


class ShellObserver(KBObserver):
    def __init__(self, shell_tree):
        self.shell_tree = shell_tree

    def append(self, location_a, location_b, shell_inst, ignore_type=False):
        """
        Updates the list of shells.
        """
        if not isinstance(shell_inst, Shell):
            return

        shell_id = str(id(shell_inst))
        self.shell_tree.listholder[shell_id] = shell_inst

        # This str() will generate the shell prompt, containing the user@host
        # notation retrieved from the remote web server by the identify_os
        # method.
        #
        # It might take some time, and HTTP requests to retrieve this
        # information, that's why it is a good idea to do it in this method
        # and not in update_gui.
        #
        # Also, this might fail because of a number of reasons, that's why
        # we're catching exceptions and showing a generic prompt on error:
        try:
            shell_inst.identify_os()
            shell_show_str = str(shell_inst)
        except OSDetectionException:
            fmt = '<%s object (ruser: "%s" | rsystem: "%s")>'
            shell_show_str = fmt % (shell_inst.get_name(),
                                    'unknown',
                                    'unknown')

        gobject.idle_add(self.shell_tree.update_gui, shell_show_str, shell_id)