svc/conntrack: Make the kickpeers coroutine more robust.

[tripe] / svc / conntrack.in

#! @PYTHON@
### -*-python-*-
###
### Service for automatically tracking network connection status
###
### (c) 2010 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 2 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, write to the Free Software Foundation,
### Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

VERSION = '@VERSION@'

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

from ConfigParser import RawConfigParser
from optparse import OptionParser
import os as OS
import sys as SYS
import socket as S
import mLib as M
import tripe as T
import dbus as D
for i in ['mainloop', 'mainloop.glib']:
  __import__('dbus.%s' % i)
import gobject as G
from struct import pack, unpack

SM = T.svcmgr
##__import__('rmcr').__debug = True

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

class struct (object):
  """A simple container object."""
  def __init__(me, **kw):
    me.__dict__.update(kw)

def toposort(cmp, things):
  """
  Generate the THINGS in an order consistent with a given partial order.

  The function CMP(X, Y) should return true if X must precede Y, and false if
  it doesn't care.  If X and Y are equal then it should return false.

  The THINGS may be any finite iterable; it is converted to a list
  internally.
  """

  ## Make sure we can index the THINGS, and prepare an ordering table.
  ## What's going on?  The THINGS might not have a helpful equality
  ## predicate, so it's easier to work with indices.  The ordering table will
  ## remember which THINGS (by index) are considered greater than other
  ## things.
  things = list(things)
  n = len(things)
  order = [{} for i in xrange(n)]
  rorder = [{} for i in xrange(n)]
  for i in xrange(n):
    for j in xrange(n):
      if i != j and cmp(things[i], things[j]):
        order[j][i] = True
        rorder[i][j] = True

  ## Now we can do the sort.
  out = []
  while True:
    done = True
    for i in xrange(n):
      if order[i] is not None:
        done = False
        if len(order[i]) == 0:
          for j in rorder[i]:
            del order[j][i]
          yield things[i]
          order[i] = None
    if done:
      break

###--------------------------------------------------------------------------
### Parse the configuration file.

## Hmm.  Should I try to integrate this with the peers database?  It's not a
## good fit; it'd need special hacks in tripe-newpeers.  And the use case for
## this service are largely going to be satellite notes, I don't think
## scalability's going to be a problem.

class Config (object):
  """
  Represents a configuration file.

  The most interesting thing is probably the `groups' slot, which stores a
  list of pairs (NAME, PATTERNS); the NAME is a string, and the PATTERNS a
  list of (TAG, PEER, ADDR, MASK) triples.  The implication is that there
  should be precisely one peer with a name matching NAME-*, and that it
  should be NAME-TAG, where (TAG, PEER, ADDR, MASK) is the first triple such
  that the host's primary IP address (if PEER is None -- or the IP address it
  would use for communicating with PEER) is within the network defined by
  ADDR/MASK.
  """

  def __init__(me, file):
    """
    Construct a new Config object, reading the given FILE.
    """
    me._file = file
    me._fwatch = M.FWatch(file)
    me._update()

  def check(me):
    """
    See whether the configuration file has been updated.
    """
    if me._fwatch.update():
      me._update()

  def _update(me):
    """
    Internal function to update the configuration from the underlying file.
    """

    ## Read the configuration.  We have no need of the fancy substitutions,
    ## so turn them all off.
    cp = RawConfigParser()
    cp.read(me._file)

    ## Save the test address.  Make sure it's vaguely sensible.  The default
    ## is probably good for most cases, in fact, since that address isn't
    ## actually in use.  Note that we never send packets to the test address;
    ## we just use it to discover routing information.
    if cp.has_option('DEFAULT', 'test-addr'):
      testaddr = cp.get('DEFAULT', 'test-addr')
      S.inet_aton(testaddr)
    else:
      testaddr = '1.2.3.4'

    ## Scan the configuration file and build the groups structure.
    groups = []
    for sec in cp.sections():
      pats = []
      for tag in cp.options(sec):
        spec = cp.get(sec, tag).split()

        ## Parse the entry into peer and network.
        if len(spec) == 1:
          peer = None
          net = spec[0]
        else:
          peer, net = spec

        ## Syntax of a net is ADDRESS/MASK, where ADDRESS is a dotted-quad,
        ## and MASK is either a dotted-quad or a single integer N indicating
        ## a mask with N leading ones followed by trailing zeroes.
        slash = net.index('/')
        addr, = unpack('>L', S.inet_aton(net[:slash]))
        if net.find('.', slash + 1) >= 0:
          mask, = unpack('>L', S.inet_aton(net[:slash]))
        else:
          n = int(net[slash + 1:], 10)
          mask = (1 << 32) - (1 << 32 - n)
        pats.append((tag, peer, addr & mask, mask))

      ## Annoyingly, RawConfigParser doesn't preserve the order of options.
      ## In order to make things vaguely sane, we topologically sort the
      ## patterns so that more specific patterns are checked first.
      pats = list(toposort(lambda (t, p, a, m), (tt, pp, aa, mm): \
                             (p and not pp) or \
                             (p == pp and m == (m | mm) and aa == (a & mm)),
                           pats))
      groups.append((sec, pats))

    ## Done.
    me.testaddr = testaddr
    me.groups = groups

### This will be a configuration file.
CF = None

###--------------------------------------------------------------------------
### Responding to a network up/down event.

def localaddr(peer):
  """
  Return the local IP address used for talking to PEER.
  """
  sk = S.socket(S.AF_INET, S.SOCK_DGRAM)
  try:
    try:
      sk.connect((peer, 1))
      addr, _ = sk.getsockname()
      addr, = unpack('>L', S.inet_aton(addr))
      return addr
    except S.error:
      return None
  finally:
    sk.close()

_kick = T.Queue()
def kickpeers():
  lastip = {}
  while True:
    upness, reason = _kick.get()

    ## Make sure the configuration file is up-to-date.  Don't worry if we
    ## can't do anything useful.
    try:
      CF.check()
    except Exception, exc:
      SM.warn('conntrack', 'config-file-error',
              exc.__class__.__name__, str(exc))

    ## Find the current list of peers.
    peers = SM.list()

    ## Work out the primary IP address.
    if upness:
      addr = localaddr(CF.testaddr)
      if addr is None:
        upness = False
    else:
      addr = None

    ## Now decide what to do.
    changes = []
    for g, pp in CF.groups:

      ## Find out which peer in the group ought to be active.
      ip = None
      map = {}
      want = None
      for t, p, a, m in pp:
        if p is None or not upness:
          ipq = addr
        else:
          ipq = localaddr(p)
        if upness and ip is None and \
              ipq is not None and (ipq & m) == a:
          map[t] = 'up'
          want = t
          ip = ipq
        else:
          map[t] = 'down'

      ## Shut down the wrong ones.
      found = False
      for p in peers:
        what = map.get(p, 'leave')
        if what == 'up':
          found = True
        elif what == 'down':
          def _(p = p):
            try:
              SM.kill(p)
            except T.TripeError, exc:
              if exc.args[0] == 'unknown-peer':
                ## Inherently racy; don't worry about this.
                pass
              else:
                raise
          changes.append(_)

      ## Start the right one if necessary.
      if want is not None and (not found or ip != lastip.get(g, None)):
        def _(want = want):
          try:
            SM.svcsubmit('connect', 'active', want)
          except T.TripeError, exc:
            SM.warn('conntrack', 'connect-failed', want, *exc.args)
        changes.append(_)
      lastip[g] = ip

    ## Commit the changes.
    if changes:
      SM.notify('conntrack', upness and 'up' or 'down', *reason)
      for c in changes: c()

def netupdown(upness, reason):
  """
  Add or kill peers according to whether the network is up or down.

  UPNESS is true if the network is up, or false if it's down.
  """

  _kick.put((upness, reason))

###--------------------------------------------------------------------------
### NetworkManager monitor.

NM_NAME = 'org.freedesktop.NetworkManager'
NM_PATH = '/org/freedesktop/NetworkManager'
NM_IFACE = NM_NAME
NMCA_IFACE = NM_NAME + '.Connection.Active'

NM_STATE_CONNECTED = 3

class NetworkManagerMonitor (object):
  """
  Watch NetworkManager signals for changes in network state.
  """

  ## Strategy.  There are two kinds of interesting state transitions for us.
  ## The first one is the global are-we-connected state, which we'll use to
  ## toggle network upness on a global level.  The second is which connection
  ## has the default route, which we'll use to tweak which peer in the peer
  ## group is active.  The former is most easily tracked using the signal
  ## org.freedesktop.NetworkManager.StateChanged; for the latter, we track
  ## org.freedesktop.NetworkManager.Connection.Active.PropertiesChanged and
  ## look for when a new connection gains the default route.

  def attach(me, bus):
    try:
      nm = bus.get_object(NM_NAME, NM_PATH)
      state = nm.Get(NM_IFACE, 'State')
      if state == NM_STATE_CONNECTED:
        netupdown(True, ['nm', 'initially-connected'])
      else:
        netupdown(False, ['nm', 'initially-disconnected'])
    except D.DBusException:
      pass
    bus.add_signal_receiver(me._nm_state, 'StateChanged', NM_IFACE,
                            NM_NAME, NM_PATH)
    bus.add_signal_receiver(me._nm_connchange,
                            'PropertiesChanged', NMCA_IFACE,
                            NM_NAME, None)

  def _nm_state(me, state):
    if state == NM_STATE_CONNECTED:
      netupdown(True, ['nm', 'connected'])
    else:
      netupdown(False, ['nm', 'disconnected'])

  def _nm_connchange(me, props):
    if props.get('Default', False):
      netupdown(True, ['nm', 'default-connection-change'])

###--------------------------------------------------------------------------
### Maemo monitor.

ICD_NAME = 'com.nokia.icd'
ICD_PATH = '/com/nokia/icd'
ICD_IFACE = ICD_NAME

class MaemoICdMonitor (object):
  """
  Watch ICd signals for changes in network state.
  """

  ## Strategy.  ICd only handles one connection at a time in steady state,
  ## though when switching between connections, it tries to bring the new one
  ## up before shutting down the old one.  This makes life a bit easier than
  ## it is with NetworkManager.  On the other hand, the notifications are
  ## relative to particular connections only, and the indicator that the old
  ## connection is down (`IDLE') comes /after/ the new one comes up
  ## (`CONNECTED'), so we have to remember which one is active.

  def attach(me, bus):
    try:
      icd = bus.get_object(ICD_NAME, ICD_PATH)
      try:
        iap = icd.get_ipinfo(dbus_interface = ICD_IFACE)[0]
        me._iap = iap
        netupdown(True, ['icd', 'initially-connected', iap])
      except D.DBusException:
        me._iap = None
        netupdown(False, ['icd', 'initially-disconnected'])
    except D.DBusException:
      me._iap = None
    bus.add_signal_receiver(me._icd_state, 'status_changed', ICD_IFACE,
                            ICD_NAME, ICD_PATH)

  def _icd_state(me, iap, ty, state, hunoz):
    if state == 'CONNECTED':
      me._iap = iap
      netupdown(True, ['icd', 'connected', iap])
    elif state == 'IDLE' and iap == me._iap:
      me._iap = None
      netupdown(False, ['icd', 'idle'])

###--------------------------------------------------------------------------
### D-Bus connection tracking.

class DBusMonitor (object):
  """
  Maintains a connection to the system D-Bus, and watches for signals.

  If the connection is initially down, or drops for some reason, we retry
  periodically (every five seconds at the moment).  If the connection
  resurfaces, we reattach the monitors.
  """

  def __init__(me):
    """
    Initialise the object and try to establish a connection to the bus.
    """
    me._mons = []
    me._loop = D.mainloop.glib.DBusGMainLoop()
    me._state = 'startup'
    me._reconnect()

  def addmon(me, mon):
    """
    Add a monitor object to watch for signals.

    MON.attach(BUS) is called, with BUS being the connection to the system
    bus.  MON should query its service's current status and watch for
    relevant signals.
    """
    me._mons.append(mon)
    if me._bus is not None:
      mon.attach(me._bus)

  def _reconnect(me, hunoz = None):
    """
    Start connecting to the bus.

    If we fail the first time, retry periodically.
    """
    if me._state == 'startup':
      T.aside(SM.notify, 'conntrack', 'dbus-connection', 'startup')
    elif me._state == 'connected':
      T.aside(SM.notify, 'conntrack', 'dbus-connection', 'lost')
    else:
      T.aside(SM.notify, 'conntrack', 'dbus-connection',
              'state=%s' % me._state)
    me._state == 'reconnecting'
    me._bus = None
    if me._try_connect():
      G.timeout_add_seconds(5, me._try_connect)

  def _try_connect(me):
    """
    Actually make a connection attempt.

    If we succeed, attach the monitors.
    """
    try:
      addr = OS.getenv('TRIPE_CONNTRACK_BUS')
      if addr == 'SESSION':
        bus = D.SessionBus(mainloop = me._loop, private = True)
      elif addr is not None:
        bus = D.bus.BusConnection(addr, mainloop = me._loop)
      else:
        bus = D.SystemBus(mainloop = me._loop, private = True)
      for m in me._mons:
        m.attach(bus)
    except D.DBusException, e:
      return True
    me._bus = bus
    me._state = 'connected'
    bus.call_on_disconnection(me._reconnect)
    T.aside(SM.notify, 'conntrack', 'dbus-connection', 'connected')
    return False

###--------------------------------------------------------------------------
### TrIPE service.

class GIOWatcher (object):
  """
  Monitor I/O events using glib.
  """
  def __init__(me, conn, mc = G.main_context_default()):
    me._conn = conn
    me._watch = None
    me._mc = mc
  def connected(me, sock):
    me._watch = G.io_add_watch(sock, G.IO_IN,
                               lambda *hunoz: me._conn.receive())
  def disconnected(me):
    G.source_remove(me._watch)
    me._watch = None
  def iterate(me):
    me._mc.iteration(True)

SM.iowatch = GIOWatcher(SM)

def init():
  """
  Service initialization.

  Add the D-Bus monitor here, because we might send commands off immediately,
  and we want to make sure the server connection is up.
  """
  global DBM
  T.Coroutine(kickpeers, name = 'kickpeers').switch()
  DBM = DBusMonitor()
  DBM.addmon(NetworkManagerMonitor())
  DBM.addmon(MaemoICdMonitor())
  G.timeout_add_seconds(30, lambda: (netupdown(True, ['interval-timer'])
                                     or True))

def parse_options():
  """
  Parse the command-line options.

  Automatically changes directory to the requested configdir, and turns on
  debugging.  Returns the options object.
  """
  op = OptionParser(usage = '%prog [-a FILE] [-d DIR]',
                    version = '%%prog %s' % VERSION)

  op.add_option('-a', '--admin-socket',
                metavar = 'FILE', dest = 'tripesock', default = T.tripesock,
                help = 'Select socket to connect to [default %default]')
  op.add_option('-d', '--directory',
                metavar = 'DIR', dest = 'dir', default = T.configdir,
                help = 'Select current diretory [default %default]')
  op.add_option('-c', '--config',
                metavar = 'FILE', dest = 'conf', default = 'conntrack.conf',
                help = 'Select configuration [default %default]')
  op.add_option('--daemon', dest = 'daemon',
                default = False, action = 'store_true',
                help = 'Become a daemon after successful initialization')
  op.add_option('--debug', dest = 'debug',
                default = False, action = 'store_true',
                help = 'Emit debugging trace information')
  op.add_option('--startup', dest = 'startup',
                default = False, action = 'store_true',
                help = 'Being called as part of the server startup')

  opts, args = op.parse_args()
  if args: op.error('no arguments permitted')
  OS.chdir(opts.dir)
  T._debug = opts.debug
  return opts

## Service table, for running manually.
def cmd_updown(upness):
  return lambda *args: T.defer(netupdown, upness, ['manual'] + list(args))
service_info = [('conntrack', VERSION, {
  'up': (0, None, '', cmd_updown(True)),
  'down': (0, None, '', cmd_updown(False))
})]

if __name__ == '__main__':
  opts = parse_options()
  CF = Config(opts.conf)
  T.runservices(opts.tripesock, service_info,
                init = init, daemon = opts.daemon)

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