agpl.py: Document and prettify.

[chopwood] / agpl.py

### -*-python-*-
###
### GNU Affero General Public License compliance
###
### (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/>.

import contextlib as CTX
import grp as GR
import os as OS
import pwd as PW
import shlex as SL
import shutil as SH
import subprocess as SUB
import sys as SYS
import tarfile as TAR
import tempfile as TF
import time as T

from cStringIO import StringIO

from auto import PACKAGE, VERSION
import util as U

###--------------------------------------------------------------------------
### Initial utilities.

@CTX.contextmanager
def tempdir():
  """
  Context manager: create and return the name of a temporary directory.

  The directory will be deleted automatically on exit from the body.
  """
  d = TF.mkdtemp()
  try: yield d
  finally: SH.rmtree(d, ignore_errors = True)

###--------------------------------------------------------------------------
### Determining which files to include.

def dirs_to_dump():
  """
  Return a list of directories containing used Python modules.

  Directories under `/usr/' but outside `/usr/local/' are excluded, since
  they are assumed to be covered by the AGPL exception for parts of the
  operating system.
  """

  ## Collect a set of directories.
  dirs = set()

  ## Work through the list of known modules, adding them to the list.
  for m in SYS.modules.itervalues():
    try: f = m.__file__
    except AttributeError: continue
    d = OS.path.realpath(OS.path.dirname(f))
    if d.startswith('/usr/') and not d.startswith('/usr/local/'): continue
    dirs.add(d)

  ## Now go through the directories again, and remove any which are wholly
  ## included within other entries.
  dirs = sorted(dirs)
  last = '!'
  dump = []
  for d in dirs:
    if d.startswith(last): continue
    dump.append(d)
    last = d

  ## We're done: return the filtered list.
  return dump

### The `DUMPERS' list consists of (PREDICATE, LISTERS) pairs.  The PREDICATE
### is a function of one argument -- a directory name -- which returns true
### if the LISTERS should be used to enumerate that directory.  The LISTERS
### are a list of functions of one argument -- again, the directory name --
### which should return an iterable of files within that directory, relative
### to its top-level.  Lister functions should not return the root directory,
### because it should obviously only be included once.  Instead, the root is
### handled separately by `dump_dir'.

def exists_subdir(subdir):
  """
  Predicate for `DUMPERS': match if the directory has a subdirectory SUBDIR.

  This is mainly useful for detecting working trees subject to version
  control.
  """
  return lambda dir: OS.path.isdir(OS.path.join(dir, subdir))

def filez(cmd):
  """
  Lister for `DUMPERS': generate the null-terminated items output by CMD.

  Run CMD, a string containing words with shell-like quoting (expected to be
  a literal in the code, so security concerns don't arise) in the directory
  of interest, yielding the invidual null-terminated strings which the
  command writes to its standard output.
  """
  def _(dir):

    ## Start the command,
    kid = SUB.Popen(SL.split(cmd), stdout = SUB.PIPE, cwd = dir)

    ## Collect and return the null-terminated items.  Strip off any leading
    ## `./' and exclude the root directory because that gets handled
    ## separately.
    left = ''
    while True:

      ## Read a new bufferload of stuff.  If there's nothing left then we're
      ## done.
      buf = kid.stdout.read(16384)
      if not buf: break

      ## Tack whatever was left over from last time on the front, and carve
      ## into null-terminated pieces.
      buf = left + buf
      i = 0
      while True:
        z = buf.find('\0', i)
        if z < 0: break
        f = buf[i:z]
        i = z + 1
        if f == '.': continue
        if f.startswith('./'): f = f[2:]
        yield f

      ## Whatever's left over will be dealt with next time through.
      left = buf[i:]

    ## If there's trailing junk left over then we should complain.
    if left:
      raise U.ExpectedError, \
            (500, "trailing junk from `%s' in `%s'" % (cmd, dir))

  ## Return the listing function.
  return _

## The list of predicates and listers.
DUMPERS = [
  (exists_subdir('.git'), [filez('git ls-files -coz --exclude-standard'),
                           filez('find .git -print0')]),
  (lambda d: True, [filez('find . ( ! -perm +004 -prune ) -o -print0')])]

###--------------------------------------------------------------------------
### Actually dumping files.

def dump_dir(name, dir, dirmap, tf, root):
  """
  Add the contents of directory DIR to the tarfile TF, under the given NAME.

  The ROOT names the toplevel of the tarball (we're not in the business of
  making tarbombs here).  The DIRMAP is a list of all of the (DIR, NAME)
  pairs being dumped, used for fixing up symbolic links between directories.
  """

  ## Find an appropriate `DUMPERS' list entry.
  for test, listers in DUMPERS:
    if test(dir): break
  else:
    raise U.ExpectedError, (500, "no dumper for `%s'" % dir)

  ## Write a tarfile entry for the toplevel.
  tf.add(dir, OS.path.join(root, name), recursive = False)

  ## Work through all of the listers.
  for lister in listers:

    ## Work through each file.
    for file in lister(dir):
      full = OS.path.join(dir, file)
      tarname = OS.path.join(root, name, file)
      skip = False

      ## Check for symbolic links.  If we find one that points to another
      ## directory we're going to dump separately then fiddle it so that it
      ## works in the directory tree we're going to make.
      if OS.path.islink(full):
        dest = OS.path.realpath(full)
        for d, local in dirmap:
          if dest.startswith(d):
            fix = OS.path.relpath(OS.path.join('/', local, dest[len(d):]),
                                  OS.path.join('/', name,
                                               OS.path.dirname(file)))
            st = OS.stat(full)
            ti = tf.gettarinfo(full, tarname)
            ti.linkname = fix
            tf.addfile(ti)
            skip = True

      ## Nothing special, so just dump the file.  Or whatever it is.
      if not skip:
        tf.add(full, tarname, recursive = False)

def source(out):
  """
  Write a tarball for the program's source code to OUT.

  This function automatically dumps all of the program's dependencies except
  for those covered by the operating-system exemption.
  """

  ## Make a tarfile writer.  There's an annoying incompatibility to bodge
  ## around.
  if SYS.version_info >= (2, 6):
    tf = TAR.open(fileobj = out, mode = 'w|gz', format = TAR.USTAR_FORMAT)
  else:
    tf = TAR.open(fileobj = out, mode = 'w|gz')
    tf.posix = True

  ## First of all, find out what needs to be dumped, and assign names to all
  ## of the various directories.
  root = '%s-%s' % (PACKAGE, VERSION)
  seen = set()
  dirmap = []
  festout = StringIO()
  for dir in dirs_to_dump():
    dir = dir.rstrip('/')
    base = OS.path.basename(dir)
    if base not in seen:
      name = base
    else:
      for i in I.count():
        name = '%s.%d' % (base, i)
        if name not in seen: break
    dirmap.append((dir + '/', name))
    festout.write('%s = %s\n' % (name, dir))

  ## Write a map of where things were in the filesystem.  This may help a
  ## user figure out how to deploy the thing.
  fest = festout.getvalue()
  ti = TAR.TarInfo(OS.path.join(root, 'MANIFEST'))
  ti.size = len(fest)
  ti.mtime = T.time()
  ti.mode = 0664
  ti.type = TAR.REGTYPE
  uid = OS.getuid(); ti.uid, ti.uname = uid, PW.getpwuid(uid).pw_name
  gid = OS.getgid(); ti.gid, ti.gname = gid, GR.getgrgid(gid).gr_name
  tf.addfile(ti, fileobj = StringIO(fest))

  ## Now actually dump all of the individual directories.
  for dir, name in dirmap:
    dump_dir(name, dir, dirmap, tf, root)

  ## We're done.
  tf.close()

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