peerdb/tripe-newpeers.in: Count enhancements correctly.

[tripe] / peerdb / tripe-newpeers.in

#! @PYTHON@
### -*-python-*-
###
### Build a CDB file from configuration file
###
### (c) 2007 Straylight/Edgeware
###

###----- Licensing notice ---------------------------------------------------
###
### This file is part of Trivial IP Encryption (TrIPE).
###
### TrIPE is free software: you can redistribute it and/or modify it under
### the terms of the GNU General Public License as published by the Free
### Software Foundation; either version 3 of the License, or (at your
### option) any later version.
###
### TrIPE 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 General Public License
### for more details.
###
### You should have received a copy of the GNU General Public License
### along with TrIPE.  If not, see <https://www.gnu.org/licenses/>.

VERSION = '@VERSION@'

###--------------------------------------------------------------------------
### External dependencies.

import ConfigParser as CP
import mLib as M
from optparse import OptionParser
import cdb as CDB
from sys import stdin, stdout, exit, argv
import re as RX
import os as OS

###--------------------------------------------------------------------------
### Utilities.

class CDBFake (object):
  """Like cdbmake, but just outputs data suitable for cdb-map."""
  def __init__(me, file = stdout):
    me.file = file
  def add(me, key, value):
    me.file.write('%s:%s\n' % (key, value))
  def finish(me):
    pass

###--------------------------------------------------------------------------
### A bulk DNS resolver.

class BulkResolver (object):
  """
  Resolve a number of DNS names in parallel.

  The BulkResovler resolves a number of hostnames in parallel.  Using it
  works in three phases:

    1. You call prepare(HOSTNAME) a number of times, to feed in the hostnames
       you're interested in.

    2. You call run() to actually drive the resolver.

    3. You call lookup(HOSTNAME) to get the address you wanted.  This will
       fail with KeyError if the resolver couldn't resolve the HOSTNAME.
  """

  def __init__(me):
    """Initialize the resolver."""
    me._resolvers = {}
    me._namemap = {}

  def prepare(me, host):
    """Prime the resolver to resolve the name HOST."""
    if host not in me._resolvers:
      me._resolvers[host] = M.SelResolveByName \
                            (host,
                             lambda name, alias, addr:
                               me._resolved(host, addr[0]),
                             lambda: me._resolved(host, None))

  def run(me):
    """Run the background DNS resolver until it's finished."""
    while me._resolvers:
      M.select()

  def lookup(me, host):
    """
    Fetch the address corresponding to HOST.
    """
    addr = me._namemap[host]
    if addr is None:
      raise KeyError(host)
    return addr

  def _resolved(me, host, addr):
    """Callback function: remember that ADDR is the address for HOST."""
    me._namemap[host] = addr
    del me._resolvers[host]

###--------------------------------------------------------------------------
### The configuration parser.

## Match a $(VAR) configuration variable reference; group 1 is the VAR.
r_ref = RX.compile(r'\$\(([^)]+)\)')

## Match a $[HOST] name resolution reference; group 1 is the HOST.
r_resolve = RX.compile(r'\$\[([^]]+)\]')

def _fmt_path(path):
  return ' -> '.join(["`%s'" % hop for hop in path])

class AmbiguousOptionError (Exception):
  def __init__(me, key, patha, vala, pathb, valb):
    me.key = key
    me.patha, me.vala = patha, vala
    me.pathb, me.valb = pathb, valb
  def __str__(me):
    return "Ambiguous answer resolving key `%s': " \
        "path %s yields `%s' but %s yields `%s'" % \
        (me.key, _fmt_path(me.patha), me.vala, _fmt_path(me.pathb), me.valb)

class InheritanceCycleError (Exception):
  def __init__(me, key, path):
    me.key = key
    me.path = path
  def __str__(me):
    return "Found a cycle %s looking up key `%s'" % \
        (_fmt_path(me.path), me.key)

class MissingKeyException (Exception):
  def __init__(me, sec, key):
    me.sec = sec
    me.key = key
  def __str__(me):
    return "Key `%s' not found in section `%s'" % (me.key, me.sec)

class MyConfigParser (CP.RawConfigParser):
  """
  A more advanced configuration parser.

  This has three major enhancements over the standard ConfigParser which are
  relevant to us.

    * It recognizes `@inherits' keys and follows them when expanding a
      value.

    * It recognizes `$(VAR)' references to configuration variables during
      expansion and processes them correctly.

    * It recognizes `$[HOST]' name-resolver requests and handles them
      correctly.

  Use:

    1. Call read(FILENAME) and/or read(FP, [FILENAME]) to slurp in the
       configuration data.

    2. Call resolve() to collect the hostnames which need to be resolved and
       actually do the name resolution.

    3. Call get(SECTION, ITEM) to collect the results, or items(SECTION) to
       iterate over them.
  """

  def __init__(me):
    """
    Initialize a new, empty configuration parser.
    """
    CP.RawConfigParser.__init__(me)
    me._resolver = BulkResolver()

  def resolve(me):
    """
    Works out all of the hostnames which need resolving and resolves them.

    Until you call this, attempts to fetch configuration items which need to
    resolve hostnames will fail!
    """
    for sec in me.sections():
      for key, value in me.items(sec, resolvep = False):
        for match in r_resolve.finditer(value):
          me._resolver.prepare(match.group(1))
    me._resolver.run()

  def _expand(me, sec, string, resolvep):
    """
    Expands $(...) and (optionally) $[...] placeholders in STRING.

    The SEC is the configuration section from which to satisfy $(...)
    requests.  RESOLVEP is a boolean switch: do we bother to tax the resolver
    or not?  This is turned off by the resolve() method while it's collecting
    hostnames to be resolved.
    """
    string = r_ref.sub \
             (lambda m: me.get(sec, m.group(1), resolvep), string)
    if resolvep:
      string = r_resolve.sub(lambda m: me._resolver.lookup(m.group(1)),
                             string)
    return string

  def has_option(me, sec, key):
    """
    Decide whether section SEC has a configuration key KEY.

    This version of the method properly handles the @inherit key.
    """
    return key == 'name' or me._get(sec, key)[0] is not None

  def _get(me, sec, key, map = None, path = None):
    """
    Low-level option-fetching method.

    Fetch the value for the named KEY from section SEC, or maybe
    (recursively) a section which SEC inherits from.

    Returns a pair VALUE, PATH.  The value is not expanded; nor do we check
    for the special `name' key.  The caller is expected to do these things.
    Returns None if no value could be found.
    """

    ## If we weren't given a memoization map or path, then we'd better make
    ## one.
    if map is None: map = {}
    if path is None: path = []

    ## If we've been this way before on another pass through then return the
    ## value we found then.  If we're still thinking about it then we've
    ## found a cycle.
    path.append(sec)
    try:
      threadp, value = map[sec]
    except KeyError:
      pass
    else:
      if threadp:
        raise InheritanceCycleError(key, path)

    ## See whether the answer is ready waiting for us.
    try:
      v = CP.RawConfigParser.get(me, sec, key)
    except CP.NoOptionError:
      pass
    else:
      p = path[:]
      path.pop()
      return v, p

    ## No, apparently, not.  Find out our list of parents.
    try:
      parents = CP.RawConfigParser.get(me, sec, '@inherit').\
          replace(',', ' ').split()
    except CP.NoOptionError:
      parents = []

    ## Initially we have no idea.
    value = None
    winner = None

    ## Go through our parents and ask them what they think.
    map[sec] = True, None
    for p in parents:

      ## See whether we get an answer.  If not, keep on going.
      v, pp = me._get(p, key, map, path)
      if v is None: continue

      ## If we got an answer, check that it matches any previous ones.
      if value is None:
        value = v
        winner = pp
      elif value != v:
        raise AmbiguousOptionError(key, winner, value, pp, v)

    ## That's the best we could manage.
    path.pop()
    map[sec] = False, value
    return value, winner

  def get(me, sec, key, resolvep = True):
    """
    Retrieve the value of KEY from section SEC.
    """

    ## Special handling for the `name' key.
    if key == 'name':
      try: value = CP.RawConfigParser.get(me, sec, key)
      except CP.NoOptionError: value = sec
    else:
      value, _ = me._get(sec, key)
      if value is None:
        raise MissingKeyException(sec, key)

    ## Expand the value and return it.
    return me._expand(sec, value, resolvep)

  def items(me, sec, resolvep = True):
    """
    Return a list of (NAME, VALUE) items in section SEC.

    This extends the default method by handling the inheritance chain.
    """

    ## Initialize for a depth-first walk of the inheritance graph.
    d = {}
    visited = {}
    basesec = sec
    stack = [sec]

    ## Visit nodes, collecting their keys.  Don't believe the values:
    ## resolving inheritance is too hard to do like this.
    while stack:
      sec = stack.pop()
      if sec in visited: continue
      visited[sec] = True

      for key, value in CP.RawConfigParser.items(me, sec):
        if key == '@inherit': stack += value.replace(',', ' ').split()
        else: d[key] = None

    ## Now collect the values for the known keys, one by one.
    items = []
    for key in d: items.append((key, me.get(basesec, key, resolvep)))

    ## And we're done.
    return items

###--------------------------------------------------------------------------
### Command-line handling.

def inputiter(things):
  """
  Iterate over command-line arguments, returning corresponding open files.

  If none were given, or one is `-', assume standard input; if one is a
  directory, scan it for files other than backups; otherwise return the
  opened files.
  """

  if not things:
    if OS.isatty(stdin.fileno()):
      M.die('no input given, and stdin is a terminal')
    yield stdin
  else:
    for thing in things:
      if thing == '-':
        yield stdin
      elif OS.path.isdir(thing):
        for item in OS.listdir(thing):
          if item.endswith('~') or item.endswith('#'):
            continue
          name = OS.path.join(thing, item)
          if not OS.path.isfile(name):
            continue
          yield file(name)
      else:
        yield file(thing)

def parse_options(argv = argv):
  """
  Parse command-line options, returning a pair (OPTS, ARGS).
  """
  M.ego(argv[0])
  op = OptionParser(usage = '%prog [-c CDB] INPUT...',
                    version = '%%prog (tripe, version %s)' % VERSION)
  op.add_option('-c', '--cdb', metavar = 'CDB',
                dest = 'cdbfile', default = None,
                help = 'Compile output into a CDB file.')
  opts, args = op.parse_args(argv)
  return opts, args

###--------------------------------------------------------------------------
### Main code.

def getconf(args):
  """
  Read the configuration files and return the accumulated result.

  We make sure that all hostnames have been properly resolved.
  """
  conf = MyConfigParser()
  for f in inputiter(args):
    conf.readfp(f)
  conf.resolve()
  return conf

def output(conf, cdb):
  """
  Output the configuration information CONF to the database CDB.

  This is where the special `user' and `auto' database entries get set.
  """
  auto = []
  for sec in sorted(conf.sections()):
    if sec.startswith('@'):
      continue
    elif sec.startswith('$'):
      label = sec
    else:
      label = 'P%s' % sec
      if conf.has_option(sec, 'auto') and \
         conf.get(sec, 'auto') in ('y', 'yes', 't', 'true', '1', 'on'):
        auto.append(sec)
      if conf.has_option(sec, 'user'):
        cdb.add('U%s' % conf.get(sec, 'user'), sec)
    url = M.URLEncode(laxp = True, semip = True)
    for key, value in sorted(conf.items(sec), key = lambda (k, v): k):
      if not key.startswith('@'):
        url.encode(key, ' '.join(M.split(value)[0]))
    cdb.add(label, url.result)
  cdb.add('%AUTO', ' '.join(auto))
  cdb.finish()

def main():
  """Main program."""
  opts, args = parse_options()
  if opts.cdbfile:
    cdb = CDB.cdbmake(opts.cdbfile, opts.cdbfile + '.new')
  else:
    cdb = CDBFake()
  conf = getconf(args[1:])
  output(conf, cdb)

if __name__ == '__main__':
  main()

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