lib/bridgedb/Storage.py

   1 # BridgeDB by Nick Mathewson.
   2 # Copyright (c) 2007-2009, The Tor Project, Inc.
   3 # See LICENSE for licensing information
   4
   5 import calendar
   6 import os
   7 import logging
   8 import bridgedb.Bridges
   9 import binascii
  10 import sqlite3
  11 import time
  12
  13 toHex = binascii.b2a_hex
  14 fromHex = binascii.a2b_hex
  15 HEX_ID_LEN = 40
  16
  17 def _escapeValue(v):
  18     return "'%s'" % v.replace("'", "''")
  19
  20 def timeToStr(t):
  21     return time.strftime("%Y-%m-%d %H:%M", time.gmtime(t))
  22 def strToTime(t):
  23     return calendar.timegm(time.strptime(t, "%Y-%m-%d %H:%M"))
  24
  25 class SqliteDict:
  26     """
  27        A SqliteDict wraps a SQLite table and makes it look like a
  28        Python dictionary.  In addition to the single key and value
  29        columns, there can be a number of "fixed" columns, such that
  30        the dictionary only contains elements of the table where the
  31        fixed columns are set appropriately.
  32     """
  33     def __init__(self, conn, cursor, table, fixedcolnames, fixedcolvalues,
  34                  keycol, valcol):
  35         assert len(fixedcolnames) == len(fixedcolvalues)
  36         self._conn = conn
  37         self._cursor = cursor
  38         keys = ", ".join(fixedcolnames+(keycol,valcol))
  39         vals = "".join("%s, "%_escapeValue(v) for v in fixedcolvalues)
  40         constraint = "WHERE %s = ?"%keycol
  41         if fixedcolnames:
  42             constraint += "".join(
  43                 " AND %s = %s"%(c,_escapeValue(v))
  44                 for c,v in zip(fixedcolnames, fixedcolvalues))
  45
  46         self._getStmt = "SELECT %s FROM %s %s"%(valcol,table,constraint)
  47         self._delStmt = "DELETE FROM %s %s"%(table,constraint)
  48         self._setStmt = "INSERT OR REPLACE INTO %s (%s) VALUES (%s?, ?)"%(
  49             table, keys, vals)
  50
  51         constraint = " AND ".join("%s = %s"%(c,_escapeValue(v))
  52                 for c,v in zip(fixedcolnames, fixedcolvalues))
  53         if constraint:
  54             whereClause = " WHERE %s"%constraint
  55         else:
  56             whereClause = ""
  57
  58         self._keysStmt = "SELECT %s FROM %s%s"%(keycol,table,whereClause)
  59
  60     def __setitem__(self, k, v):
  61         self._cursor.execute(self._setStmt, (k,v))
  62     def __delitem__(self, k):
  63         self._cursor.execute(self._delStmt, (k,))
  64         if self._cursor.rowcount == 0:
  65             raise KeyError(k)
  66     def __getitem__(self, k):
  67         self._cursor.execute(self._getStmt, (k,))
  68         val = self._cursor.fetchone()
  69         if val == None:
  70             raise KeyError(k)
  71         else:
  72             return val[0]
  73     def has_key(self, k):
  74         self._cursor.execute(self._getStmt, (k,))
  75         return self._cursor.rowcount != 0
  76     def get(self, k, v=None):
  77         self._cursor.execute(self._getStmt, (k,))
  78         val = self._cursor.fetchone()
  79         if val == None:
  80             return v;
  81         else:
  82             return val[0]
  83     def setdefault(self, k, v):
  84         try:
  85             r = self[k]
  86         except KeyError:
  87             r = self[k] = v
  88         return r
  89     def keys(self):
  90         self._cursor.execute(self._keysStmt)
  91         return [ key for (key,) in self._cursor.fetchall() ]
  92
  93     def commit(self):
  94         self._conn.commit()
  95     def rollback(self):
  96         self._conn.rollback()
  97
  98
  99 #  The old DB system was just a key->value mapping DB, with special key
 100 #  prefixes to indicate which database they fell into.
 101 #
 102 #     sp|<ID> -- given to bridgesplitter; maps bridgeID to ring name.
 103 #     em|<emailaddr> -- given to emailbaseddistributor; maps email address
 104 #            to concatenated ID.
 105 #     fs|<ID> -- Given to BridgeTracker, maps to time when a router was
 106 #            first seen (YYYY-MM-DD HH:MM)
 107 #     ls|<ID> -- given to bridgetracker, maps to time when a router was
 108 #            last seen (YYYY-MM-DD HH:MM)
 109 #
 110 # We no longer want to use em| at all, since we're not doing that kind
 111 # of persistence any more.
 112
 113 # Here is the SQL schema.
 114
 115 SCHEMA1_SCRIPT = """
 116  CREATE TABLE Config (
 117      key PRIMARY KEY NOT NULL,
 118      value
 119  );
 120
 121  CREATE TABLE Bridges (
 122      id INTEGER PRIMARY KEY NOT NULL,
 123      hex_key,
 124      address,
 125      or_port,
 126      distributor,
 127      first_seen,
 128      last_seen
 129  );
 130
 131  CREATE UNIQUE INDEX BridgesKeyIndex ON Bridges ( hex_key );
 132
 133  CREATE TABLE EmailedBridges (
 134      email PRIMARY KEY NOT NULL,
 135      when_mailed
 136  );
 137
 138  CREATE INDEX EmailedBridgesWhenMailed on EmailedBridges ( email );
 139
 140  INSERT INTO Config VALUES ( 'schema-version', 1 );
 141 """
 142
 143
 144 class Database:
 145     def __init__(self, sqlite_fname, db_fname=None):
 146         if db_fname is None:
 147             self._conn = openDatabase(sqlite_fname)
 148         else:
 149             self._conn = openOrConvertDatabase(sqlite_fname, db_fname)
 150         self._cur = self._conn.cursor()
 151
 152     def commit(self):
 153         self._conn.commit()
 154
 155     def close(self):
 156         self._cur.close()
 157         self._conn.close()
 158
 159     def insertBridgeAndGetRing(self, bridge, setRing, seenAt):
 160         '''updates info about bridge, setting ring to setRing if none was set.
 161            Returns the bridge's ring.
 162         '''
 163         cur = self._cur
 164
 165         t = timeToStr(seenAt)
 166         h = bridge.fingerprint
 167         assert len(h) == HEX_ID_LEN
 168
 169         cur.execute("SELECT id, distributor "
 170                     "FROM Bridges WHERE hex_key = ?", (h,))
 171         v = cur.fetchone()
 172         if v is not None:
 173             idx, ring = v
 174             # Update last_seen and address.
 175             cur.execute("UPDATE Bridges SET address = ?, or_port = ?, "
 176                         "last_seen = ? WHERE id = ?",
 177                         (bridge.ip, bridge.orport, timeToStr(seenAt), idx))
 178             return ring
 179         else:
 180             # Insert it.
 181             cur.execute("INSERT INTO Bridges (hex_key, address, or_port, "
 182                         "distributor, first_seen, last_seen) "
 183                         "VALUES (?, ?, ?, ?, ?, ?)",
 184                         (h, bridge.ip, bridge.orport, setRing, t, t))
 185             return setRing
 186
 187     def cleanEmailedBridges(self, expireBefore):
 188         cur = self._cur
 189         t = timeToStr(expireBefore)
 190
 191         cur.execute("DELETE FROM EmailedBridges WHERE when_mailed < ?", (t,));
 192
 193     def getEmailTime(self, addr):
 194         cur = self._cur
 195         cur.execute("SELECT when_mailed FROM EmailedBridges WHERE "
 196                     "email = ?", (addr,))
 197         v = cur.fetchone()
 198         if v is None:
 199             return None
 200         return strToTime(v[0])
 201
 202     def setEmailTime(self, addr, whenMailed):
 203         cur = self._cur
 204         t = timeToStr(whenMailed)
 205         cur.execute("INSERT OR REPLACE INTO EmailedBridges "
 206                     "(email,when_mailed) VALUES (?,?)", (addr, t))
 207
 208 def openDatabase(sqlite_file):
 209     conn = sqlite3.Connection(sqlite_file)
 210     cur = conn.cursor()
 211     try:
 212         try:
 213             cur.execute("SELECT value FROM Config WHERE key = 'schema-version'")
 214             val, = cur.fetchone()
 215             if val != 1:
 216                 logging.warn("Unknown schema version %s in database.", val)
 217         except sqlite3.OperationalError:
 218             logging.warn("No Config table found in DB; creating tables")
 219             cur.executescript(SCHEMA1_SCRIPT)
 220             conn.commit()
 221     finally:
 222         cur.close()
 223     return conn
 224
 225
 226 def openOrConvertDatabase(sqlite_file, db_file):
 227     """Open a sqlite database, converting it from a db file if needed."""
 228     if os.path.exists(sqlite_file):
 229         return openDatabase(sqlite_file)
 230
 231     conn = sqlite3.Connection(sqlite_file)
 232     cur = conn.cursor()
 233     cur.executescript(SCHEMA1_SCRIPT)
 234     conn.commit()
 235
 236     import anydbm
 237
 238     try:
 239         db = anydbm.open(db_file, 'r')
 240     except anydbm.error:
 241         return conn
 242
 243     try:
 244         # We handle all the sp| keys first, since other tables have
 245         # dependencies on Bridges.
 246         for k in db.keys():
 247             v = db[k]
 248             if k.startswith("sp|"):
 249                 assert len(k) == 23
 250                 cur.execute("INSERT INTO Bridges ( hex_key, distributor ) "
 251                             "VALUES (?, ?)", (toHex(k[3:]),v))
 252         # Now we handle the other key types.
 253         for k in db.keys():
 254             v = db[k]
 255             if k.startswith("fs|"):
 256                 assert len(k) == 23
 257                 cur.execute("UPDATE Bridges SET first_seen = ? "
 258                             "WHERE hex_key = ?", (v, toHex(k[3:])))
 259             elif k.startswith("ls|"):
 260                 assert len(k) == 23
 261                 cur.execute("UPDATE Bridges SET last_seen = ? "
 262                             "WHERE hex_key = ?", (v, toHex(k[3:])))
 263             #elif k.startswith("em|"):
 264             #    keys = list(toHex(i) for i in
 265             #        bridgedb.Bridges.chopString(v, bridgedb.Bridges.ID_LEN))
 266             #    cur.executemany("INSERT INTO EmailedBridges ( email, id ) "
 267             #                    "SELECT ?, id FROM Bridges WHERE hex_key = ?",
 268             #                    [(k[3:],i) for i in keys])
 269             elif k.startswith("sp|") or k.startswith("em|"):
 270                 pass
 271             else:
 272                 logging.warn("Unrecognized key %r", k)
 273     except:
 274         conn.rollback()
 275         conn.close()
 276         os.unlink(sqlite_file)
 277         raise
 278
 279     conn.commit()
 280     return conn
 281
 282 _THE_DB = None
 283
 284 def setGlobalDB(db):
 285     global _THE_DB
 286     _THE_DB = db
 287
 288 def getDB(db):
 289     return _THE_DB
 290