svc/conntrack.in (kickpeers): Refactor and reformat the search loop.

[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
import re as RX
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)

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

CONFSYNTAX = [
  ('COMMENT', RX.compile(r'^\s*($|[;#])')),
  ('GRPHDR', RX.compile(r'^\s*\[(.*)\]\s*$')),
  ('ASSGN', RX.compile(r'\s*([\w.-]+)\s*[:=]\s*(|\S|\S.*\S)\s*$'))]

class ConfigError (Exception):
  def __init__(me, file, lno, msg):
    me.file = file
    me.lno = lno
    me.msg = msg
  def __str__(me):
    return '%s:%d: %s' % (me.file, me.lno, me.msg)

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

    if T._debug: print '# reread config'

    ## Initial state.
    testaddr = None
    groups = {}
    grpname = None
    grplist = []

    ## Open the file and start reading.
    with open(me._file) as f:
      lno = 0
      for line in f:
        lno += 1
        for tag, rx in CONFSYNTAX:
          m = rx.match(line)
          if m: break
        else:
          raise ConfigError(me._file, lno, 'failed to parse line: %r' % line)

        if tag == 'COMMENT':
          ## A comment.  Ignore it and hope it goes away.

          continue

        elif tag == 'GRPHDR':
          ## A group header.  Flush the old group and start a new one.
          newname = m.group(1)

          if grpname is not None: groups[grpname] = grplist
          if newname in groups:
            raise ConfigError(me._file, lno,
                              "duplicate group name `%s'" % newname)
          grpname = newname
          grplist = []

        elif tag == 'ASSGN':
           ## An assignment.  Deal with it.
          name, value = m.group(1), m.group(2)

          if grpname is None:
            ## We're outside of any group, so this is a global configuration
            ## tweak.

            if name == 'test-addr':
              for astr in value.split():
                try:
                  a = parse_address(astr)
                except Exception, e:
                  raise ConfigError(me._file, lno,
                                    "invalid IP address `%s': %s" %
                                    (astr, e))
                if testaddr is not None:
                  raise ConfigError(me._file, lno, 'duplicate test-address')
                testaddr = a
            else:
              raise ConfigError(me._file, lno,
                                "unknown global option `%s'" % name)

          else:
            ## Parse a pattern and add it to the group.
            spec = value.split()
            i = 0

            ## Check for an explicit target address.
            if i >= len(spec) or spec[i].find('/') >= 0:
              peer = None
            else:
              try:
                peer = parse_address(spec[i])
              except Exception, e:
                raise ConfigError(me._file, lno,
                                  "invalid IP address `%s': %s" %
                                  (spec[i], e))
              i += 1

            ## Parse the local network.
            if len(spec) != i + 1:
              raise ConfigError(me._file, lno, 'no network defined')
            try:
              net = parse_net(spec[i])
            except Exception, e:
              raise ConfigError(me._file, lno,
                                "invalid IP network `%s': %s" %
                                (spec[i], e))

            ## Add this entry to the list.
            grplist.append((name, peer, net))

    ## Fill in the default test address if necessary.
    if testaddr is None: testaddr = TESTADDR

    ## Done.
    if grpname is not None: groups[grpname] = grplist
    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.
      statemap = {}
      want = None
      matchp = False
      for t, p, n in pp:
        if p is None or not upness: ip = addr
        else: ip = localaddr(p)
        if T._debug:
          info = 'peer=%s; target=%s; net=%s; local=%s' % (
            t, p or '(default)', n, straddr(ip))
        if upness and not matchp and \
              ip is not None and ip.withinp(n):
          if T._debug: print '#     %s: SELECTED' % info
          statemap[t] = 'up'
          select.append('%s=%s' % (g, t))
          if t == 'down' or t.startswith('down/'): want = None
          else: want = t
          matchp = True
        else:
          statemap[t] = 'down'
          if T._debug: print '#     %s: skipped' % info

      ## Shut down the wrong ones.
      found = False
      if T._debug: print '#   peer-map = %r' % statemap
      for p in peers:
        what = statemap.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 --------------------------------------------------