test updates for blocklist - tarka

plugins/BlocklistImport/__init__.py

@@ -71,27 +71,43 @@ class BlocklistImport:
         if fetch:
             print "Fetching",self.config.get('url')
             self.gtkprog.start_download()
+            try:
             filename, headers = urllib.urlretrieve(self.config.get('url'),
                                                    filename=self.blockfile,
                                                    reporthook=self._download_update)
+            except IOError, (errno, strerr):
+                err = ui.GTKError("Couldn't download URL: %s"%strerr)
+                self.gtkprog.stop()
+                return
 
         self.gtkprog.start_import()
 
         self.core.reset_ip_filter()
         ltype = self.config.get('listtype')
         print "importing with",ltype
-        reader = readers[ltype][1](self.blockfile)
 
+        try:
+        reader = readers[ltype][1](self.blockfile)
+        except IOError, (errno, strerr):
+            err = ui.GTKError("Couldn't open blocklist file: %s"%strerr)
+            self.gtkprog.stop()
+            return
+
+        print "Starting import"
         ips = reader.next()
         curr = 0
         while ips and not self.cancelled:
             self.core.add_range_to_ip_filter(*ips)
             ips = reader.next()
             curr += 1
+            if curr % 100 == 0:
+                self.gtkprog.import_prog(text="Imported %s IPs"%curr)
+            else:
             self.gtkprog.import_prog()
 
         reader.close()
         self.gtkprog.end_import()
+        print "Import complete"
 
         self.gtkprog.stop()
 

plugins/BlocklistImport/pg.txt

@@ -0,0 +1,4 @@
+General Electric Company:3.0.0.0-3.255.255.255
+
+s0-0.ciscoseattle.bbnplanet.net:4.0.25.146-4.0.25.148
+

plugins/BlocklistImport/splist.txt

@@ -0,0 +1,4 @@
+Bogon,:0.0.0.0-3.255.255.255
+
+s0-0.c:4.0.25.146-4.0.25.148
+

plugins/BlocklistImport/test.py

@@ -0,0 +1,20 @@
+import unittest
+from text import TextReader, GZMuleReader
+
+
+class ImportTests(unittest.TestCase):
+
+    def testpgtext(self):
+        tr = TextReader("pg.txt")
+        ips = tr.next()
+        self.assertEqual("3.0.0.0", ips[0])
+        self.assertEqual("3.255.255.255", ips[1])
+
+    def testMule(self):
+        mr = GZMuleReader("nipfilter.dat.gz")
+        ips = mr.next()
+        self.assertEqual("0.0.0.0", ips[0])
+        self.assertEqual("3.255.255.255", ips[1])
+
+if __name__ == '__main__':
+    unittest.main()

plugins/BlocklistImport/ui.py

@@ -114,7 +114,9 @@ class GTKProgress(gtk.Dialog):
         self.progress.set_pulse_step(0.0075)
         self.update()
 
-    def import_prog(self):
+    def import_prog(self, text=None):
+        if text:
+            self.progress.set_text(text)
         self.progress.pulse()
         self.update()
 
@@ -142,3 +144,30 @@ class GTKProgress(gtk.Dialog):
     def update(self):
         while gtk.events_pending():
             not gtk.main_iteration(block=True)
+
+
+class GTKError(gtk.Dialog):
+    def __init__(self, message):
+        gtk.Dialog.__init__(self, title="Error",
+                            flags=gtk.DIALOG_MODAL,
+                            buttons=(gtk.STOCK_OK, gtk.RESPONSE_ACCEPT))
+
+        # Setup
+        self.set_border_width(12)
+        self.vbox.set_spacing(6)
+
+        # List source
+        label = gtk.Label()
+        label.set_text(message)
+        self.vbox.pack_start(label)
+
+        self.connect('response', self.ok)
+        self.connect('close', self.cancel)
+
+        self.show_all()
+
+    def ok(self, dialog, response):
+        self.hide_all()
+
+    def cancel(self, dialog):
+        self.hide_all()