chiark - git - mdw - chopwood/blob - util.py

   1 ### -*-python-*-
   2 ###
   3 ### Miscellaneous utilities
   4 ###
   5 ### (c) 2013 Mark Wooding
   6 ###
   7
   8 ###----- Licensing notice ---------------------------------------------------
   9 ###
  10 ### This file is part of Chopwood: a password-changing service.
  11 ###
  12 ### Chopwood is free software; you can redistribute it and/or modify
  13 ### it under the terms of the GNU Affero General Public License as
  14 ### published by the Free Software Foundation; either version 3 of the
  15 ### License, or (at your option) any later version.
  16 ###
  17 ### Chopwood is distributed in the hope that it will be useful,
  18 ### but WITHOUT ANY WARRANTY; without even the implied warranty of
  19 ### MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  20 ### GNU Affero General Public License for more details.
  21 ###
  22 ### You should have received a copy of the GNU Affero General Public
  23 ### License along with Chopwood; if not, see
  24 ### <http://www.gnu.org/licenses/>.
  25
  26 from __future__ import with_statement
  27
  28 import base64 as BN
  29 import contextlib as CTX
  30 import fcntl as F
  31 import os as OS
  32 import re as RX
  33 import signal as SIG
  34 import sys as SYS
  35 import time as T
  36
  37 try: import threading as TH
  38 except ImportError: import dummy_threading as TH
  39
  40 ###--------------------------------------------------------------------------
  41 ### Some basics.
  42
  43 def identity(x):
  44   """The identity function: returns its argument."""
  45   return x
  46
  47 def constantly(x):
  48   """The function which always returns X."""
  49   return lambda: x
  50
  51 class struct (object):
  52   """A simple object for storing data in attributes."""
  53   DEFAULTS = {}
  54   def __init__(me, *args, **kw):
  55     cls = me.__class__
  56     for k, v in kw.iteritems(): setattr(me, k, v)
  57     try:
  58       slots = cls.__slots__
  59     except AttributeError:
  60       if args: raise ValueError, 'no slots defined'
  61     else:
  62       if len(args) > len(slots): raise ValueError, 'too many arguments'
  63       for k, v in zip(slots, args): setattr(me, k, v)
  64       for k in slots:
  65         if hasattr(me, k): continue
  66         try: setattr(me, k, cls.DEFAULTS[k])
  67         except KeyError: raise ValueError, "no value for `%s'" % k
  68
  69 class Tag (object):
  70   """An object whose only purpose is to be distinct from other objects."""
  71   def __init__(me, name):
  72     me._name = name
  73   def __repr__(me):
  74     return '#<%s %r>' % (type(me).__name__, me._name)
  75
  76 class DictExpanderClass (type):
  77   """
  78   Metaclass for classes with autogenerated members.
  79
  80   If the class body defines a dictionary `__extra__' then the key/value pairs
  81   in this dictionary are promoted into attributes of the class.  This is much
  82   easier -- and safer -- than fiddling about with `locals'.
  83   """
  84   def __new__(cls, name, supers, dict):
  85     try:
  86       ex = dict['__extra__']
  87     except KeyError:
  88       pass
  89     else:
  90       for k, v in ex.iteritems():
  91         dict[k] = v
  92       del dict['__extra__']
  93     return super(DictExpanderClass, cls).__new__(cls, name, supers, dict)
  94
  95 class ExpectedError (Exception):
  96   """
  97   A (concrete) base class for various errors we expect to encounter.
  98
  99   The `msg' attribute carries a human-readable message explaining what the
 100   problem actually is.  The really important bit, though, is the `code'
 101   attribute, which carries an HTTP status code to be reported to the user
 102   agent, if we're running through CGI.
 103   """
 104   def __init__(me, code, msg):
 105     me.code = code
 106     me.msg = msg
 107   def __str__(me):
 108     return '%s (%d)' % (me.msg, me.code)
 109
 110 def register(dict, name):
 111   """A decorator: add the decorated function to DICT, under the key NAME."""
 112   def _(func):
 113     dict[name] = func
 114     return func
 115   return _
 116
 117 class StringSubst (object):
 118   """
 119   A string substitution.  Initialize with a dictionary mapping source strings
 120   to target strings.  The object is callable, and maps strings in the obvious
 121   way.
 122   """
 123   def __init__(me, map):
 124     me._map = map
 125     me._rx = RX.compile('|'.join(RX.escape(s) for s in map))
 126   def __call__(me, s):
 127     return me._rx.sub(lambda m: me._map[m.group(0)], s)
 128
 129 def readline(what, file = SYS.stdin):
 130   """Read a single line from FILE (default stdin) and return it."""
 131   try: line = SYS.stdin.readline()
 132   except IOError, e: raise ExpectedError, (500, str(e))
 133   if not line.endswith('\n'):
 134     raise ExpectedError, (500, "Failed to read %s" % what)
 135   return line[:-1]
 136
 137 class EscapeHatch (BaseException):
 138   """Exception used by the `Escape' context manager"""
 139   def __init__(me): pass
 140
 141 class Escape (object):
 142   """
 143   A context manager.  Executes its body until completion or the `Escape'
 144   object itself is invoked as a function.  Other exceptions propagate
 145   normally.
 146   """
 147   def __init__(me):
 148     me.exc = EscapeHatch()
 149   def __call__(me):
 150     raise me.exc
 151   def __enter__(me):
 152     return me
 153   def __exit__(me, exty, exval, extb):
 154     return exval is me.exc
 155
 156 class Fluid (object):
 157   """
 158   Stores `fluid' variables which can be temporarily bound to new values, and
 159   later restored.
 160
 161   A caller may use the object's attributes for storing arbitrary values
 162   (though storing a `bind' value would be silly).  The `bind' method provides
 163   a context manager which binds attributes to other values during its
 164   execution.  This works even with multiple threads.
 165   """
 166
 167   ## We maintain two stores for variables.  One is a global store, `_g'; the
 168   ## other is a thread-local store `_t'.  We look for a variable first in the
 169   ## thread-local store, and then if necessary in the global store.  Binding
 170   ## works by remembering the old state of the variable on entry, setting it
 171   ## in the thread-local store (always), and then restoring the old state on
 172   ## exit.
 173
 174   ## A special marker for unbound variables.  If a variable is bound to a
 175   ## value, rebound temporarily with `bind', and then deleted, we must
 176   ## pretend that it's not there, and then restore it again afterwards.  We
 177   ## use this tag to mark variables which have been deleted while they're
 178   ## rebound.
 179   UNBOUND = Tag('unbound-variable')
 180
 181   def __init__(me, **kw):
 182     """Create a new set of fluid variables, initialized from the keywords."""
 183     me.__dict__.update(_g = struct(),
 184                        _t = TH.local())
 185     for k, v in kw.iteritems():
 186       setattr(me._g, k, v)
 187
 188   def __getattr__(me, k):
 189     """Return the current value stored with K, or raise AttributeError."""
 190     try: v = getattr(me._t, k)
 191     except AttributeError: v = getattr(me._g, k)
 192     if v is Fluid.UNBOUND: raise AttributeError, k
 193     return v
 194
 195   def __setattr__(me, k, v):
 196     """Associate the value V with the variable K."""
 197     if hasattr(me._t, k): setattr(me._t, k, v)
 198     else: setattr(me._g, k, v)
 199
 200   def __delattr__(me, k):
 201     """
 202     Forget about the variable K, so that attempts to read it result in an
 203     AttributeError.
 204     """
 205     if hasattr(me._t, k): setattr(me._t, k, Fluid.UNBOUND)
 206     else: delattr(me._g, k)
 207
 208   def __dir__(me):
 209     """Return a list of the currently known variables."""
 210     seen = set()
 211     keys = []
 212     for s in [me._t, me._g]:
 213       for k in dir(s):
 214         if k in seen: continue
 215         seen.add(k)
 216         if getattr(s, k) is not Fluid.UNBOUND: keys.append(k)
 217     return keys
 218
 219   @CTX.contextmanager
 220   def bind(me, **kw):
 221     """
 222     A context manager: bind values to variables according to the keywords KW,
 223     and execute the body; when the body exits, restore the rebound variables
 224     to their previous values.
 225     """
 226
 227     ## A list of things to do when we finish.
 228     unwind = []
 229
 230     def _delattr(k):
 231       ## Remove K from the thread-local store.  Only it might already have
 232       ## been deleted, so be careful.
 233       try: delattr(me._t, k)
 234       except AttributeError: pass
 235
 236     def stash(k):
 237       ## Stash a function for restoring the old state of K.  We do this here
 238       ## rather than inline only because Python's scoping rules are crazy and
 239       ## we need to ensure that all of the necessary variables are
 240       ## lambda-bound.
 241       try: ov = getattr(me._t, k)
 242       except AttributeError: unwind.append(lambda: _delattr(k))
 243       else: unwind.append(lambda: setattr(me._t, k, ov))
 244
 245     ## Rebind the variables.
 246     for k, v in kw.iteritems():
 247       stash(k)
 248       setattr(me._t, k, v)
 249
 250     ## Run the body, and restore.
 251     try: yield me
 252     finally:
 253       for f in unwind: f()
 254
 255 class Cleanup (object):
 256   """
 257   A context manager for stacking other context managers.
 258
 259   By itself, it does nothing.  Attach other context managers with `enter' or
 260   loose cleanup functions with `add'.  On exit, contexts are left and
 261   cleanups performed in reverse order.
 262   """
 263   def __init__(me):
 264     me._cleanups = []
 265   def __enter__(me):
 266     return me
 267   def __exit__(me, exty, exval, extb):
 268     trap = False
 269     for c in reversed(me._cleanups):
 270       if c(exty, exval, extb): trap = True
 271     return trap
 272   def enter(me, ctx):
 273     v = ctx.__enter__()
 274     me._cleanups.append(ctx.__exit__)
 275     return v
 276   def add(me, func):
 277     me._cleanups.append(lambda exty, exval, extb: func())
 278
 279 ###--------------------------------------------------------------------------
 280 ### Encodings.
 281
 282 class Encoding (object):
 283   """
 284   A pairing of injective encoding on binary strings, with its appropriate
 285   partial inverse.
 286
 287   The two functions are available in the `encode' and `decode' attributes.
 288   See also the `ENCODINGS' dictionary.
 289   """
 290   def __init__(me, encode, decode):
 291     me.encode = encode
 292     me.decode = decode
 293
 294 ENCODINGS = {
 295   'base64': Encoding(lambda s: BN.b64encode(s),
 296                      lambda s: BN.b64decode(s)),
 297   'base32': Encoding(lambda s: BN.b32encode(s).lower(),
 298                      lambda s: BN.b32decode(s, casefold = True)),
 299   'hex': Encoding(lambda s: BN.b16encode(s).lower(),
 300                   lambda s: BN.b16decode(s, casefold = True)),
 301   None: Encoding(identity, identity)
 302 }
 303
 304 ###--------------------------------------------------------------------------
 305 ### Time and timeouts.
 306
 307 def update_time():
 308   """
 309   Reset our idea of the current time, as kept in the global variable `NOW'.
 310   """
 311   global NOW
 312   NOW = int(T.time())
 313 update_time()
 314
 315 class Alarm (Exception):
 316   """
 317   Exception used internally by the `timeout' context manager.
 318
 319   If you're very unlucky, you might get one of these at top level.
 320   """
 321   pass
 322
 323 class Timeout (ExpectedError):
 324   """
 325   Report a timeout, from the `timeout' context manager.
 326   """
 327   def __init__(me, what):
 328     ExpectedError.__init__(me, 500, "Timeout %s" % what)
 329
 330 ## Set `DEADLINE' to be the absolute time of the next alarm.  We'll keep this
 331 ## up to date in `timeout'.
 332 delta, _ = SIG.getitimer(SIG.ITIMER_REAL)
 333 if delta == 0: DEADLINE = None
 334 else: DEADLINE = NOW + delta
 335
 336 def _alarm(sig, tb):
 337   """If we receive `SIGALRM', raise the alarm."""
 338   raise Alarm
 339 SIG.signal(SIG.SIGALRM, _alarm)
 340
 341 @CTX.contextmanager
 342 def timeout(delta, what):
 343   """
 344   A context manager which interrupts execution of its body after DELTA
 345   seconds, if it doesn't finish before then.
 346
 347   If execution is interrupted, a `Timeout' exception is raised, carrying WHY
 348   (a gerund phrase) as part of its message.
 349   """
 350
 351   global DEADLINE
 352   when = NOW + delta
 353   if DEADLINE is not None and when >= DEADLINE:
 354     yield
 355     update_time()
 356   else:
 357     od = DEADLINE
 358     try:
 359       DEADLINE = when
 360       SIG.setitimer(SIG.ITIMER_REAL, delta)
 361       yield
 362     except Alarm:
 363       raise Timeout, what
 364     finally:
 365       update_time()
 366       DEADLINE = od
 367       if od is None: SIG.setitimer(SIG.ITIMER_REAL, 0)
 368       else: SIG.setitimer(SIG.ITIMER_REAL, DEADLINE - NOW)
 369
 370 ###--------------------------------------------------------------------------
 371 ### File locking.
 372
 373 @CTX.contextmanager
 374 def lockfile(lock, t = None):
 375   """
 376   Acquire an exclusive lock on a named file LOCK while executing the body.
 377
 378   If T is zero, fail immediately if the lock can't be acquired; if T is none,
 379   then wait forever if necessary; otherwise give up after T seconds.
 380   """
 381   fd = -1
 382   try:
 383     fd = OS.open(lock, OS.O_WRONLY | OS.O_CREAT, 0600)
 384     if timeout is None:
 385       F.lockf(fd, F.LOCK_EX)
 386     elif timeout == 0:
 387       F.lockf(fd, F.LOCK_EX | F.LOCK_NB)
 388     else:
 389       with timeout(t, "waiting for lock file `%s'" % lock):
 390         F.lockf(fd, F.LOCK_EX)
 391     yield None
 392   finally:
 393     if fd != -1: OS.close(fd)
 394
 395 ###--------------------------------------------------------------------------
 396 ### Database utilities.
 397
 398 ### Python's database API is dreadful: it exposes far too many
 399 ### implementation-specific details to the programmer, who may well want to
 400 ### write code which works against many different databases.
 401 ###
 402 ### One particularly frustrating problem is the variability of placeholder
 403 ### syntax in SQL statements: there's no universal convention, just a number
 404 ### of possible syntaxes, at least one of which will be implemented (and some
 405 ### of which are mutually incompatible).  Because not doing this invites all
 406 ### sorts of misery such as SQL injection vulnerabilties, we introduce a
 407 ### simple abstraction.  A database parameter-type object keeps track of one
 408 ### particular convention, providing the correct placeholders to be inserted
 409 ### into the SQL command string, and the corresponding arguments, in whatever
 410 ### way is necessary.
 411 ###
 412 ### The protocol is fairly simple.  An object of the appropriate class is
 413 ### instantiated for each SQL statement, providing it with a dictionary
 414 ### mapping placeholder names to their values.  The object's `sub' method is
 415 ### called for each placeholder found in the statement, with a match object
 416 ### as an argument; the match object picks out the name of the placeholder in
 417 ### question in group 1, and the method returns a piece of syntax appropriate
 418 ### to the database backend.  Finally, the collected arguments are made
 419 ### available, in whatever format is required, in the object's `args'
 420 ### attribute.
 421
 422 ## Turn simple Unix not-quite-glob patterns into SQL `LIKE' patterns.
 423 ## Match using: x LIKE y ESCAPE '\\'
 424 globtolike = StringSubst({
 425   '\\*': '*', '%': '\\%', '*': '%',
 426   '\\?': '?', '_': '\\_', '?': '_'
 427 })
 428
 429 class LinearParam (object):
 430   """
 431   Abstract parent class for `linear' parameter conventions.
 432
 433   A linear convention is one where the arguments are supplied as a list, and
 434   placeholders are either all identical (with semantics `insert the next
 435   argument'), or identify their argument by its position within the list.
 436   """
 437   def __init__(me, kw):
 438     me._i = 0
 439     me.args = []
 440     me._kw = kw
 441   def sub(me, match):
 442     name = match.group(1)
 443     me.args.append(me._kw[name])
 444     marker = me._format()
 445     me._i += 1
 446     return marker
 447 class QmarkParam (LinearParam):
 448   def _format(me): return '?'
 449 class NumericParam (LinearParam):
 450   def _format(me): return ':%d' % me._i
 451 class FormatParam (LinearParam):
 452   def _format(me): return '%s'
 453
 454 class DictParam (object):
 455   """
 456   Abstract parent class for `dictionary' parameter conventions.
 457
 458   A dictionary convention is one where the arguments are provided as a
 459   dictionary, and placeholders contain a key name identifying the
 460   corresponding value in that dictionary.
 461   """
 462   def __init__(me, kw):
 463     me.args = kw
 464   def sub(me, match):
 465     name = match.group(1)
 466     return me._format(name)
 467 def NamedParam (object):
 468   def _format(me, name): return ':%s' % name
 469 def PyFormatParam (object):
 470   def _format(me, name): return '%%(%s)s' % name
 471
 472 ### Since we're doing a bunch of work to paper over idiosyncratic placeholder
 473 ### syntax, we might as well also sort out other problems.  The `DB_FIXUPS'
 474 ### dictionary maps database module names to functions which might need to do
 475 ### clever stuff at connection setup time.
 476
 477 DB_FIXUPS = {}
 478
 479 @register(DB_FIXUPS, 'sqlite3')
 480 def fixup_sqlite3(db):
 481   """
 482   Unfortunately, SQLite learnt about FOREIGN KEY constraints late, and so
 483   doesn't enforce them unless explicitly told to.
 484   """
 485   c = db.cursor()
 486   c.execute("PRAGMA foreign_keys = ON")
 487
 488 class SimpleDBConnection (object):
 489   """
 490   Represents a database connection, while trying to hide the differences
 491   between various kinds of database backends.
 492   """
 493
 494   __metaclass__ = DictExpanderClass
 495
 496   ## A map from placeholder convention names to classes implementing them.
 497   PLACECLS = {
 498     'qmark': QmarkParam,
 499     'numeric': NumericParam,
 500     'named': NamedParam,
 501     'format': FormatParam,
 502     'pyformat': PyFormatParam
 503   }
 504
 505   ## A pattern for our own placeholder syntax.
 506   R_PLACE = RX.compile(r'\$(\w+)')
 507
 508   def __init__(me, modname, modargs):
 509     """
 510     Make a new database connection, using the module MODNAME, and passing its
 511     `connect' function the MODARGS -- which may be either a list or a
 512     dictionary.
 513     """
 514
 515     ## Get the module, and create a connection.
 516     mod = __import__(modname)
 517     if isinstance(modargs, dict): me._db = mod.connect(**modargs)
 518     else: me._db = mod.connect(*modargs)
 519
 520     ## Apply any necessary fixups.
 521     try: fixup = DB_FIXUPS[modname]
 522     except KeyError: pass
 523     else: fixup(me._db)
 524
 525     ## Grab hold of other interesting things.
 526     me.Error = mod.Error
 527     me.Warning = mod.Warning
 528     me._placecls = me.PLACECLS[mod.paramstyle]
 529
 530   def execute(me, command, **kw):
 531     """
 532     Execute the SQL COMMAND.  The keyword arguments are used to provide
 533     values corresponding to `$NAME' placeholders in the COMMAND.
 534
 535     Return the receiver, so that iterator protocol is convenient.
 536     """
 537     me._cur = me._db.cursor()
 538     plc = me._placecls(kw)
 539     subst = me.R_PLACE.sub(plc.sub, command)
 540     ##PRINT('*** %s : %r' % (subst, plc.args))
 541     me._cur.execute(subst, plc.args)
 542     return me
 543
 544   def __iter__(me):
 545     """Iterator protocol: simply return the receiver."""
 546     return me
 547   def next(me):
 548     """Iterator protocol: return the next row from the current query."""
 549     row = me.fetchone()
 550     if row is None: raise StopIteration
 551     return row
 552
 553   def __enter__(me):
 554     """
 555     Context protocol: begin a transaction.
 556     """
 557     ##PRINT('<<< BEGIN')
 558     return
 559   def __exit__(me, exty, exval, tb):
 560     """Context protocol: commit or roll back a transaction."""
 561     if exty:
 562       ##PRINT('>*> ROLLBACK')
 563       me.rollback()
 564     else:
 565       ##PRINT('>>> COMMIT')
 566       me.commit()
 567
 568   ## Import a number of methods from the underlying connection.
 569   __extra__ = {}
 570   for _name in ['fetchone', 'fetchmany', 'fetchall']:
 571     def _(name, extra):
 572       extra[name] = lambda me, *args, **kw: \
 573                     getattr(me._cur, name)(*args, **kw)
 574     _(_name, __extra__)
 575   for _name in ['commit', 'rollback']:
 576     def _(name, extra):
 577       extra[name] = lambda me, *args, **kw: \
 578                     getattr(me._db, name)(*args, **kw)
 579     _(_name, __extra__)
 580   del _name, _
 581
 582 ###----- That's all, folks --------------------------------------------------