Wrap up database interface as a class.

svn path=/trunk/; revision=5450

1 files changed, 234 insertions, 220 deletions

diff --git a/scripts/rp-sqlite b/scripts/rp-sqlite
index aca57d88..56f0c440 100755
--- a/scripts/rp-sqlite
+++ b/scripts/rp-sqlite
@@ -34,12 +34,6 @@ import rpki.POW
 import rpki.x509
 import rpki.resource_set
 
-fn2map = dict(cer = rpki.x509.X509,
-              crl = rpki.x509.CRL,
-              mft = rpki.x509.SignedManifest,
-              roa = rpki.x509.ROA,
-              gbr = rpki.x509.Ghostbuster)
-
 sqlite3.register_adapter(rpki.POW.IPAddress,
                          lambda x: buffer("_" + x.toBytes()))
 
@@ -48,255 +42,275 @@ sqlite3.register_converter("RangeVal",
 
 
 def main():
-  db, cur = initialize_database()
-  test(cur)
-  db.close()
+  rpdb = RPDB()
+  test(rpdb)
+  rpdb.close()
+
 
+def test(rpdb):
+  fn2s = [None] + rpdb.fn2map.keys()
 
-def test(cur):
   print "Testing range functions"
-  for fn2 in [None] + fn2map.keys():
+  for fn2 in fn2s:
     if fn2 is not None:
       print
       print "Restricting search to type", fn2
     print
     print "Looking for range that should include adrilankha and psg again"
-    for r in find_by_range(cur, "147.28.0.19", "147.28.0.62", fn2):
+    for r in rpdb.find_by_range("147.28.0.19", "147.28.0.62", fn2):
       print r, r.uris
     print
     print "Looking for range that should include adrilankha"
-    for r in find_by_range(cur, "147.28.0.19", "147.28.0.19", fn2):
+    for r in rpdb.find_by_range("147.28.0.19", "147.28.0.19", fn2):
       print r, r.uris
     print
     print "Looking for range that should include ASN 3130"
-    for r in find_by_range(cur, 3130, 3130, fn2):
+    for r in rpdb.find_by_range(3130, 3130, fn2):
       print r, r.uris
   print
   print "Moving on to resource sets"
-  for expr in ("147.28.0.19-147.28.0.62",
-               "3130",
-               "2001:418:1::19/128",
-               "147.28.0.19-147.28.0.62,198.180.150.50/32",
-               "3130,147.28.0.19-147.28.0.62,198.180.150.50/32",
-               "2001:418:1::62/128,198.180.150.50/32,2001:418:8006::50/128",
-               "147.28.0.19-147.28.0.62,2001:418:1::19/128,2001:418:1::62/128,198.180.150.50/32,2001:418:8006::50/128"):
-    print
-    print "Trying", expr
-    for r in find_by_resource_bag(cur, rpki.resource_set.resource_bag.from_str(expr)):
-      print r, r.uris
-    for fn2 in fn2map:
+  for fn2 in fn2s:
+    if fn2 is not None:
       print
-      print "Trying", fn2, expr
-      for r in find_by_resource_bag(cur, rpki.resource_set.resource_bag.from_str(expr), fn2):
+      print "Restricting search to type", fn2
+    for expr in ("147.28.0.19-147.28.0.62",
+                 "3130",
+                 "2001:418:1::19/128",
+                 "147.28.0.19-147.28.0.62,198.180.150.50/32",
+                 "3130,147.28.0.19-147.28.0.62,198.180.150.50/32",
+                 "2001:418:1::62/128,198.180.150.50/32,2001:418:8006::50/128",
+                 "147.28.0.19-147.28.0.62,2001:418:1::19/128,2001:418:1::62/128,198.180.150.50/32,2001:418:8006::50/128"):
+      print
+      print "Trying", expr
+      for r in rpdb.find_by_resource_bag(rpki.resource_set.resource_bag.from_str(expr), fn2):
         print r, r.uris
     
 
-def find_by_ski(cur, ski, fn2 = None):
-  return find_results(cur, fn2,
-                      """
-                      SELECT id, fn2, der
-                      FROM object
-                      WHERE ski = ?
-                      """,
-                      buffer(ski))
-
-
-def find_by_aki(cur, aki, fn2 = None):
-  return find_results(cur, fn2,
-                      """
-                      SELECT id, fn2, der
-                      FROM object
-                      WHERE aki = ?
-                      """,
-                      buffer(aki))
-
-
-# It's easiest to understand overlap conditions by understanding
-# non-overlap then inverting and and applying De Morgan's law.  Ranges
-# A and B do not overlap if either A.min > B.max or A.max < B.min;
-# therefore they do overlap if A.min <= B.max and A.max >= B.min.
-
-def find_by_range(cur, range_min, range_max = None, fn2 = None):
-  if range_max is None:
-    range_max = range_min
-  if isinstance(range_min, (str, unicode)):
-    range_min = long(range_min) if range_min.isdigit() else rpki.POW.IPAddress(range_min)
-  if isinstance(range_max, (str, unicode)):
-    range_max = long(range_max) if range_max.isdigit() else rpki.POW.IPAddress(range_max)
-  assert isinstance(range_min, (int, long, rpki.POW.IPAddress))
-  assert isinstance(range_max, (int, long, rpki.POW.IPAddress))
-  return find_results(cur, fn2,
-                      """
-                      SELECT object.id, fn2, der
-                      FROM object, range
-                      WHERE ? <= max AND ? >= min AND object.id = range.id
-                      """,
-                      range_min,
-                      range_max)
-
-
-def find_by_resource_bag(cur, bag, fn2 = None):
-  assert bag.asn or bag.v4 or bag.v6
-  qset = []
-  aset = []
-  for rset in (bag.asn, bag.v4, bag.v6):
-    if rset:
-      for r in rset:
-        qset.append("(? <= max AND ? >= min)")
-        aset.append(r.min)
-        aset.append(r.max)
-  return find_results(*([
-    cur, fn2,
-    """
-    SELECT object.id, fn2, der
-    FROM object, range
-    WHERE object.id = range.id AND (%s)
-    """ % (" OR ".join(qset))
-    ] + aset))
-
-
-def find_results(cur, fn2, query, *args):
-  if fn2 is not None:
-    assert fn2 in fn2map
-    query += " AND fn2 = ?"
-    args = args + (fn2,)
-  query += " GROUP BY object.id"
-  results = []
-  cur.execute(query, args)
-  selections = cur.fetchall()
-  for rowid, fn2, der in selections:
-    obj = fn2map[fn2](DER = der)
-    cur.execute("SELECT uri FROM uri WHERE id = ?", (rowid,))
-    obj.uris = [u[0] for u in cur.fetchall()]
-    obj.uri = obj.uris[0] if len(obj.uris) == 1 else None
-    results.append(obj)
-  return results
-
-def initialize_database(db_name = "rp-sqlite.db",
-                        rcynic_root = os.path.expanduser("~/rpki/subvert-rpki.hactrn.net/trunk/"
-                                                         "rcynic/rcynic-data/unauthenticated"),
-                        delete_old_db = True,
-                        spinner = 100):
-
-  # For now just wire in the database name and rcynic root, fix this
-  # later if overall approach seems usable.  Might even end up just
-  # being an in-memory SQL database, who knows?
-
-  if delete_old_db:
-    try:
-      os.unlink(db_name)
-    except:
-      pass
-
-  db = sqlite3.connect(db_name, detect_types = sqlite3.PARSE_DECLTYPES)
-  db.text_factory = str
-
-  cur = db.cursor()
-  cur.executescript('''
-        PRAGMA foreign_keys = on;
-
-        CREATE TABLE object (
-                id INTEGER PRIMARY KEY NOT NULL,
-                der BLOB NOT NULL,
-                fn2 TEXT NOT NULL,
-                ski BLOB,
-                aki BLOB,
-                inherits BOOLEAN NOT NULL,
-                UNIQUE (der));
-
-        CREATE TABLE uri (
-                id INTEGER NOT NULL,
-                uri TEXT NOT NULL,
-                UNIQUE (uri),
-                FOREIGN KEY (id) REFERENCES object(id)
-                        ON DELETE CASCADE
-                        ON UPDATE CASCADE);
-
-        CREATE INDEX uri_index ON uri(id);
-
-        CREATE TABLE range (
-                id INTEGER NOT NULL,
-                min RangeVal NOT NULL,
-                max RangeVal NOT NULL,
-                UNIQUE (id, min, max),
-                FOREIGN KEY (id) REFERENCES object(id)
-                        ON DELETE CASCADE
-                        ON UPDATE CASCADE);
-
-        CREATE INDEX range_index ON range(min, max);
-        ''')
-
-  nobj = 0
-
-  for root, dirs, files in os.walk(rcynic_root):
-    for fn in files:
-      fn = os.path.join(root, fn)
-      fn2 = os.path.splitext(fn)[1][1:]
-
+class RPDB(object):
+  """
+  Relying party database.
+  """
+
+  fn2map = dict(cer = rpki.x509.X509,
+                crl = rpki.x509.CRL,
+                mft = rpki.x509.SignedManifest,
+                roa = rpki.x509.ROA,
+                gbr = rpki.x509.Ghostbuster)
+
+  def __init__(self,
+               db_name = "rp-sqlite.db",
+               rcynic_root = os.path.expanduser("~/rpki/subvert-rpki.hactrn.net/trunk/"
+                                                "rcynic/rcynic-data/unauthenticated"),
+               delete_old_db = True,
+               spinner = 100):
+
+    # For now just wire in the database name and rcynic root, fix this
+    # later if overall approach seems usable.  Might even end up just
+    # being an in-memory SQL database, who knows?
+
+    if delete_old_db:
       try:
-        obj = fn2map[fn2](DER_file = fn)
+        os.unlink(db_name)
       except:
-        continue
-
-      if spinner and nobj % spinner == 0:
-        sys.stderr.write("\r%s %d..." % ("|\\-/"[(nobj/spinner) & 3], nobj))
+        pass
+
+    self.db = sqlite3.connect(db_name, detect_types = sqlite3.PARSE_DECLTYPES)
+    self.db.text_factory = str
+
+    self.cur = self.db.cursor()
+    self.cur.executescript('''
+          PRAGMA foreign_keys = on;
+
+          CREATE TABLE object (
+                  id INTEGER PRIMARY KEY NOT NULL,
+                  der BLOB NOT NULL,
+                  fn2 TEXT NOT NULL,
+                  ski BLOB,
+                  aki BLOB,
+                  inherits BOOLEAN NOT NULL,
+                  UNIQUE (der));
+
+          CREATE TABLE uri (
+                  id INTEGER NOT NULL,
+                  uri TEXT NOT NULL,
+                  UNIQUE (uri),
+                  FOREIGN KEY (id) REFERENCES object(id)
+                          ON DELETE CASCADE
+                          ON UPDATE CASCADE);
+
+          CREATE INDEX uri_index ON uri(id);
+
+          CREATE TABLE range (
+                  id INTEGER NOT NULL,
+                  min RangeVal NOT NULL,
+                  max RangeVal NOT NULL,
+                  UNIQUE (id, min, max),
+                  FOREIGN KEY (id) REFERENCES object(id)
+                          ON DELETE CASCADE
+                          ON UPDATE CASCADE);
+
+          CREATE INDEX range_index ON range(min, max);
+          ''')
+
+    nobj = 0
+
+    for root, dirs, files in os.walk(rcynic_root):
+      for fn in files:
+        fn = os.path.join(root, fn)
+        fn2 = os.path.splitext(fn)[1][1:]
 
-      nobj += 1
-
-      inherits = False
-
-      if fn2 == "crl":
-        ski = None
-        aki = buffer(obj.get_AKI())
-        cer = None
-        bag = None
-
-      else:
-        if fn2 == "cer":
-          cer = obj
-        else:
-          cer = rpki.x509.X509(POW = obj.get_POW().certs()[0])
-        ski = buffer(cer.get_SKI())
         try:
-          aki = buffer(cer.get_AKI())
+          obj = self.fn2map[fn2](DER_file = fn)
         except:
-          aki = None
-        bag = cer.get_3779resources()
-        inherits = bag.asn.inherit or bag.v4.inherit or bag.v6.inherit
+          continue
 
-      der = buffer(obj.get_DER())
-      uri = "rsync://" + fn[len(rcynic_root) + 1:]
+        if spinner and nobj % spinner == 0:
+          sys.stderr.write("\r%s %d..." % ("|\\-/"[(nobj/spinner) & 3], nobj))
 
-      try:
-        cur.execute("INSERT INTO object (der, fn2, ski, aki, inherits) VALUES (?, ?, ?, ?, ?)",
-                    (der, fn2, ski, aki, inherits))
-        rowid = cur.lastrowid
+        nobj += 1
 
-      except sqlite3.IntegrityError:
-        cur.execute("SELECT id FROM object WHERE der = ? AND fn2 = ?", (der, fn2))
-        rows = cur.fetchall()
-        rowid = rows[0][0]
-        assert len(rows) == 1
+        inherits = False
 
-      else:
-        if bag is not None:
-          for rset in (bag.asn, bag.v4, bag.v6):
-            if rset is not None:
-              cur.executemany("REPLACE INTO range (id, min, max) VALUES (?, ?, ?)",
-                              ((rowid, i.min, i.max) for i in rset))
+        if fn2 == "crl":
+          ski = None
+          aki = buffer(obj.get_AKI())
+          cer = None
+          bag = None
 
-      cur.execute("INSERT INTO uri (id, uri) VALUES (?, ?)",
-                  (rowid, uri))
+        else:
+          if fn2 == "cer":
+            cer = obj
+          else:
+            cer = rpki.x509.X509(POW = obj.get_POW().certs()[0])
+          ski = buffer(cer.get_SKI())
+          try:
+            aki = buffer(cer.get_AKI())
+          except:
+            aki = None
+          bag = cer.get_3779resources()
+          inherits = bag.asn.inherit or bag.v4.inherit or bag.v6.inherit
+
+        der = buffer(obj.get_DER())
+        uri = "rsync://" + fn[len(rcynic_root) + 1:]
 
-  if spinner:
-    sys.stderr.write("\r= %d objects, committing..." % nobj)
+        try:
+          self.cur.execute("INSERT INTO object (der, fn2, ski, aki, inherits) VALUES (?, ?, ?, ?, ?)",
+                           (der, fn2, ski, aki, inherits))
+          rowid = self.cur.lastrowid
 
-  db.commit()
+        except sqlite3.IntegrityError:
+          self.cur.execute("SELECT id FROM object WHERE der = ? AND fn2 = ?", (der, fn2))
+          rows = self.cur.fetchall()
+          rowid = rows[0][0]
+          assert len(rows) == 1
 
-  if spinner:
-    sys.stderr.write("done.\n")
+        else:
+          if bag is not None:
+            for rset in (bag.asn, bag.v4, bag.v6):
+              if rset is not None:
+                self.cur.executemany("REPLACE INTO range (id, min, max) VALUES (?, ?, ?)",
+                                     ((rowid, i.min, i.max) for i in rset))
+
+        self.cur.execute("INSERT INTO uri (id, uri) VALUES (?, ?)",
+                         (rowid, uri))
+
+    if spinner:
+      sys.stderr.write("\r= %d objects, committing..." % nobj)
+
+    self.db.commit()
+
+    if spinner:
+      sys.stderr.write("done.\n")
+
+
+  def find_by_ski(self, ski, fn2 = None):
+    return self._find_results(
+      fn2,
+      """
+      SELECT id, fn2, der
+      FROM object
+      WHERE ski = ?
+      """,
+      [buffer(ski)])
+
+
+  def find_by_aki(self, aki, fn2 = None):
+    return self._find_results(
+      fn2,
+      """
+      SELECT id, fn2, der
+      FROM object
+      WHERE aki = ?
+      """,
+      [buffer(aki)])
+
+
+  # It's easiest to understand overlap conditions by understanding
+  # non-overlap then inverting and and applying De Morgan's law.  Ranges
+  # A and B do not overlap if either A.min > B.max or A.max < B.min;
+  # therefore they do overlap if A.min <= B.max and A.max >= B.min.
+
+  def find_by_range(self, range_min, range_max = None, fn2 = None):
+    if range_max is None:
+      range_max = range_min
+    if isinstance(range_min, (str, unicode)):
+      range_min = long(range_min) if range_min.isdigit() else rpki.POW.IPAddress(range_min)
+    if isinstance(range_max, (str, unicode)):
+      range_max = long(range_max) if range_max.isdigit() else rpki.POW.IPAddress(range_max)
+    assert isinstance(range_min, (int, long, rpki.POW.IPAddress))
+    assert isinstance(range_max, (int, long, rpki.POW.IPAddress))
+    return self._find_results(
+      fn2,
+      """
+      SELECT object.id, fn2, der
+      FROM object, range
+      WHERE ? <= max AND ? >= min AND object.id = range.id
+      """,
+      [range_min, range_max])
+
+
+  def find_by_resource_bag(self, bag, fn2 = None):
+    assert bag.asn or bag.v4 or bag.v6
+    qset = []
+    aset = []
+    for rset in (bag.asn, bag.v4, bag.v6):
+      if rset:
+        for r in rset:
+          qset.append("(? <= max AND ? >= min)")
+          aset.append(r.min)
+          aset.append(r.max)
+    return self._find_results(
+      fn2,
+      """
+      SELECT object.id, fn2, der
+      FROM object, range
+      WHERE object.id = range.id AND (%s)
+      """ % (" OR ".join(qset)),
+      aset)
+
+
+  def _find_results(self, fn2, query, args):
+    if fn2 is not None:
+      assert fn2 in self.fn2map
+      query += " AND fn2 = ?"
+      args.append(fn2)
+    query += " GROUP BY object.id"
+    results = []
+    self.cur.execute(query, args)
+    selections = self.cur.fetchall()
+    for rowid, fn2, der in selections:
+      obj = self.fn2map[fn2](DER = der)
+      self.cur.execute("SELECT uri FROM uri WHERE id = ?", (rowid,))
+      obj.uris = [u[0] for u in self.cur.fetchall()]
+      obj.uri = obj.uris[0] if len(obj.uris) == 1 else None
+      results.append(obj)
+    return results
+
+
+  def close(self):
+    self.cur.close()
+    self.db.close()
 
-  return db, cur
 
 if __name__ == "__main__":
   main()