svc/conntrack.in: Maintain config groups in a dictionary.

[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 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.

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)
try: from gi.repository import GLib as G
except ImportError: 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

###--------------------------------------------------------------------------
### Address manipulation.

class InetAddress (object):
  def __init__(me, addrstr, maskstr = None):
    me.addr = me._addrstr_to_int(addrstr)
    if maskstr is None:
      me.mask = -1
    elif maskstr.isdigit():
      me.mask = (1 << 32) - (1 << 32 - int(maskstr))
    else:
      me.mask = me._addrstr_to_int(maskstr)
    if me.addr&~me.mask:
      raise ValueError('network contains bits set beyond mask')
  def _addrstr_to_int(me, addrstr):
    return unpack('>L', S.inet_aton(addrstr))[0]
  def _int_to_addrstr(me, n):
    return S.inet_ntoa(pack('>L', n))
  def sockaddr(me, port = 0):
    if me.mask != -1: raise ValueError('not a simple address')
    return me._int_to_addrstr(me.addr), port
  def __str__(me):
    addrstr = me._int_to_addrstr(me.addr)
    if me.mask == -1:
      return addrstr
    else:
      inv = me.mask ^ ((1 << 32) - 1)
      if (inv&(inv + 1)) == 0:
        return '%s/%d' % (addrstr, 32 - inv.bit_length())
      else:
        return '%s/%s' % (addrstr, me._int_to_addrstr(me.mask))
  def withinp(me, net):
    if (me.mask&net.mask) != net.mask: return False
    if (me.addr ^ net.addr)&net.mask: return False
    return True
  def eq(me, other):
    if me.mask != other.mask: return False
    if me.addr != other.addr: return False
    return True
  @classmethod
  def from_sockaddr(cls, sa):
    addr, port = (lambda a, p: (a, p))(*sa)
    return cls(addr), port

def parse_address(addrstr, maskstr = None):
  return InetAddress(addrstr, maskstr)

def parse_net(netstr):
  try: sl = netstr.index('/')
  except ValueError: raise ValueError('missing mask')
  return parse_address(netstr[:sl], netstr[sl + 1:])

def straddr(a): return a is None and '#<none>' or str(a)

###--------------------------------------------------------------------------
### 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.

TESTADDR = InetAddress('1.2.3.4')

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, NET) triples.  The implication is that there should be
  precisely one peer from the set, and that it should be named TAG, where
  (TAG, PEER, NET) 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 NET.
  """

  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)
    if T._debug: print '# reread config'

    ## 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 = InetAddress(cp.get('DEFAULT', 'test-addr'))
    else:
      testaddr = TESTADDR

    ## 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 = InetAddress(spec[0])
          net = spec[1]

        ## 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.
        net = parse_net(net)
        pats.append((tag, peer, net))

      ## 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, n), (tt, pp, nn): \
                             (p and not pp) or \
                             (p == pp and n.withinp(nn)),
                           pats))
      groups[sec] = pats

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

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

def cmd_showconfig():
  T.svcinfo('test-addr=%s' % CF.testaddr)
def cmd_showgroups():
  for g in sorted(CF.groups.iterkeys()):
    T.svcinfo(g)
def cmd_showgroup(g):
  try: pats = CF.groups[g]
  except KeyError: raise T.TripeJobError('unknown-group', g)
  for t, p, n in pats:
    T.svcinfo('peer', t,
              'target', p and str(p) or '(default)',
              'net', str(n))

###--------------------------------------------------------------------------
### 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.sockaddr(1))
      addr = sk.getsockname()
      return InetAddress.from_sockaddr(addr)[0]
    except S.error:
      return None
  finally:
    sk.close()

_kick = T.Queue()
_delay = None

def cancel_delay():
  global _delay
  if _delay is not None:
    if T._debug: print '# cancel delayed kick'
    G.source_remove(_delay)
    _delay = None

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))

def delay_netupdown(upness, reason):
  global _delay
  cancel_delay()
  def _func():
    global _delay
    if T._debug: print '# delayed %s: %s' % (upness, reason)
    _delay = None
    netupdown(upness, reason)
    return False
  if T._debug: print '# delaying %s: %s' % (upness, reason)
  _delay = G.timeout_add(2000, _func)

def kickpeers():
  while True:
    upness, reason = _kick.get()
    if T._debug: print '# kickpeers %s: %s' % (upness, reason)
    select = []
    cancel_delay()

    ## 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
    if not T._debug: pass
    elif addr: print '#   local address = %s' % straddr(addr)
    else: print '#   offline'

    ## Now decide what to do.
    changes = []
    for g, pp in CF.groups.iteritems():
      if T._debug: print '#   check group %s' % g

      ## Find out which peer in the group ought to be active.
      ip = None
      map = {}
      want = None
      for t, p, n in pp:
        if p is None or not upness:
          ipq = addr
        else:
          ipq = localaddr(p)
        if T._debug:
          info = 'peer=%s; target=%s; net=%s; local=%s' % (
            t, p or '(default)', n, straddr(ipq))
        if upness and ip is None and \
              ipq is not None and ipq.withinp(n):
          if T._debug: print '#     %s: SELECTED' % info
          map[t] = 'up'
          select.append('%s=%s' % (g, t))
          if t == 'down' or t.startswith('down/'):
            want = None
          else:
            want = t
          ip = ipq
        else:
          map[t] = 'down'
          if T._debug: print '#     %s: skipped' % info

      ## Shut down the wrong ones.
      found = False
      if T._debug: print '#   peer-map = %r' % map
      for p in peers:
        what = map.get(p, 'leave')
        if what == 'up':
          found = True
          if T._debug: print '#   peer %s: already up' % p
        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
          if T._debug: print '#   peer %s: bring down' % p
          changes.append(_)

      ## Start the right one if necessary.
      if want is not None and not found:
        def _(want = want):
          try:
            list(SM.svcsubmit('connect', 'active', want))
          except T.TripeError, exc:
            SM.warn('conntrack', 'connect-failed', want, *exc.args)
        if T._debug: print '#   peer %s: bring up' % want
        changes.append(_)

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

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

DBPROPS_IFACE = 'org.freedesktop.DBus.Properties'

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

NM_STATE_CONNECTED = 3 #obsolete
NM_STATE_CONNECTED_LOCAL = 50
NM_STATE_CONNECTED_SITE = 60
NM_STATE_CONNECTED_GLOBAL = 70
NM_CONNSTATES = set([NM_STATE_CONNECTED,
                     NM_STATE_CONNECTED_LOCAL,
                     NM_STATE_CONNECTED_SITE,
                     NM_STATE_CONNECTED_GLOBAL])

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', dbus_interface = DBPROPS_IFACE)
      if state in NM_CONNSTATES:
        netupdown(True, ['nm', 'initially-connected'])
      else:
        netupdown(False, ['nm', 'initially-disconnected'])
    except D.DBusException, e:
      if T._debug: print '# exception attaching to network-manager: %s' % e
    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 in NM_CONNSTATES:
      delay_netupdown(True, ['nm', 'connected'])
    else:
      delay_netupdown(False, ['nm', 'disconnected'])

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

##--------------------------------------------------------------------------
### Connman monitor.

CM_NAME = 'net.connman'
CM_PATH = '/'
CM_IFACE = 'net.connman.Manager'

class ConnManMonitor (object):
  """
  Watch ConnMan signls for changes in network state.
  """

  ## Strategy.  Everything seems to be usefully encoded in the `State'
  ## property.  If it's `offline', `idle' or `ready' then we don't expect a
  ## network connection.  During handover from one network to another, the
  ## property passes through `ready' to `online'.

  def attach(me, bus):
    try:
      cm = bus.get_object(CM_NAME, CM_PATH)
      props = cm.GetProperties(dbus_interface = CM_IFACE)
      state = props['State']
      netupdown(state == 'online', ['connman', 'initially-%s' % state])
    except D.DBusException, e:
      if T._debug: print '# exception attaching to connman: %s' % e
    bus.add_signal_receiver(me._cm_state, 'PropertyChanged',
                            CM_IFACE, CM_NAME, CM_PATH)

  def _cm_state(me, prop, value):
    if prop != 'State': return
    delay_netupdown(value == 'online', ['connman', value])

###--------------------------------------------------------------------------
### 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, e:
      if T._debug: print '# exception attaching to icd: %s' % e
      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
      delay_netupdown(True, ['icd', 'connected', iap])
    elif state == 'IDLE' and iap == me._iap:
      me._iap = None
      delay_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(ConnManMonitor())
  DBM.addmon(MaemoICdMonitor())
  G.timeout_add_seconds(30, lambda: (_delay is not None or
                                     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)),
  'show-config': (0, 0, '', cmd_showconfig),
  'show-groups': (0, 0, '', cmd_showgroups),
  'show-group': (1, 1, 'GROUP', cmd_showgroup)
})]

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 --------------------------------------------------