"""
vuln_list.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 gtk
import gobject

import w3af.core.data.kb.knowledge_base as kb

from w3af.core.ui.gui import helpers
from w3af.core.ui.gui.tabs.exploit.utils import get_exploitable_vulns
from w3af.core.ui.gui.tabs.exploit.utils import TextDialogConsumer

from w3af.core.data.kb.info import Info
from w3af.core.data.kb.kb_observer import KBObserver
from w3af.core.controllers.exceptions import (BaseFrameworkException,
                                              NoVulnerabilityFoundException,
                                              ExploitFailedException)


class VulnerabList(gtk.TreeView):
    """A tree showing all the found vulnerabilities.

    :author: Facundo Batista <facundobatista =at= taniquetil.com.ar>
    """
    def __init__(self, w3af, exploitlist):
        """
        :param w3af: The w3af core.
        :param exploitlist: The widget that keeps the list of exploits
        """
        self.w3af = w3af
        self.exploitlist = exploitlist

        # simple List Store with the following columns:
        #     * The string to show (<b>SQL Injection</b>)
        #     * The real name for this vuln (SQL Injection)
        #     * The unique vulnerability id (hash)
        #     * The icon
        self.liststore = gtk.ListStore(str, str, str, gtk.gdk.Pixbuf)
        gtk.TreeView.__init__(self, self.liststore)

        # Stores the uniq ids for vulnerabilities which are exploitable by
        # the currently selected plugin
        self.applicable = []

        # the text & icon column
        tvcolumn = gtk.TreeViewColumn(_("Vulnerabilities"))
        cell = gtk.CellRendererPixbuf()
        tvcolumn.pack_start(cell, expand=False)
        tvcolumn.add_attribute(cell, "pixbuf", 3)
        
        cell = gtk.CellRendererText()
        tvcolumn.pack_start(cell, expand=True)
        tvcolumn.add_attribute(cell, "markup", 0)
        self.append_column(tvcolumn)

        # drag and drop setup, this is the DESTINATION
        target = [("explot-activ", 0, 1)]
        self.enable_model_drag_dest(target, gtk.gdk.ACTION_COPY)
        self.connect("drag-data-received", self._drag_dropped)

        self.connect('cursor-changed', self._changed_selection)

        # get the knowledge base and go live
        kb.kb.add_observer(VulnerabilityObserver(self))

        self.show()

    def _changed_selection(self, *w):
        """
        Changed which exploit is selected.
        
        We want to make the selected vulnerability bold and also let the
        exploit list know about this new selection.
        """
        (path, column) = self.get_cursor()
        vuln = self.get_instance(path)
        
        if vuln is None:
            # No vuln returned by get_instance, make all text regular,
            # in other words, NOT bold.
            for row in self.liststore:
                show, name, uniq_id, icon = row
                row[0] = name
            return
            
        self.exploitlist.set_filter(vuln)

        for row in self.liststore:
            show, name, uniq_id, icon = row
            
            if uniq_id == vuln.get_uniq_id():
                row[0] = make_bold(name)
            else:
                row[0] = name

    def set_filter(self, exploit):
        """Sets a new filter and update the list.

        :param exploit: which exploit is selected/filtered
        """
        vulns = get_exploitable_vulns(exploit)
        
        # Store the vulnerability ids for later
        self.applicable = [v.get_uniq_id() for v in vulns]
        
        # Make bold all the vulnerabilities in the list store which are in
        # self.applicable . In other words, bold the ones which can be
        # exploited (identified by uniq_id
        for row in self.liststore:
            show, name, uniq_id, icon = row
            
            if uniq_id in self.applicable:
                row[0] = make_bold(name)
            else:
                row[0] = name

    def get_name_with_bold(self, info_inst):
        """
        :return: The best obj name possible
        """
        name = info_inst.get_name()
        
        if info_inst.get_uniq_id() in self.applicable:
            showname = make_bold(name)
        else:
            showname = name
            
        return showname, name
        
    def add_item_to_lstore(self, showname, obj_name, obj_severity, idinstance):
        """Updates the GUI with the given information.
        """
        newicon = helpers.KB_ICONS.get(("vuln", obj_severity))
        if newicon is not None:
            newicon = newicon.get_pixbuf()

        self.liststore.append((showname, obj_name, idinstance, newicon))

    def get_instance(self, path):
        """Extracts the instance from the tree.

        :param path: where the user is in the tree
        :return: The instance
        """
        info_uniq_id = self.liststore[path][2]
        info_inst = kb.kb.get_by_uniq_id(info_uniq_id)
        return info_inst
    
    def _drag_dropped(self, tv, drag_context, x, y, selection_data, info, timestamp):
        """Something was dropped (after a drag) on us."""
        droppoint = tv.get_dest_row_at_pos(x, y)
        if droppoint is None:
            return True

        # collect info about source and dest
        (destpath, where) = droppoint
        sourcepath = self.exploitlist.get_cursor()[0]
        sourcerow = self.exploitlist.liststore[sourcepath]

        # it should select a destination row
        if where not in (gtk.TREE_VIEW_DROP_INTO_OR_AFTER, gtk.TREE_VIEW_DROP_INTO_OR_BEFORE):
            self.w3af.mainwin.sb(
                _("You must drop into a row, not in the middle of two"))
            return

        # get real objects
        exploit = self.exploitlist.get_plugin_inst(sourcepath)
        dstvuln = self.get_instance(destpath)
        if dstvuln is None:
            self.w3af.mainwin.sb(
                _("You must select a vulnerability as destination"))
            return

        self._execute_exploit(exploit, dstvuln)
        return

    def _execute_exploit(self, expl, vuln):
        """Exploits a vulnerability.

        This raises a text dialog that informs how the exploit
        is going until it finishes.

        This method is going to:
            a) Create the TextDialog
            b) spawn a thread to launch the exploit process
            c) spawn a thread to read from the output manager queue

        b and c both write messages to the TextDialog.

        :param expl: the exploit to use
        :param vuln: the vulnerability to exploit
        """
        dlg = TextDialogConsumer("Exploit!")
        
        # Start the generator that launches the exploit
        exploit_task = self._launch_exploit(dlg, expl, vuln)
        gobject.idle_add(exploit_task.next)

        return

    def _launch_exploit(self, dlg, expl, vuln):
        """
        Launch the exploit and write messages to the TextDialog.

        :param dlg: The TextDialog.
        """
        # get the info, and see if we can go for it
        dlg.add_message("Checking if this attack plugin can exploit this"
                        " vulnerability...\n")
        vuln_id_list = vuln.get_id()

        yield True

        can_exploit = False

        try:
            can_exploit = expl.can_exploit(vuln_id_list)
        except BaseFrameworkException, e:
            dlg.add_message(_("\nERROR: "))
            dlg.add_message(str(e) + '\n')
            dlg.done()  # set button to sensitive
            dlg.dialog_run()  # wait for user response
            yield False

        if not can_exploit:
            dlg.add_message(_("Sorry, this attack plugin can not exploit this"
                              " vulnerability\n"))
            dlg.done()  # set button to sensitive
            dlg.dialog_run()  # wait for user response
            yield False

        # ok, go for it!
        dlg.add_message(_("Starting exploitation process...\n"))
        yield True

        try:
            expl.exploit(vuln_id_list)
            # print the console messages to the dialog
            yield True
        except NoVulnerabilityFoundException, e:
            dlg.add_message(str(e) + '\n')
        except ExploitFailedException, e:
            dlg.add_message(str(e) + '\n')
        else:
            dlg.add_message(_("Done\n"))
            yield True

        dlg.done()  # set button to sensitive
        dlg.dialog_run()  # wait for user response

        yield False


def make_bold(text):
    return '<b>%s</b>' % text


class VulnerabilityObserver(KBObserver):
    def __init__(self, vulnerability_list):
        self.vulnerability_list = vulnerability_list
        
    def append(self, location_a, location_b, info_inst, ignore_type=False):
        """
        Gets called by the KB when one of the plugins writes something to it.
        
        We've subscribed using kb.kb.add_observer()
        
        :return: None, the information we'll show to the user is stored in an
                 internal variable.
        """
        if not isinstance(info_inst, Info):
            return
        
        obj_name = info_inst.get_name()
        obj_severity = info_inst.get_severity()
        
        # Note that I can't use id(instance) here since the
        # instances that come here are created from the SQLite DB
        # and have different instances, even though they hold the
        # same information.
        idinstance = info_inst.get_uniq_id()

        showname, _ = self.vulnerability_list.get_name_with_bold(info_inst)

        gobject.idle_add(self.vulnerability_list.add_item_to_lstore,
                         showname, obj_name, obj_severity, idinstance)