root/branch/max/umitGUI/ScriptManager.py @ 735

Revision 735, 14.0 kB (checked in by maxim-gavrilov, 6 years ago)

merge from trunk@732

Line 
1# Copyright (C) 2007 Adriano Monteiro Marques <py.adriano@gmail.com>
2#
3# Author: Maxim Gavrilov <lovelymax@gmail.com>
4#
5# This program is free software; you can redistribute it and/or modify
6# it under the terms of the GNU General Public License as published by
7# the Free Software Foundation; either version 2 of the License, or
8# (at your option) any later version.
9#
10# This program is distributed in the hope that it will be useful,
11# but WITHOUT ANY WARRANTY; without even the implied warranty of
12# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13# GNU General Public License for more details.
14#
15# You should have received a copy of the GNU General Public License
16# along with this program; if not, write to the Free Software
17# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
18
19import os, os.path
20import re
21import sys
22from gtksourceview import SourceView, SourceLanguagesManager, SourceBuffer
23
24#from umitCore.Paths import Search
25Search = None
26from umitCore.I18N import _
27
28import gtk
29from umitGUI.FileChoosers import DirectoryChooserDialog, AllFilesFileChooserDialog
30from higwidgets.higwindows import HIGWindow
31from higwidgets.higboxes import HIGVBox, HIGHBox, HIGSpacer, hig_box_space_holder
32from higwidgets.higexpanders import HIGExpander
33from higwidgets.higlabels import HIGSectionLabel, HIGEntryLabel, HIGDialogLabel
34from higwidgets.higscrollers import HIGScrolledWindow
35from higwidgets.higtextviewers import HIGTextView
36from higwidgets.higbuttons import HIGButton
37from higwidgets.higtables import HIGTable
38from higwidgets.higdialogs import HIGAlertDialog, HIGDialog
39
40# aux alert function
41def _alert(header, text):
42    alert = HIGAlertDialog(
43        message_format='<b>%s</b>' % header,
44        secondary_text=text
45        )
46    alert.run()
47    alert.destroy()
48
49# aux nmap_fetchfile functions
50# XXX: will be moved into Paths.Search or/and into Python/Nmap wrapper
51class NmapFetch(object):
52    def __init__(self):
53        self.dirs = self.__fetchdirs()
54
55    def fetchdirs(self):
56        return self.dirs
57
58    def fetchfile(self, filename):
59        for path in self.fetchdirs():
60            fullpath = os.path.join(path, filename)
61            if Search.check_access(fullpath, os.R_OK):
62                return fullpath
63        return None
64
65    def get_file_list(self):
66        result = []
67        for path in self.fetchdirs():
68            try:
69                for filename in Search.get_file_list(path):
70                    fullpath = os.path.join(path, filename)
71                    if Search.check_access(fullpath, os.R_OK):
72                        result.append(fullpath)
73            except OSError:
74                pass
75        return result
76
77    def nmap_path(self, path):
78        fullpath = os.path.abspath(path)
79        for p in self.fetchdirs():
80            if fullpath.startswith(p):
81                return fullpath[len(p)+1:] # XXX: check +1 (removing last slash) on Windows
82        return fullpath
83
84    def __fetchdirs(self):
85        # standart Nmap searching directories (see nmap.cc:nmap_fetchfile function)
86        def varpath():
87            return os.path.expandvars("${NMAPDIR}")
88        def uidpath():
89            return os.path.join(pwd.getpwuid(os.getuid()).pw_dir, ".nmap")
90        def euidpath():
91            return os.path.join(pwd.getpwuid(os.geteuid()).pw_dir, ".nmap")
92        def userpath():
93            return os.path.expanduser("~")
94        def datadirpath_win():
95            return "c:\\nmap"
96        def datadirpath():
97            return "/usr/share/nmap/"
98        def datadirpath2():
99            return "/usr/local/share/nmap/"
100        def currentpath():
101            return "."
102
103        if sys.platform != 'win32':
104            import pwd
105            checklist = [varpath, uidpath, euidpath, datadirpath, datadirpath2, currentpath]
106        else:
107            checklist = [varpath, userpath, datadirpath_win, currentpath]
108
109        paths = [os.path.abspath(f()) for f in checklist]
110        # XXX: not stable
111        return list(set(paths))
112
113class NmapFetchScripts(NmapFetch):
114    def __init__(self):
115        NmapFetch.__init__(self)
116        self.dirs = [os.path.join(d, "scripts") for d in self.dirs]
117
118    def get_file_list(self):
119        return [f for f in NmapFetch.get_file_list(self) if f.endswith(".nse")]
120
121# Model classes
122# XXX: will be moved to umitCore/
123class ScriptParseException(Exception):
124    pass
125
126class Script(object):
127    def __init__(self, path):
128        f = file(path, 'r')
129        self.data = f.read()
130        f.close()
131
132        self.path = NmapFetchScripts().nmap_path(path)
133        self.id = self._get_attr('id')
134        self.desc = self._get_attr('description')
135
136    def _get_attr(self, attr):
137        r = re.findall(attr + '\s*=\s*"([^\"]+)"', self.data)
138        # XXX: Exception vs False result - that is the Question?!
139        if not r:
140            raise ScriptParseException("Can't parse file %s" % self.path)
141        res = r[0]
142        res = res.replace('\\', ' ')
143        res = res.replace('\n', ' ')
144        res = res.strip()
145        return res
146
147    # for set-element support
148    def __eq__(self, other):
149        return self.path.__eq__(other.path)
150
151    def __hash__(self):
152        return self.path.__hash__()
153
154class ScriptSelection(object):
155    def __init__(self, selected = ""):
156        self.d = dict([(s, False) for s in script_manager])
157        self.set_selected(selected)
158
159    def set_selected(self, selected):
160        for filename in selected.split(";"):
161            script = script_manager.find_by_path(filename)
162            if self.d.has_key(script):
163                self.d[script] = True
164
165    def get_selected(self):
166        return ";".join([script.path for script in self.d.keys() if self.d[script]])
167
168    def is_selected(self, script):
169        return self.d.get(script, False)
170
171    def select(self, script):
172        self.d[script] = True
173
174    def unselect(self, script):
175        self.d[script] = False
176
177class ScriptManager(set):
178    def __init__(self):
179        for filename in NmapFetchScripts().get_file_list():
180            try:
181                self.add_file(filename)
182            except ScriptParseException:
183                pass
184
185    def add_file(self, filename):
186        script = Script(filename)
187        self.add(script)
188
189    def add_dir(self, dirname):
190        res = False
191        for name in Search.get_file_list(dirname):
192            if name.endswith('.nse'):
193                try:
194                    self.add_file(os.path.join(dirname, name))
195                    res = True
196                except ScriptParseException:
197                    pass
198        return res
199
200    def find_by_path(self, filename):
201        for s in self:
202            if s.path == filename:
203                return s
204        return None
205
206# Singletone
207#script_manager = ScriptManager()
208
209# GUI classes
210class ScriptManagerWindow(HIGWindow):
211    def __init__(self):
212        HIGWindow.__init__(self)
213        self.set_title(_("Script Manager"))
214        self.set_position(gtk.WIN_POS_CENTER)
215        self.create_widgets()
216        self.update_model()
217
218    def create_list(self):
219        scroll = HIGScrolledWindow()
220        scroll.set_shadow_type(gtk.SHADOW_IN)
221        self.model = gtk.ListStore(object, str)
222        self.list_view = gtk.TreeView(self.model)
223        scroll.add(self.list_view)
224
225        self.model.set_sort_column_id(1, gtk.SORT_ASCENDING)
226        self.list_view.set_headers_visible(False)
227        cell = gtk.CellRendererText()
228        col = gtk.TreeViewColumn('id', cell, text=1)
229        self.list_view.append_column(col)
230        self.list_view.set_search_column(1)
231        self.list_view.connect('cursor-changed', self.id_select_cb)
232        return scroll
233
234    def create_text(self):
235        scroll = HIGScrolledWindow()
236        scroll.set_shadow_type(gtk.SHADOW_IN)
237        self.text_buffer = SourceBuffer()
238        self.text_buffer.set_highlight(True)
239        lang = SourceLanguagesManager().get_language_from_mime_type('text/x-lua')
240        self.text_buffer.set_language(lang)
241        text_view = SourceView(self.text_buffer)
242        text_view.set_editable(True)
243        text_view.set_auto_indent(True)
244        scroll.add(text_view)
245        return scroll
246
247    def create_buttons(self):
248        hbox = HIGHBox()
249        btn_add_file = HIGButton(_("Add file"), stock=gtk.STOCK_OPEN)
250        btn_add_dir = HIGButton(_("Add dir"), stock=gtk.STOCK_DIRECTORY)
251        btn_delete = HIGButton(stock=gtk.STOCK_DELETE)
252        btn_close = HIGButton(stock=gtk.STOCK_CLOSE)
253        hbox.pack_start(btn_add_file)
254        hbox.pack_start(btn_add_dir)
255        hbox.pack_start(btn_delete)
256        hbox.pack_start(btn_close)
257        hbox.set_border_width(5)
258        hbox.set_spacing(6)
259        btn_add_file.connect('clicked', self.add_file_cb)
260        btn_add_dir.connect('clicked', self.add_dir_cb)
261        btn_delete.connect('clicked', self.delete_cur_cb)
262        btn_close.connect('clicked', lambda x,y=None:self.destroy())
263        return hbox
264
265    def create_widgets(self):
266        self.set_size_request(600, 400)
267        hpane = gtk.HPaned()
268        hpane.pack1(self.create_list(), False)
269        hpane.pack2(self.create_text(), True)
270        #hbox = HIGHBox()
271        #hbox.pack_start(self.create_list())
272        #hbox.pack_start(self.create_text())
273        main_vbox = HIGVBox()
274        main_vbox.pack_start(hpane, True, True)
275        main_vbox.pack_start(self.create_buttons(), False, False)
276        self.add(main_vbox)
277
278    def update_model(self):
279        self.model.clear()
280        for script in script_manager:
281            self.model.append((script, script.id))
282
283    def id_select_cb(self, list_view):
284        (model, it) = list_view.get_selection().get_selected()
285        if it:
286            script = model[it][0]
287            self.text_buffer.set_text(script.data)
288
289    def add_file_cb(self, widget):
290        file_chooser = AllFilesFileChooserDialog(_("Select Script"))
291        file_chooser.run()
292        name = file_chooser.get_filename()
293        file_chooser.destroy()
294        if name:
295            self.add_file(name)
296
297    def add_dir_cb(self, widget):
298        file_chooser = DirectoryChooserDialog(_("Select Scripts Directory"))
299        file_chooser.run()
300        dirname = file_chooser.get_filename()
301        file_chooser.destroy()
302        if dirname:
303            self.add_dir(dirname)
304
305    def delete_cur_cb(self, widget):
306        (model, it) = self.list_view.get_selection().get_selected()
307        if it:
308            cur = model.get_path(it)[0]
309            script_manager.remove(model[it][0])
310            model.remove(it)
311            self.list_view.get_selection().select_path((cur,))
312            self.id_select_cb(self.list_view)
313
314    def add_file(self, filename):
315        if Search.check_access(filename, os.R_OK):
316            try:
317                script_manager.add_file(filename)
318                self.update_model()
319            except:
320                _alert(
321                    _('File is not a NSE script file'),
322                    _("Selected file is not a NSE script file. Umit can not \
323parse this file. Please, select another."))
324        else:
325            _alert(
326                _('Can not open selected file'),
327                _("Umit can not open selected file. Please, select another."))
328
329    def add_dir(self, dirname):
330        if script_manager.add_dir(dirname):
331            self.update_model()
332        else:
333            _alert(
334                _('Can not find any NSE script files in the directory'),
335                _("Umit can not find any NSE script files in the selected \
336directory. Please, select another."))
337
338class ScriptChooserDialog(HIGDialog):
339    def __init__(self, selected = ""):
340        HIGDialog.__init__(self, _("Select Necessory Scripts"), None,
341                           gtk.DIALOG_MODAL,
342                           (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
343                            gtk.STOCK_OK, gtk.RESPONSE_OK))
344        self.selection = ScriptSelection(selected)
345        self.set_size_request(400, 400)
346        self.create_widgets()
347        self.update_model()
348
349    def create_list(self):
350        scroll = HIGScrolledWindow()
351        scroll.set_shadow_type(gtk.SHADOW_IN)
352        self.model = gtk.ListStore(object, bool, str)
353        self.list_view = gtk.TreeView(self.model)
354        scroll.add(self.list_view)
355
356        self.model.set_sort_column_id(2, gtk.SORT_ASCENDING)
357        self.list_view.set_headers_visible(False)
358        self.list_view.set_search_column(2)
359
360        cell = gtk.CellRendererToggle()
361        cell.connect('toggled', self.toggled_cb, self.model)
362        col = gtk.TreeViewColumn('b', cell, active=1)
363        self.list_view.append_column(col)
364
365        cell = gtk.CellRendererText()
366        col = gtk.TreeViewColumn('id', cell, text=2)
367        self.list_view.append_column(col)
368
369        self.list_view.connect('cursor-changed', self.id_select_cb)
370        return scroll
371
372    def create_text(self):
373        scroll = HIGScrolledWindow()
374        scroll.set_shadow_type(gtk.SHADOW_IN)
375        text_view = HIGTextView()
376        text_view.set_editable(False)
377        scroll.add(text_view)
378        self.text_buffer = text_view.get_buffer()
379        return scroll
380
381    def create_widgets(self):
382        vpane = gtk.VPaned()
383        vpane.pack1(self.create_list(), True)
384        vpane.pack2(self.create_text(), False)
385        self.vbox.add(vpane)
386        self.show_all()
387
388    def update_model(self):
389        self.model.clear()
390        for script in script_manager:
391            self.model.append((
392                script,
393                self.selection.is_selected(script),
394                script.id))
395
396    def id_select_cb(self, list_view):
397        (model, it) = list_view.get_selection().get_selected()
398        if it:
399            script = model[it][0]
400            self.text_buffer.set_text(script.desc)
401
402    def toggled_cb(self, cell, path, model):
403        model[path][1] = not model[path][1]
404        if model[path][1]:
405            self.selection.select(model[path][0])
406        else:
407            self.selection.unselect(model[path][0])
408
409    def get_scripts(self):
410        return self.selection.get_selected()
411
412if __name__ == "__main__":
413    sm = ScriptManagerWindow()
414    sm.show_all()
415    sm.connect("destroy", gtk.main_quit)
416    gtk.main()
417
418    #sd = ScriptChooserDialog("")
419    #if sd.run() == gtk.RESPONSE_OK:
420    #    print "OK:", sd.get_scripts()
421    #else:
422    #    print "Cancel"
423    #sd.destroy()
Note: See TracBrowser for help on using the browser.