agpl.py (filez): Check the exit code from the command.

[chopwood] / util.py

### -*-python-*-
###
### Miscellaneous utilities
###
### (c) 2013 Mark Wooding
###

###----- Licensing notice ---------------------------------------------------
###
### This file is part of Chopwood: a password-changing service.
###
### Chopwood is free software; you can redistribute it and/or modify
### it under the terms of the GNU Affero General Public License as
### published by the Free Software Foundation; either version 3 of the
### License, or (at your option) any later version.
###
### Chopwood is distributed in the hope that it will be useful,
### but WITHOUT ANY WARRANTY; without even the implied warranty of
### MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
### GNU Affero General Public License for more details.
###
### You should have received a copy of the GNU Affero General Public
### License along with Chopwood; if not, see
### <http://www.gnu.org/licenses/>.

from __future__ import with_statement

import base64 as BN
import contextlib as CTX
import fcntl as F
import os as OS
import re as RX
import signal as SIG
import sys as SYS
import time as T

try: import threading as TH
except ImportError: import dummy_threading as TH

###--------------------------------------------------------------------------
### Some basics.

def identity(x):
  """The identity function: returns its argument."""
  return x

def constantly(x):
  """The function which always returns X."""
  return lambda: x

class struct (object):
  """A simple object for storing data in attributes."""
  DEFAULTS = {}
  def __init__(me, *args, **kw):
    cls = me.__class__
    for k, v in kw.iteritems(): setattr(me, k, v)
    try:
      slots = cls.__slots__
    except AttributeError:
      if args: raise ValueError, 'no slots defined'
    else:
      if len(args) > len(slots): raise ValueError, 'too many arguments'
      for k, v in zip(slots, args): setattr(me, k, v)
      for k in slots:
        if hasattr(me, k): continue
        try: setattr(me, k, cls.DEFAULTS[k])
        except KeyError: raise ValueError, "no value for `%s'" % k

class Tag (object):
  """An object whose only purpose is to be distinct from other objects."""
  def __init__(me, name):
    me._name = name
  def __repr__(me):
    return '#<%s %r>' % (type(me).__name__, me._name)

class DictExpanderClass (type):
  """
  Metaclass for classes with autogenerated members.

  If the class body defines a dictionary `__extra__' then the key/value pairs
  in this dictionary are promoted into attributes of the class.  This is much
  easier -- and safer -- than fiddling about with `locals'.
  """
  def __new__(cls, name, supers, dict):
    try:
      ex = dict['__extra__']
    except KeyError:
      pass
    else:
      for k, v in ex.iteritems():
        dict[k] = v
      del dict['__extra__']
    return super(DictExpanderClass, cls).__new__(cls, name, supers, dict)

class ExpectedError (Exception):
  """
  A (concrete) base class for various errors we expect to encounter.

  The `msg' attribute carries a human-readable message explaining what the
  problem actually is.  The really important bit, though, is the `code'
  attribute, which carries an HTTP status code to be reported to the user
  agent, if we're running through CGI.
  """
  def __init__(me, code, msg):
    me.code = code
    me.msg = msg
  def __str__(me):
    return '%s (%d)' % (me.msg, me.code)

def register(dict, name):
  """A decorator: add the decorated function to DICT, under the key NAME."""
  def _(func):
    dict[name] = func
    return func
  return _

class StringSubst (object):
  """
  A string substitution.  Initialize with a dictionary mapping source strings
  to target strings.  The object is callable, and maps strings in the obvious
  way.
  """
  def __init__(me, map):
    me._map = map
    me._rx = RX.compile('|'.join(RX.escape(s) for s in map))
  def __call__(me, s):
    return me._rx.sub(lambda m: me._map[m.group(0)], s)

def readline(what, file = SYS.stdin):
  """Read a single line from FILE (default stdin) and return it."""
  try: line = SYS.stdin.readline()
  except IOError, e: raise ExpectedError, (500, str(e))
  if not line.endswith('\n'):
    raise ExpectedError, (500, "Failed to read %s" % what)
  return line[:-1]

class EscapeHatch (BaseException):
  """Exception used by the `Escape' context manager"""
  def __init__(me): pass

class Escape (object):
  """
  A context manager.  Executes its body until completion or the `Escape'
  object itself is invoked as a function.  Other exceptions propagate
  normally.
  """
  def __init__(me):
    me.exc = EscapeHatch()
  def __call__(me):
    raise me.exc
  def __enter__(me):
    return me
  def __exit__(me, exty, exval, extb):
    return exval is me.exc

class Fluid (object):
  """
  Stores `fluid' variables which can be temporarily bound to new values, and
  later restored.

  A caller may use the object's attributes for storing arbitrary values
  (though storing a `bind' value would be silly).  The `bind' method provides
  a context manager which binds attributes to other values during its
  execution.  This works even with multiple threads.
  """

  ## We maintain two stores for variables.  One is a global store, `_g'; the
  ## other is a thread-local store `_t'.  We look for a variable first in the
  ## thread-local store, and then if necessary in the global store.  Binding
  ## works by remembering the old state of the variable on entry, setting it
  ## in the thread-local store (always), and then restoring the old state on
  ## exit.

  ## A special marker for unbound variables.  If a variable is bound to a
  ## value, rebound temporarily with `bind', and then deleted, we must
  ## pretend that it's not there, and then restore it again afterwards.  We
  ## use this tag to mark variables which have been deleted while they're
  ## rebound.
  UNBOUND = Tag('unbound-variable')

  def __init__(me, **kw):
    """Create a new set of fluid variables, initialized from the keywords."""
    me.__dict__.update(_g = struct(),
                       _t = TH.local())
    for k, v in kw.iteritems():
      setattr(me._g, k, v)

  def __getattr__(me, k):
    """Return the current value stored with K, or raise AttributeError."""
    try: v = getattr(me._t, k)
    except AttributeError: v = getattr(me._g, k)
    if v is Fluid.UNBOUND: raise AttributeError, k
    return v

  def __setattr__(me, k, v):
    """Associate the value V with the variable K."""
    if hasattr(me._t, k): setattr(me._t, k, v)
    else: setattr(me._g, k, v)

  def __delattr__(me, k):
    """
    Forget about the variable K, so that attempts to read it result in an
    AttributeError.
    """
    if hasattr(me._t, k): setattr(me._t, k, Fluid.UNBOUND)
    else: delattr(me._g, k)

  def __dir__(me):
    """Return a list of the currently known variables."""
    seen = set()
    keys = []
    for s in [me._t, me._g]:
      for k in dir(s):
        if k in seen: continue
        seen.add(k)
        if getattr(s, k) is not Fluid.UNBOUND: keys.append(k)
    return keys

  @CTX.contextmanager
  def bind(me, **kw):
    """
    A context manager: bind values to variables according to the keywords KW,
    and execute the body; when the body exits, restore the rebound variables
    to their previous values.
    """

    ## A list of things to do when we finish.
    unwind = []

    def _delattr(k):
      ## Remove K from the thread-local store.  Only it might already have
      ## been deleted, so be careful.
      try: delattr(me._t, k)
      except AttributeError: pass

    def stash(k):
      ## Stash a function for restoring the old state of K.  We do this here
      ## rather than inline only because Python's scoping rules are crazy and
      ## we need to ensure that all of the necessary variables are
      ## lambda-bound.
      try: ov = getattr(me._t, k)
      except AttributeError: unwind.append(lambda: _delattr(k))
      else: unwind.append(lambda: setattr(me._t, k, ov))

    ## Rebind the variables.
    for k, v in kw.iteritems():
      stash(k)
      setattr(me._t, k, v)

    ## Run the body, and restore.
    try: yield me
    finally:
      for f in unwind: f()

class Cleanup (object):
  """
  A context manager for stacking other context managers.

  By itself, it does nothing.  Attach other context managers with `enter' or
  loose cleanup functions with `add'.  On exit, contexts are left and
  cleanups performed in reverse order.
  """
  def __init__(me):
    me._cleanups = []
  def __enter__(me):
    return me
  def __exit__(me, exty, exval, extb):
    trap = False
    for c in reversed(me._cleanups):
      if c(exty, exval, extb): trap = True
    return trap
  def enter(me, ctx):
    v = ctx.__enter__()
    me._cleanups.append(ctx.__exit__)
    return v
  def add(me, func):
    me._cleanups.append(lambda exty, exval, extb: func())

###--------------------------------------------------------------------------
### Encodings.

class Encoding (object):
  """
  A pairing of injective encoding on binary strings, with its appropriate
  partial inverse.

  The two functions are available in the `encode' and `decode' attributes.
  See also the `ENCODINGS' dictionary.
  """
  def __init__(me, encode, decode):
    me.encode = encode
    me.decode = decode

ENCODINGS = {
  'base64': Encoding(lambda s: BN.b64encode(s),
                     lambda s: BN.b64decode(s)),
  'base32': Encoding(lambda s: BN.b32encode(s).lower(),
                     lambda s: BN.b32decode(s, casefold = True)),
  'hex': Encoding(lambda s: BN.b16encode(s).lower(),
                  lambda s: BN.b16decode(s, casefold = True)),
  None: Encoding(identity, identity)
}

###--------------------------------------------------------------------------
### Time and timeouts.

def update_time():
  """
  Reset our idea of the current time, as kept in the global variable `NOW'.
  """
  global NOW
  NOW = int(T.time())
update_time()

class Alarm (Exception):
  """
  Exception used internally by the `timeout' context manager.

  If you're very unlucky, you might get one of these at top level.
  """
  pass

class Timeout (ExpectedError):
  """
  Report a timeout, from the `timeout' context manager.
  """
  def __init__(me, what):
    ExpectedError.__init__(me, 500, "Timeout %s" % what)

## Set `DEADLINE' to be the absolute time of the next alarm.  We'll keep this
## up to date in `timeout'.
delta, _ = SIG.getitimer(SIG.ITIMER_REAL)
if delta == 0: DEADLINE = None
else: DEADLINE = NOW + delta

def _alarm(sig, tb):
  """If we receive `SIGALRM', raise the alarm."""
  raise Alarm
SIG.signal(SIG.SIGALRM, _alarm)

@CTX.contextmanager
def timeout(delta, what):
  """
  A context manager which interrupts execution of its body after DELTA
  seconds, if it doesn't finish before then.

  If execution is interrupted, a `Timeout' exception is raised, carrying WHY
  (a gerund phrase) as part of its message.
  """

  global DEADLINE
  when = NOW + delta
  if DEADLINE is not None and when >= DEADLINE:
    yield
    update_time()
  else:
    od = DEADLINE
    try:
      DEADLINE = when
      SIG.setitimer(SIG.ITIMER_REAL, delta)
      yield
    except Alarm:
      raise Timeout, what
    finally:
      update_time()
      DEADLINE = od
      if od is None: SIG.setitimer(SIG.ITIMER_REAL, 0)
      else: SIG.setitimer(SIG.ITIMER_REAL, DEADLINE - NOW)

###--------------------------------------------------------------------------
### File locking.

@CTX.contextmanager
def lockfile(lock, t = None):
  """
  Acquire an exclusive lock on a named file LOCK while executing the body.

  If T is zero, fail immediately if the lock can't be acquired; if T is none,
  then wait forever if necessary; otherwise give up after T seconds.
  """
  fd = -1
  try:
    fd = OS.open(lock, OS.O_WRONLY | OS.O_CREAT, 0600)
    if timeout is None:
      F.lockf(fd, F.LOCK_EX)
    elif timeout == 0:
      F.lockf(fd, F.LOCK_EX | F.LOCK_NB)
    else:
      with timeout(t, "waiting for lock file `%s'" % lock):
        F.lockf(fd, F.LOCK_EX)
    yield None
  finally:
    if fd != -1: OS.close(fd)

###--------------------------------------------------------------------------
### Database utilities.

### Python's database API is dreadful: it exposes far too many
### implementation-specific details to the programmer, who may well want to
### write code which works against many different databases.
###
### One particularly frustrating problem is the variability of placeholder
### syntax in SQL statements: there's no universal convention, just a number
### of possible syntaxes, at least one of which will be implemented (and some
### of which are mutually incompatible).  Because not doing this invites all
### sorts of misery such as SQL injection vulnerabilties, we introduce a
### simple abstraction.  A database parameter-type object keeps track of one
### particular convention, providing the correct placeholders to be inserted
### into the SQL command string, and the corresponding arguments, in whatever
### way is necessary.
###
### The protocol is fairly simple.  An object of the appropriate class is
### instantiated for each SQL statement, providing it with a dictionary
### mapping placeholder names to their values.  The object's `sub' method is
### called for each placeholder found in the statement, with a match object
### as an argument; the match object picks out the name of the placeholder in
### question in group 1, and the method returns a piece of syntax appropriate
### to the database backend.  Finally, the collected arguments are made
### available, in whatever format is required, in the object's `args'
### attribute.

## Turn simple Unix not-quite-glob patterns into SQL `LIKE' patterns.
## Match using: x LIKE y ESCAPE '\\'
globtolike = StringSubst({
  '\\*': '*', '%': '\\%', '*': '%',
  '\\?': '?', '_': '\\_', '?': '_'
})

class LinearParam (object):
  """
  Abstract parent class for `linear' parameter conventions.

  A linear convention is one where the arguments are supplied as a list, and
  placeholders are either all identical (with semantics `insert the next
  argument'), or identify their argument by its position within the list.
  """
  def __init__(me, kw):
    me._i = 0
    me.args = []
    me._kw = kw
  def sub(me, match):
    name = match.group(1)
    me.args.append(me._kw[name])
    marker = me._format()
    me._i += 1
    return marker
class QmarkParam (LinearParam):
  def _format(me): return '?'
class NumericParam (LinearParam):
  def _format(me): return ':%d' % me._i
class FormatParam (LinearParam):
  def _format(me): return '%s'

class DictParam (object):
  """
  Abstract parent class for `dictionary' parameter conventions.

  A dictionary convention is one where the arguments are provided as a
  dictionary, and placeholders contain a key name identifying the
  corresponding value in that dictionary.
  """
  def __init__(me, kw):
    me.args = kw
  def sub(me, match):
    name = match.group(1)
    return me._format(name)
def NamedParam (object):
  def _format(me, name): return ':%s' % name
def PyFormatParam (object):
  def _format(me, name): return '%%(%s)s' % name

### Since we're doing a bunch of work to paper over idiosyncratic placeholder
### syntax, we might as well also sort out other problems.  The `DB_FIXUPS'
### dictionary maps database module names to functions which might need to do
### clever stuff at connection setup time.

DB_FIXUPS = {}

@register(DB_FIXUPS, 'sqlite3')
def fixup_sqlite3(db):
  """
  Unfortunately, SQLite learnt about FOREIGN KEY constraints late, and so
  doesn't enforce them unless explicitly told to.
  """
  c = db.cursor()
  c.execute("PRAGMA foreign_keys = ON")

class SimpleDBConnection (object):
  """
  Represents a database connection, while trying to hide the differences
  between various kinds of database backends.
  """

  __metaclass__ = DictExpanderClass

  ## A map from placeholder convention names to classes implementing them.
  PLACECLS = {
    'qmark': QmarkParam,
    'numeric': NumericParam,
    'named': NamedParam,
    'format': FormatParam,
    'pyformat': PyFormatParam
  }

  ## A pattern for our own placeholder syntax.
  R_PLACE = RX.compile(r'\$(\w+)')

  def __init__(me, modname, modargs):
    """
    Make a new database connection, using the module MODNAME, and passing its
    `connect' function the MODARGS -- which may be either a list or a
    dictionary.
    """

    ## Get the module, and create a connection.
    mod = __import__(modname)
    if isinstance(modargs, dict): me._db = mod.connect(**modargs)
    else: me._db = mod.connect(*modargs)

    ## Apply any necessary fixups.
    try: fixup = DB_FIXUPS[modname]
    except KeyError: pass
    else: fixup(me._db)

    ## Grab hold of other interesting things.
    me.Error = mod.Error
    me.Warning = mod.Warning
    me._placecls = me.PLACECLS[mod.paramstyle]

  def execute(me, command, **kw):
    """
    Execute the SQL COMMAND.  The keyword arguments are used to provide
    values corresponding to `$NAME' placeholders in the COMMAND.

    Return the receiver, so that iterator protocol is convenient.
    """
    me._cur = me._db.cursor()
    plc = me._placecls(kw)
    subst = me.R_PLACE.sub(plc.sub, command)
    ##PRINT('*** %s : %r' % (subst, plc.args))
    me._cur.execute(subst, plc.args)
    return me

  def __iter__(me):
    """Iterator protocol: simply return the receiver."""
    return me
  def next(me):
    """Iterator protocol: return the next row from the current query."""
    row = me.fetchone()
    if row is None: raise StopIteration
    return row

  def __enter__(me):
    """
    Context protocol: begin a transaction.
    """
    ##PRINT('<<< BEGIN')
    return
  def __exit__(me, exty, exval, tb):
    """Context protocol: commit or roll back a transaction."""
    if exty:
      ##PRINT('>*> ROLLBACK')
      me.rollback()
    else:
      ##PRINT('>>> COMMIT')
      me.commit()

  ## Import a number of methods from the underlying connection.
  __extra__ = {}
  for _name in ['fetchone', 'fetchmany', 'fetchall']:
    def _(name, extra):
      extra[name] = lambda me, *args, **kw: \
                    getattr(me._cur, name)(*args, **kw)
    _(_name, __extra__)
  for _name in ['commit', 'rollback']:
    def _(name, extra):
      extra[name] = lambda me, *args, **kw: \
                    getattr(me._db, name)(*args, **kw)
    _(_name, __extra__)
  del _name, _

###----- That's all, folks --------------------------------------------------