backend.py: Separate out the main work of `_update'.

[chopwood] / backend.py

### -*-python-*-
###
### Password backends
###
### (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 os as OS; ENV = OS.environ

import config as CONF; CFG = CONF.CFG
import util as U

###--------------------------------------------------------------------------
### Relevant configuration.

CONF.DEFAULTS.update(

  ## A directory in which we can create lockfiles.
  LOCKDIR = OS.path.join(ENV['HOME'], 'var', 'lock', 'chpwd'))

###--------------------------------------------------------------------------
### Protocol.
###
### A password backend knows how to fetch and modify records in some password
### database, e.g., a flat passwd(5)-style password file, or a table in some
### proper grown-up SQL database.
###
### A backend's `lookup' method retrieves the record for a named user from
### the database, returning it in a record object, or raises `UnknownUser'.
### The record object maintains `user' (the user name, as supplied to
### `lookup') and `passwd' (the encrypted password, in whatever form the
### underlying database uses) attributes, and possibly others.  The `passwd'
### attribute (at least) may be modified by the caller.  The record object
### has a `write' method, which updates the corresponding record in the
### database.
###
### The concrete record objects defined here inherit from `BasicRecord',
### which keeps track of its parent backend, and implements `write' by
### calling the backend's `_update' method.  Some backends require that their
### record objects implement additional private protocols.

class UnknownUser (U.ExpectedError):
  """The named user wasn't found in the database."""
  def __init__(me, user):
    U.ExpectedError.__init__(me, 500, "Unknown user `%s'" % user)
    me.user = user

class BasicRecord (object):
  """
  A handy base class for record classes.

  Keep track of the backend in `_be', and call its `_update' method to write
  ourselves back.
  """
  def __init__(me, backend):
    me._be = backend
  def write(me):
    me._be._update(me)

class TrivialRecord (BasicRecord):
  """
  A trivial record which simply remembers `user' and `passwd' attributes.

  Additional attributes can be set on the object if this is convenient.
  """
  def __init__(me, user, passwd, *args, **kw):
    super(TrivialRecord, me).__init__(*args, **kw)
    me.user = user
    me.passwd = passwd

###--------------------------------------------------------------------------
### Flat files.

class FlatFileRecord (BasicRecord):
  """
  A record from a flat-file database (like a passwd(5) file).

  Such a file carries one record per line; each record is split into fields
  by a delimiter character, specified by the DELIM constructor argument.

  The FMAP argument to the constructor maps names to field index numbers.
  The standard `user' and `passwd' fields must be included in this map if the
  object is to implement the protocol correctly (though the `FlatFileBackend'
  is careful to do this).
  """

  def __init__(me, line, delim, fmap, *args, **kw):
    """
    Initialize the record, splitting the LINE into fields separated by DELIM,
    and setting attributes under control of FMAP.
    """
    super(FlatFileRecord, me).__init__(*args, **kw)
    line = line.rstrip('\n')
    fields = line.split(delim)
    me._delim = delim
    me._fmap = fmap
    me._raw = fields
    for k, v in fmap.iteritems():
      setattr(me, k, fields[v])

  def _format(me):
    """
    Format the record as a line of text.

    The flat-file format is simple, but rather fragile with respect to
    invalid characters, and often processed by substandard software, so be
    careful not to allow bad characters into the file.
    """
    fields = me._raw
    for k, v in me._fmap.iteritems():
      val = getattr(me, k)
      for badch, what in [(me._delim, "delimiter `%s'" % me._delim),
                          ('\n', 'newline character'),
                          ('\0', 'null character')]:
        if badch in val:
          raise U.ExpectedError, \
                (500, "New `%s' field contains %s" % (k, what))
      fields[v] = val
    return me._delim.join(fields) + '\n'

class FlatFileBackend (object):
  """
  Password storage in a flat passwd(5)-style file.

  The FILE constructor argument names the file.  Such a file carries one
  record per line; each record is split into fields by a delimiter character,
  specified by the DELIM constructor argument.

  The file is updated by writing a new version alongside, as `FILE.new', and
  renaming it over the old version.  If a LOCK file is named then an
  exclusive fcntl(2)-style lock is taken out on `LOCKDIR/LOCK' (creating the
  file if necessary) during the update operation.  Use of a lockfile is
  strongly recommended.

  The DELIM constructor argument specifies the delimiter character used when
  splitting lines into fields.  The USER and PASSWD arguments give the field
  numbers (starting from 0) for the user-name and hashed-password fields;
  additional field names may be given using keyword arguments: the values of
  these fields are exposed as attributes `f_NAME' on record objects.
  """

  def __init__(me, file, lock = None,
               delim = ':', user = 0, passwd = 1, **fields):
    """
    Construct a new flat-file backend object.  See the class documentation
    for details.
    """
    me._lock = lock
    me._file = file
    me._delim = delim
    fmap = dict(user = user, passwd = passwd)
    for k, v in fields.iteritems(): fmap['f_' + k] = v
    me._fmap = fmap

  def lookup(me, user):
    """Return the record for the named USER."""
    with open(me._file) as f:
      for line in f:
        rec = me._parse(line)
        if rec.user == user:
          return rec
    raise UnknownUser, user

  def _rewrite(me, op, rec):
    """
    Rewrite the file, according to OP.

    The OP may be one of the following.

    `create'            There must not be a record matching REC; add a new
                        one.

    `remove'            There must be a record matching REC: remove it.

    `update'            There must be a record matching REC: write REC in its
                        place.
    """

    ## The main update function.
    def doit():

      ## Make sure we preserve the file permissions, and in particular don't
      ## allow a window during which the new file has looser permissions than
      ## the old one.
      st = OS.stat(me._file)
      tmp = me._file + '.new'
      fd = OS.open(tmp, OS.O_WRONLY | OS.O_CREAT | OS.O_EXCL, st.st_mode)

      ## This is the fiddly bit.
      lose = True
      try:

        ## Copy the old file to the new one, changing the user's record if
        ## and when we encounter it.
        found = False
        with OS.fdopen(fd, 'w') as f_out:
          with open(me._file) as f_in:
            for line in f_in:
              r = me._parse(line)
              if r.user != rec.user:
                f_out.write(line)
              elif op == 'create':
                raise U.ExpectedError, \
                    (500, "Record for `%s' already exists" % rec.user)
              else:
                found = True
                if op != 'remove': f_out.write(rec._format())
          if found:
            pass
          elif op == 'create':
            f_out.write(rec._format())
          else:
            raise U.ExpectedError, \
                (500, "Record for `%s' not found" % rec.user)

        ## Update the permissions on the new file.  Don't try to fix the
        ## ownership (we shouldn't be running as root) or the group (the
        ## parent directory should have the right permissions already).
        OS.chmod(tmp, st.st_mode)
        OS.rename(tmp, me._file)
        lose = False
      except OSError, e:
        ## I suppose that system errors are to be expected at this point.
        raise U.ExpectedError, \
              (500, "Failed to update `%s': %s" % (me._file, e))
      finally:
        ## Don't try to delete the new file if we succeeded: it might belong
        ## to another instance of us.
        if lose:
          try: OS.unlink(tmp)
          except: pass

    ## If there's a lockfile, then acquire it around the meat of this
    ## function; otherwise just do the job.
    if me._lock is None:
      doit()
    else:
      with U.lockfile(OS.path.join(CFG.LOCKDIR, me._lock), 5):
        doit()

  def _parse(me, line):
    """Convenience function for constructing a record."""
    return FlatFileRecord(line, me._delim, me._fmap, backend = me)

  def _update(me, rec):
    """Update the record REC in the file."""
    me._rewrite('update', rec)

CONF.export('FlatFileBackend')

###--------------------------------------------------------------------------
### SQL databases.

class DatabaseBackend (object):
  """
  Password storage in a SQL database table.

  We assume that there's a single table mapping user names to (hashed)
  passwords: we won't try anything complicated involving joins.

  We need to know a database module MODNAME and arguments MODARGS to pass to
  the `connect' function.  We also need to know the TABLE to search, and the
  USER and PASSWD field names.  Additional field names can be passed to the
  constructor: these will be read from the database and attached as
  attributes `f_NAME' to the record returned by `lookup'.  Changes to these
  attributes are currently not propagated back to the database.
  """

  def __init__(me, modname, modargs, table, user, passwd, *fields):
    """
    Create a database backend object.  See the class docstring for details.
    """
    me._table = table
    me._user = user
    me._passwd = passwd
    me._fields = list(fields)

    ## We don't connect immediately.  That would be really bad if we had lots
    ## of database backends running at a time, because we probably only want
    ## to use one.
    me._db = None
    me._modname = modname
    me._modargs = modargs

  def _connect(me):
    """Set up the lazy connection to the database."""
    if me._db is None:
      me._db = U.SimpleDBConnection(me._modname, me._modargs)

  def lookup(me, user):
    """Return the record for the named USER."""
    me._connect()
    me._db.execute("SELECT %s FROM %s WHERE %s = $user" %
                   (', '.join([me._passwd] + me._fields),
                    me._table, me._user),
                   user = user)
    row = me._db.fetchone()
    if row is None: raise UnknownUser, user
    passwd = row[0]
    rec = TrivialRecord(backend = me, user = user, passwd = passwd)
    for f, v in zip(me._fields, row[1:]):
      setattr(rec, 'f_' + f, v)
    return rec

  def _update(me, rec):
    """Update the record REC in the database."""
    me._connect()
    with me._db:
      me._db.execute(
        "UPDATE %s SET %s = $passwd WHERE %s = $user" % (
          me._table, me._passwd, me._user),
        user = rec.user, passwd = rec.passwd)

CONF.export('DatabaseBackend')

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