add support for importing mrt format ribs

remove default=0 value for timestamps, as that is not a valid initial
value for a datetime object.

since there is no default argument, the call to the Timestamp
constructor needs to specify the value.

svn path=/branches/tk161/; revision=4177

3 files changed, 205 insertions, 71 deletions

diff --git a/rpkid/portal-gui/scripts/rpkigui-import-routes.py b/rpkid/portal-gui/scripts/rpkigui-import-routes.py
index 3df4b83f..77c4a569 100644
--- a/rpkid/portal-gui/scripts/rpkigui-import-routes.py
+++ b/rpkid/portal-gui/scripts/rpkigui-import-routes.py
@@ -14,120 +14,253 @@
 # PERFORMANCE OF THIS SOFTWARE.
 #
 
-import sys, itertools, re
-import struct
+import itertools
 import _mysql_exceptions
+import optparse
+import os.path
+import re
+import sys
+import struct
+import subprocess
 
 from django.db import transaction, connection
 
 from rpki.resource_set import resource_range_ipv4, resource_range_ipv6
 import rpki.gui.app.timestamp
 
-f = open(sys.argv[1])
-
-prefixes = {}
-
-ip_re = re.compile(r'^[0-9a-fA-F:.]+/\d{1,3}$')
+# globals
+DEBUG=False
+VERBOSE=False
+BGPDUMP='bgpdump'
+PREFIXES = {}
 
 class InvalidPrefix(Exception):
     pass
 
-last_prefix = None
-last_asn = None
+def parse_text(f):
+    ip_re = re.compile(r'^[0-9a-fA-F:.]+/\d{1,3}$')
+    last_prefix = None
+    last_asn = None
 
-for row in itertools.islice(f, 5, None):
-    try:
-        cols = row.split()
+    for row in itertools.islice(f, 5, None):
+        try:
+            cols = row.split()
 
-        prefix = cols[1]
+            prefix = cols[1]
+
+            # index -1 is i/e/? for igp/egp
+            origin_as = cols[-2]
+
+            # FIXME: skip AS_SETs
+            if origin_as[0] == '{':
+                continue
+
+            # the output may contain multiple paths to the same origin.
+            # if this is the same prefix as the last entry, we don't need
+            # to validate it again.
+            if prefix != last_prefix:
+                # validate the prefix since the "sh ip bgp" output is sometimes corrupt
+                # by no space between the prefix and the next hop IP address.
+
+                if not ip_re.match(prefix):
+                    net, bits = prefix.split('/')
+                    if len(bits) > 2 and int(bits[0]) <= 3:
+                        print 'mask for %s looks fishy...' % prefix,
+                        prefix = '%s/%s' % (net, bits[0:2])
+                        print 'assuming it should be %s' % prefix
+                    if not ip_re.match(prefix):
+                        raise InvalidPrefix(prefix)
+                last_prefix = prefix
+            elif origin_as == last_asn:
+                # we are only interested in origins, so skip alternate paths
+                # to same origin as last entry.
+                continue
+            last_asn = origin_as
 
-        # index -1 is i/e/? for igp/egp
-        origin_as = cols[-2]
+            asns = PREFIXES.get(prefix)
+            if not asns:
+                asns = set()
+                PREFIXES[prefix] = asns
+            asns.add(int(origin_as))
 
-        # FIXME: skip AS_SETs
-        if origin_as[0] == '{':
+        except InvalidPrefix, e:
+            print >>sys.stderr, 'skipping bad entry: ' + row,
+            print >>sys.stderr, e
+
+def parse_mrt(f):
+    # filter input through bgpdump
+    pipe = subprocess.Popen([BGPDUMP, '-m', '-v', '-'], stdin=f, stdout=subprocess.PIPE)
+
+    last_prefix = None
+    last_as = None
+    for e in pipe.stdout.readlines():
+        a = e.split('|')
+        prefix = a[5]
+        try:
+            origin_as = int(a[6].split()[-1])
+        except ValueError:
+            # skip AS_SETs
             continue
 
-        # the output may contain multiple paths to the same origin.
-        # if this is the same prefix as the last entry, we don't need
-        # to validate it again.
         if prefix != last_prefix:
-            # validate the prefix since the "sh ip bgp" output is sometimes corrupt
-            # by no space between the prefix and the next hop IP address.
-
-            if not ip_re.match(prefix):
-                net, bits = prefix.split('/')
-                if len(bits) > 2 and int(bits[0]) <= 3:
-                    print 'mask for %s looks fishy...' % prefix,
-                    prefix = '%s/%s' % (net, bits[0:2])
-                    print 'assuming it should be %s' % prefix
-                if not ip_re.match(prefix):
-                    raise InvalidPrefix(prefix)
             last_prefix = prefix
-        elif origin_as == last_asn:
-            # we are only interested in origins, so skip alternate paths
-            # to same origin as last entry.
+        elif last_as == origin_as:
             continue
-        last_asn = origin_as
+        last_as = origin_as
 
-        asns = prefixes.get(prefix)
+        asns = PREFIXES.get(prefix)
         if not asns:
             asns = set()
-            prefixes[prefix] = asns
-        asns.add(int(origin_as))
+            PREFIXES[prefix] = asns
+        asns.add(origin_as)
 
-    except InvalidPrefix, e:
-        print >>sys.stderr, 'skipping bad entry: ' + row,
-        print >>sys.stderr, e
-
-f.close()
+    pipe.wait()
+    if pipe.returncode:
+        raise ProgException('bgpdump exited with code %d' % pipe.returncode)
 
 def commit():
+    "Write the PREFIXES dict into the appropriate database table."
+    # auto-detect the IP version
+    sample_ip = PREFIXES.iterkeys().next()
+    if ':' in sample_ip:
+        ip_version = 6
+        range_class = resource_range_ipv6
+        value_xform = lambda v: struct.pack('!QQ', (long(v) >> 64) & 0xffffffffffffffffL, long(v) & 0xFFFFFFFFFFFFFFFFL)
+        table = 'routeview_routeoriginv6'
+    else:
+        ip_version = 4
+        range_class = resource_range_ipv4
+        value_xform = long
+        table = 'routeview_routeorigin'
+
+    log('Inserting data into table...')
+
+    debug('auto-detected IP version %d prefixes' % ip_version)
+
     cursor = connection.cursor()
 
     try:
-        print 'Dropping existing staging table...'
-        cursor.execute('DROP TABLE IF EXISTS routeview_routeorigin_new')
+        debug('Dropping existing staging table...')
+        cursor.execute('DROP TABLE IF EXISTS %s_new' % table)
     except _mysql_exceptions.Warning:
         pass
 
-    print 'Creating staging table...'
-    cursor.execute('CREATE TABLE routeview_routeorigin_new LIKE routeview_routeorigin')
+    debug('Creating staging table...')
+    cursor.execute('CREATE TABLE %(table)s_new LIKE %(table)s' % { 'table': table })
 
-    print 'Disabling autocommit...'
+    debug('Disabling autocommit...')
     cursor.execute('SET autocommit=0')
 
-    print 'Adding rows to table...'
-    for prefix, asns in prefixes.iteritems():
-        family = 6 if ':' in prefix else 4
-        cls = resource_range_ipv6 if family == 6 else resource_range_ipv4
-        rng = cls.parse_str(prefix)
-
-        if family == 4:
-            xform = long
-        else:
-            xform = lambda v: struct.pack('!QQ', (long(v) >> 64) &0xffffffffffffffffL, long(v) & 0xFFFFFFFFFFFFFFFFL)
+    debug('Adding rows to table...')
+    sql = "INSERT INTO %s_new SET asn=%%s, prefix_min=%%s, prefix_max=%%s" % table
 
-        cursor.executemany("INSERT INTO routeview_routeorigin_new SET asn=%s, prefix_min=%s, prefix_max=%s",
-                    [(asn, xform(rng.min), xform(rng.max)) for asn in asns])
+    for prefix, asns in PREFIXES.iteritems():
+        rng = range_class.parse_str(prefix)
+        cursor.executemany(sql, [(asn, value_xform(rng.min),
+            value_xform(rng.max)) for asn in asns])
 
-    print 'Committing...'
+    debug('Committing...')
     cursor.execute('COMMIT')
 
     try:
-        print 'Dropping old table...'
-        cursor.execute('DROP TABLE IF EXISTS routeview_routeorigin_old')
+        debug('Dropping old table...')
+        cursor.execute('DROP TABLE IF EXISTS %s_old' % table)
     except _mysql_exceptions.Warning:
         pass
 
-    print 'Swapping staging table with live table...'
-    cursor.execute('RENAME TABLE routeview_routeorigin TO routeview_routeorigin_old, routeview_routeorigin_new TO routeview_routeorigin')
+    debug('Swapping staging table with live table...')
+    cursor.execute('RENAME TABLE %(table)s TO %(table)s_old, %(table)s_new TO %(table)s' % { 'table': table })
 
     transaction.commit_unless_managed()
 
-commit()
+    debug('Updating timestamp metadata...')
+    rpki.gui.app.timestamp.update('bgp_v4_import' if ip_version == 4 else 'bgp_v6_import')
+
+def debug(s):
+    if DEBUG: print s
+
+def log(s):
+    if VERBOSE: print s
+
+class ProgException(Exception): pass
+
+class BadArgument(ProgException): pass
+
+class UnknownInputType(ProgException): pass
+
+class PipeFailed(ProgException): pass
+
+if __name__ == '__main__':
+    parser = optparse.OptionParser(usage='%prog [options] PATH',
+            description="""This tool is used to import the IPv4/6 BGP table dumps
+from routeviews.org into the RPKI Web Portal database.  If the
+input file is a bzip2 compressed file, it will be decompressed
+automatically.""")
+    parser.add_option('-t', '--type', dest='filetype', metavar='TYPE',
+            help='Specify the input file type (auto, text, mrt) [Default: %default]')
+    parser.add_option('-d', '--debug', dest='debug', action='store_true',
+            help='Enabling debugging output [Default: %default]')
+    parser.add_option('-v', '--verbose', dest='verbose', action='store_true',
+            help='Enable verbose output [Default: %default]')
+    parser.add_option('-u', '--bunzip2', dest='bunzip', metavar='PROG',
+            help='Specify bunzip2 program to use')
+    parser.add_option('-b', '--bgpdump', dest='bgpdump', metavar='PROG',
+            help='Specify path to bgdump binary')
+    parser.set_defaults(debug=False, verbose=False, filetype='auto')
+    options, args = parser.parse_args()
+
+    DEBUG = options.debug
+    VERBOSE = options.verbose
+    if options.bgpdump: BGPDUMP=os.path.expanduser(options.bgpdump)
+
+    try:
+        if len(args) != 1:
+            raise BadArgument('no filename specified, or more than one filename specified')
+        filename = args[0]
+
+        if options.filetype == 'auto':
+            # try to determine input type from filename, based on the default filenames from
+            # archive.routeviews.org
+            bname = os.path.basename(filename)
+            if bname.startswith('oix-full-snapshot-latest'):
+                filetype = 'text'
+            elif bname.startswith('rib.'):
+                filetype = 'mrt'
+            else:
+                raise UnknownInputType('unable to automatically determine input file type')
+            debug('auto-detected import format as "%s"' % filetype)
+        else:
+            filetype = options.filetype
+
+        pipe = None
+        if filename.endswith('.bz2'):
+            bunzip = 'bunzip2' if not options.bunzip else os.path.expanduser(options.bunzip)
+            debug('Decompressing input file on the fly...')
+            pipe = subprocess.Popen([bunzip, '--stdout', filename], stdout=subprocess.PIPE)
+            input_file = pipe.stdout
+        else:
+            input_file = open(filename)
+
+        try:
+            log('Reading data...')
+            dispatch = { 'text': parse_text, 'mrt': parse_mrt }
+            dispatch[filetype](input_file)
+        except KeyError:
+            raise UnknownInputType('"%s" is an unknown input file type' % filetype)
+
+        if pipe:
+            debug('Waiting for child to exit...')
+            pipe.wait()
+            if pipe.returncode:
+                raise PipeFailed('Child exited code %d' % pipe.returncode)
+            pipe = None
+        else:
+            input_file.close()
+
+        commit()
 
-print 'Updating timestamp metadata...'
-rpki.gui.app.timestamp.update('bgp_v4_import')
+        sys.exit(0)
 
-sys.exit(0)
+    except ProgException, e:
+        print 'Error:', e
+        sys.exit(1)
diff --git a/rpkid/rpki/gui/app/models.py b/rpkid/rpki/gui/app/models.py
index f81526e8..b3333986 100644
--- a/rpkid/rpki/gui/app/models.py
+++ b/rpkid/rpki/gui/app/models.py
@@ -261,7 +261,7 @@ class Timestamp(models.Model):
     set timestamps rather than updating this model directly."""
 
     name = models.CharField(max_length=30, primary_key=True)
-    ts = models.DateTimeField(null=False, default=0)
+    ts = models.DateTimeField(null=False)
 
     def __unicode__(self):
         return '%s: %s' % (self.name, self.ts)
diff --git a/rpkid/rpki/gui/app/timestamp.py b/rpkid/rpki/gui/app/timestamp.py
index 93f1d032..959f2025 100644
--- a/rpkid/rpki/gui/app/timestamp.py
+++ b/rpkid/rpki/gui/app/timestamp.py
@@ -19,6 +19,7 @@ from datetime import datetime
 
 def update(name):
     "Set the timestamp value for the given name to the current time."
-    obj, created = models.Timestamp.objects.get_or_create(name=name)
+    q = models.Timestamp.objects.filter(name=name)
+    obj = q[0] if q else models.Timestamp(name=name)
     obj.ts = datetime.utcnow()
     obj.save()