#!/usr/bin/env python

# Daemon that listens to updates about the source IP prefixes from which snmp access
# is allowed. In case of change, it will update the snmp configuration file accordingly. 
# Also, after a change, it will notify snmpd to re-read its config file (service reload).

import os
import re
import sys
import time
import redis

service="snmpd"
config_file_path="/etc/snmp"
redis_key="SNMP_ALLOW_LIST"   # the redis list we listen to
subscription='__keyspace@0__:%s' % redis_key
temporization_duration = 3 # how long we wait for changes to settle (ride out a bursts of changes in redis_key)
fake_infinite = 9999 # How often we wake up when nothing is going on --get_message()'s timeout has no 'infinite' value
# after these operations we may need to revisit existing ssh connections because they removed or modified existing entries
delete_operations = ["lrem", "lpop", "rpop", "blpop", "brpop", "brpoplpush", "rpoplpush", "ltrim", "del", "lset"]

r = redis.StrictRedis(host='localhost')
p = r.pubsub()

# If redis is not up yet, this can fail, so wait for redis to be available
while True:
   try:
      p.subscribe(subscription)
      break
   except redis.exceptions.ConnectionError:
     time.sleep(3)

# We could loose contact with redis at a later stage, in which case we will exit with
# return code -2 and supervisor will restart us, at which point we are back in the
# while loop above waiting for redis to be ready.
try:

  # By default redis does enable events, so enable them
  r.config_set("notify-keyspace-events", "KAE")
  

  # To update the configuration file 
  #
  # Example config file for reference:
  #   root@sonic:/# cat /etc/snmp/snmpd.conf
  #   <...some snmp config, like udp port to use etc...>
  #   rocommunity public 172.20.61.0/24
  #   rocommunity public 172.20.60.0/24
  #   rocommunity public 127.00.00.0/8
  #   <...some more snmp config...>
  #   root@sonic:/#
  #
  # snmpd.conf supports include file, like so:
  #  includeFile /etc/snmp/community.conf
  #  includeDir /etc/snmp/config.d
  # which could make file massaging simpler, but even then we still deal with lines 
  # that have shared "masters", since some other entity controls the community strings 
  # part of that line.
  # If other database attributes need to be written to the snmp config file, then
  # it should be done by this daemon as well (sure, we could inotify on the file
  # and correct it back, but that's glitchy).
  
  def write_configuration_file(v):
    filename="%s/%s.conf" % (config_file_path, service)
    filename_tmp = filename + ".tmp"
    f=open(filename, "r")
    snmpd_config = f.read()
    f.close()
    f=open(filename_tmp, "w")
    this_community = "not_a_community"
    for l in snmpd_config.split('\n'):
      m = re.match("^(..)community (\S+)", l)
      if not m:
          f.write(l)
          f.write("\n")
      else:
        if not l.startswith(this_community): # already handled community (each community is duplicated per allow entry)
           this_community="%scommunity %s" % (m.group(1), m.group(2))   
           if len(v):
             for value in v:
               f.write("%s %s\n" % (this_community, value))
           else:
             f.write("%s\n" % this_community)
    f.close()
    os.rename(filename_tmp, filename)
    os.system("kill -HUP $(pgrep snmpd) > /dev/null 2> /dev/null || :")

  # write initial configuration
  write_configuration_file(r.lrange(redis_key, 0, -1))
  
  # listen for changes and rewrite configuration file if needed, after some temporization
  #
  # How those subscribed to messages look like, for reference:
  #   {'pattern': None, 'type': 'subscribe', 'channel': '__keyspace@0__:SNMP_PERMIT_LIST', 'data': 1L}
  #   {'pattern': None, 'type': 'message',   'channel': '__keyspace@0__:SNMP_PERMIT_LIST', 'data': 'rpush'}
  #   {'pattern': None, 'type': 'message',   'channel': '__keyspace@0__:SNMP_PERMIT_LIST', 'data': 'lpush'}
  #   {'pattern': None, 'type': 'message',   'channel': '__keyspace@0__:SNMP_PERMIT_LIST', 'data': 'lrem'}
  #   {'pattern': None, 'type': 'message',   'channel': '__keyspace@0__:SNMP_PERMIT_LIST', 'data': 'lset'}
  #   {'pattern': None, 'type': 'message',   'channel': '__keyspace@0__:SNMP_PERMIT_LIST', 'data': 'del'}

  select_timeout = fake_infinite
  config_changed = False
  while True:
     try:
       m = p.get_message(timeout=select_timeout)
     except Exception:
       sys.exit(-2)
     # temporization: no change after 'timeout' seconds -> commit any accumulated changes
     if not m and config_changed:
        write_configuration_file(r.lrange(redis_key, 0, -1))
        config_changed = False
        select_timeout = fake_infinite
     if m and m['type'] == "message":
        if m['channel'] != subscription:
           print "WTF: unexpected case"
           continue
        config_changed = True
        select_timeout = temporization_duration
        # some debugs for now
        print "-------------------- config change: ",
        if m["data"] in delete_operations:
           print "DELETE"
        else:
           print ""
        v = r.lrange(redis_key, 0, -1)
        for value in v:
           print value

except redis.exceptions.ConnectionError as e:
  sys.exit(-2)