rwhoisd/MemDB.py

   1 # This file is part of python-rwhoisd
   2 #
   3 # Copyright (C) 2003, David E. Blacka
   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, but
  11 # WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  13 # 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
  18 # USA
  19
  20 import bisect, types
  21 import MemIndex, Cidr
  22 from Rwhois import rwhoisobject
  23
  24 class MemDB:
  25
  26     def __init__(self):
  27
  28         # a dictonary holding the various attribute indexes.  The keys
  29         # are lowercase attribute names, values are MemIndex or
  30         # CidrMemIndex objects.
  31         self.indexes = {}
  32
  33         # a dictonary holding the actual rwhoisobjects.  keys are
  34         # string IDs, values are rwhoisobject instances.
  35         self.main_index = {}
  36
  37         # dictionary holding all of the seen attributes.  keys are
  38         # lowercase attribute names, value is a character indicating
  39         # the index type (if indexed), or None if not indexed.  Index
  40         # type characters a 'N' for normal string index, 'C' for CIDR
  41         # index.
  42         self.attrs = {}
  43
  44         # Lists containing attribute names that have indexes by type.
  45         # This exists so unconstrained searches can just iterate over
  46         # them.
  47         self.normal_indexes = []
  48         self.cidr_indexes   = []
  49
  50         # dictonary holding all of the seen class names.  keys are
  51         # lowercase classnames, value is always None.
  52         self.classes = {}
  53
  54         # dictionary holding all of the seen auth-areas.  keys are
  55         # lowercase authority area names, value is always None.
  56         self.authareas = {}
  57
  58     def init_schema(self, schema_file):
  59         """Initialize the schema from a schema file.  Currently the
  60         schema file is a list of 'attribute_name = index_type' pairs,
  61         one per line.  index_type is one of N or C, where N means a
  62         normal string index, and C means a CIDR index.
  63
  64         It should be noted that this database implementation
  65         implements a global namespace for attributes, which isn't
  66         really correct according to RFC 2167.  RFC 2167 dictates that
  67         different authority area are actually autonomous and thus have
  68         separate schemas."""
  69
  70         # initialize base schema
  71
  72         self.attrs['id']         = "N"
  73         self.attrs['auth-area']  = None
  74         self.attrs['class-name'] = None
  75         self.attrs['updated']    = None
  76         self.attrs['referred-auth-area'] = "R"
  77
  78         sf = open(schema_file, "r")
  79
  80         for line in sf.xreadlines():
  81             line = line.strip()
  82             if not line or line.startswith("#"): continue
  83
  84             attr, it = line.split("=")
  85             self.attrs[attr.strip().lower()] = it.strip()[0].upper()
  86
  87         for attr, index_type in self.attrs.items():
  88             if index_type == "N":
  89                 # normal index
  90                 self.indexes[attr] = MemIndex.MemIndex()
  91                 self.normal_indexes.append(attr)
  92             elif index_type == "A":
  93                 # "all" index -- both a normal and a cidr index
  94                 self.indexes[attr] = MemIndex.ComboMemIndex()
  95                 self.normal_indexes.append(attr)
  96                 self.cidr_indexes.append(attr)
  97             elif index_type == "R":
  98                 # referral index, an all index that must be searched
  99                 # explictly by attribute
 100                 self.indexes[attr] = MemIndex.ComboMemIndex()
 101             elif index_type == "C":
 102                 # a cidr index
 103                 self.indexes[attr] = MemIndex.CidrMemIndex()
 104                 self.cidr_indexes.append(attr)
 105         return
 106
 107     def add_object(self, obj):
 108         """Add an rwhoisobject to the raw indexes, including the
 109         master index."""
 110
 111         # add the object to the main index
 112         id = obj.getid()
 113         if not id: return
 114         id = id.lower()
 115
 116         self.main_index[id] = obj
 117
 118         for a,v in obj.items():
 119             # note the attribute.
 120             index_type = self.attrs.setdefault(a, None)
 121             v = v.lower()
 122             # make sure that we note the auth-area and class
 123             if a == 'auth-area':
 124                 self.authareas.setdefault(v, None)
 125             elif a == 'class-name':
 126                 self.classes.setdefault(v, None)
 127
 128             if index_type:
 129                 index = self.indexes[a]
 130                 index.add(v, id)
 131
 132     def load_data(self, data_file):
 133         """Load data from rwhoisd-style TXT files (i.e., attr:value,
 134         records separated with a "---" bare line)."""
 135
 136         df = open(data_file, "r")
 137         obj = rwhoisobject()
 138
 139         for line in df.xreadlines():
 140             line = line.strip()
 141             if line.startswith("#"): continue
 142             if not line or line.startswith("---"):
 143                 # we've reached the end of an object, so index it.
 144                 self.add_object(obj)
 145                 # reset obj
 146                 obj = rwhoisobject()
 147                 continue
 148
 149             a, v = line.split(":", 1)
 150             obj.add_attr(a, v.lstrip())
 151
 152         self.add_object(obj)
 153         return
 154
 155     def index_data(self):
 156         """Prepare the indexes for searching.  Currently, this isn't
 157         strictly necessary (the indexes will prepare themselves when
 158         necessary), but it should eliminate a penalty on initial
 159         searches"""
 160
 161         for i in self.indexes.values():
 162             i.prepare()
 163         return
 164
 165     def is_attribute(self, attr):
 166         return self.attrs.has_key(attr.lower())
 167
 168     def is_indexed_attr(self, attr):
 169         if self.is_attribute(attr):
 170             return self.attrs[attr.lower()]
 171         return False
 172
 173     def is_objectclass(self, objectclass):
 174         return self.classes.has_key(objectclass.lower())
 175
 176     def is_autharea(self, aa):
 177         return self.authareas.has_key(aa.lower())
 178
 179     def get_authareas(self):
 180         return self.authareas.keys()
 181
 182     def fetch_objects(self, id_list):
 183         return [ self.main_index[x] for x in id_list
 184                  if self.main_index.has_key(x) ]
 185
 186     def search_attr(self, attr, value, max = 0):
 187
 188         """Search for a value in a particular attribute's index.  If
 189         the attribute is cidr indexed, an attempt to convert value
 190         into a Cidr object will be made.  Returns a list of object ids
 191         (or an empty list if nothing was found)"""
 192
 193         attr = attr.lower()
 194         index_type = self.attrs.get(attr)
 195         index = self.indexes.get(attr)
 196         if not index: return []
 197
 198         super_prefix_match = False
 199         if value.endswith("**"):
 200             super_prefix_match = True
 201
 202         prefix_match = False
 203         if value.endswith("*"):
 204             value = value.rstrip("*")
 205             prefix_match = True
 206
 207         if index_type == 'C' and not isinstance(value, Cidr.Cidr):
 208             value = Cidr.valid_cidr(value)
 209         else:
 210             value = value.strip().lower()
 211
 212         if index_type == 'C' and super_prefix_match:
 213             return index.find_subnets(value, max)
 214
 215         res = index.find(value, prefix_match, max)
 216         return IndexResult(res)
 217
 218     def search_normal(self, value, max = 0):
 219         """Search for a value in the 'normal' (string keyed) indexes.
 220         Returns a list of object ids, or an empty list if nothing was
 221         found."""
 222
 223         res = IndexResult()
 224
 225         for attr in self.normal_indexes:
 226             res.extend(self.search_attr(attr, value, max))
 227             if max:
 228                 if len(res) >= max:
 229                     res.truncate(max)
 230                     return res
 231         return res
 232
 233     def search_cidr(self, value, max = 0):
 234         """Search for a value in the cidr indexes.  Returns a list of
 235         object ids, or an empty list if nothing was found."""
 236
 237         res = IndexResult()
 238         for attr in self.cidr_indexes:
 239             res.extend(self.search_attr(attr, value, max))
 240             if max:
 241                 if len(res) >= max:
 242                     res.truncate(max)
 243                     return res
 244         return res
 245
 246     def search_referral(self, value, max = 0):
 247         """Given a heirarchal value, search for referrals.  Returns a
 248         list of object ids or an empty list."""
 249
 250         return self.search_attr("referred-auth-area", value, max)
 251
 252     def object_iterator(self):
 253         return self.main_index.itervalues()
 254
 255 class IndexResult:
 256     def __init__(self, list=None):
 257         if not list: list = []
 258         self.data = list
 259         self._dict = dict(zip(self.data, self.data))
 260
 261     def __len__(self):
 262         return len(self.data)
 263
 264     def extend(self, list):
 265         if isinstance(list, type(self)):
 266             list = list.list()
 267         new_els = [ x for x in list if not self._dict.has_key(x) ]
 268         self.data.extend(new_els)
 269         self._dict.update(dict(zip(new_els, new_els)))
 270
 271     def list(self):
 272         return self.data
 273
 274     def truncate(self, n=0):
 275         to_del = self.data[n:]
 276         for i in to_del: del self._dict[i]
 277         self.data = self.data[:n]
 278
 279
 280 # test driver
 281 if __name__ == "__main__":
 282     import sys
 283     db = MemDB()
 284
 285     print "loading schema:", sys.argv[1]
 286     db.init_schema(sys.argv[1])
 287     for data_file in sys.argv[2:]:
 288         print "loading data file:", data_file
 289         db.load_data(data_file)
 290     db.index_data()
 291
 292     print "Schema: authority areas"
 293     for a in db.authareas.keys():
 294         print "   %s" % a
 295     print "Schema: classes"
 296     for c in db.classes.keys():
 297         print "   %s" % c
 298     print "Schema: attributes"
 299     for a in db.attrs.keys():
 300         print "   %s" % a
 301
 302     print "Is 'Network' a class?", db.is_objectclass("Network")
 303
 304 #    for k, v in db.main_index.items():
 305 #        print "main_index[", k, "]:", v
 306
 307     print "searching for a.com"
 308     res = db.search_attr("domain-name", "a.com")
 309     print res.list()
 310     print [ str(x) for x in db.fetch_objects(res.list()) ]
 311
 312     print "searching for doe"
 313     res = db.search_normal("doe")
 314     print res.list()
 315     print [ str(x) for x in db.fetch_objects(res.list()) ]
 316
 317     print "searching for 10.0.0.2"
 318     res = db.search_cidr("10.0.0.2")
 319     print res.list()
 320     print [ str(x) for x in db.fetch_objects(res.list()) ]
 321
 322     print "searching for fddi.a.com"
 323     res = db.search_normal("fddi.a.com")
 324     print res.list()
 325
 326     print "searching referral index for fddi.a.com"
 327     res = db.search_attr("referred-auth-area", "fddi.a.com")
 328     print res.list()
 329     print [ str(x) for x in db.fetch_objects(res.list()) ]
 330
 331