#!/usr/bin/env python

# Daemon that listens to updates from ConfigDB about the source IP prefixes from which
# SNMP connections are allowed. In case of change, it will update the SNMP configuration
# file accordingly. After a change, it will notify snmpd to re-read its config file
# via SIGHUP.
#
# This daemon is meant to be run on Arista platforms only. Service ACLs on all other
# platforms will be managed by caclmgrd.
#

import os
import re
import signal
import subprocess
import sys
import syslog
import time
from swsssdk import ConfigDBConnector

VERSION = "1.0"

SYSLOG_IDENTIFIER = "snmpd-config-updater"


# ============================== Classes ==============================

class Process(object):
    def __init__(self, pid):
        self.pid = pid
        self.path = '/proc/%d/status' % pid
        self.status = None

    def read_proc_status(self):
        data = {}
        with open(self.path) as f:
            for line in f.readlines():
                key, value = line.split(':', 1)
                data[ key ] = value.strip()
        self.status = data

    def get_proc_signals(self):
        assert self.status
        sigBlk = int(self.status[ 'SigBlk' ], 16)
        sigIgn = int(self.status[ 'SigIgn' ], 16)
        sigCgt = int(self.status[ 'SigCgt' ], 16)
        return (sigBlk, sigIgn, sigCgt)

    def handle_signal(self, sig):
        sigBlk, sigIgn, sigCgt = self.get_proc_signals()
        mask = 1 << ( sig - 1 )
        if mask & sigBlk:
            return True
        if mask & sigIgn:
            return True
        if mask & sigCgt:
            return True
        return False

    def send_signal(self, sig):
        log_info('Sending signal %s to %d' % (sig, self.pid))
        os.kill(self.pid, sig)

    def safe_send_signal(self, sig):
        self.read_proc_status()
        if not self.handle_signal(sig):
            return False
        self.send_signal(sig)
        return True

    def wait_send_signal(self, sig, interval=0.1):
        while not self.safe_send_signal(sig):
            log_info('Process %s has not yet registered %s' % (self.pid, sig))
            time.sleep(interval)

    @staticmethod
    def by_name(name):
        try:
            pid = subprocess.check_output([ 'pidof', '-s', name ])
        except subprocess.CalledProcessError:
            return None
        return Process(int(pid.rstrip()))

class ConfigUpdater(object):
    SERVICE = "snmpd"
    CONFIG_FILE_PATH = "/etc/snmp"

    ACL_TABLE = "ACL_TABLE"
    ACL_RULE = "ACL_RULE"

    ACL_TABLE_TYPE_CTRLPLANE = "CTRLPLANE"

    ACL_SERVICE_SNMP = "SNMP"

    def get_src_ip_allow_list(self):
        src_ip_allow_list = []

        # Get current ACL tables and rules from Config DB
        tables_db_info = self.config_db.get_table(self.ACL_TABLE)
        rules_db_info = self.config_db.get_table(self.ACL_RULE)

        # Walk the ACL tables
        for (table_name, table_data) in tables_db_info.iteritems():
            # Ignore non-control-plane ACL tables
            if table_data["type"] != self.ACL_TABLE_TYPE_CTRLPLANE:
                continue

            # Ignore non-SNMP service ACLs
            if self.ACL_SERVICE_SNMP not in table_data["services"]:
                continue

            acl_rules = {}

            for ((rule_table_name, rule_id), rule_props) in rules_db_info.iteritems():
                if rule_table_name == table_name:
                    acl_rules[rule_props["PRIORITY"]] = rule_props

            # For each ACL rule in this table (in descending order of priority)
            for priority in sorted(acl_rules.iterkeys(), reverse=True):
                rule_props = acl_rules[priority]

                if "PACKET_ACTION" not in rule_props:
                    log_error("ACL rule does not contain PACKET_ACTION property")
                    continue

                # We're only interested in ACCEPT rules
                if rule_props["PACKET_ACTION"] != "ACCEPT":
                    continue

                if "SRC_IP" in rule_props and rule_props["SRC_IP"]:
                    src_ip_allow_list.append(rule_props["SRC_IP"])

        return src_ip_allow_list

    # 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).
    #
    # src_ip_allow_list may contain individual IP addresses or blocks of
    # IP addresses using CIDR notation.
    def write_configuration_file(self, src_ip_allow_list):
        filename = "%s/%s.conf" % (self.CONFIG_FILE_PATH, self.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 line in snmpd_config.split('\n'):
            m = re.match("^(..)community (\S+)", line)
            if not m:
                f.write(line)
                f.write("\n")
            else:
                if not line.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(src_ip_allow_list):
                        for value in src_ip_allow_list:
                            f.write("%s %s\n" % (this_community, value))
                    else:
                        f.write("%s\n" % this_community)
        f.close()

        os.rename(filename_tmp, filename)

        # Force snmpd process to reload its configuration if it is running
        proc = Process.by_name(self.SERVICE)
        if proc:
           proc.wait_send_signal(signal.SIGHUP)

    def notification_handler(self, key, data):
        log_info("ACL configuration changed. Updating {} config accordingly...".format(self.SERVICE))
        self.write_configuration_file(self.get_src_ip_allow_list())

    def run(self):
        # Open a handle to the Config database
        self.config_db = ConfigDBConnector()
        self.config_db.connect()

        # Write initial configuration
        self.write_configuration_file(self.get_src_ip_allow_list())

        # Subscribe to notifications when ACL tables or rules change
        self.config_db.subscribe(self.ACL_TABLE,
                lambda table, key, data: self.notification_handler(key, data))
        self.config_db.subscribe(self.ACL_RULE,
                lambda table, key, data: self.notification_handler(key, data))

        # Indefinitely listen for Config DB notifications
        self.config_db.listen()


# ========================== Syslog wrappers ==========================

def log_info(msg):
    syslog.openlog(SYSLOG_IDENTIFIER)
    syslog.syslog(syslog.LOG_INFO, msg)
    syslog.closelog()


def log_warning(msg):
    syslog.openlog(SYSLOG_IDENTIFIER)
    syslog.syslog(syslog.LOG_WARNING, msg)
    syslog.closelog()


def log_error(msg):
    syslog.openlog(SYSLOG_IDENTIFIER)
    syslog.syslog(syslog.LOG_ERR, msg)
    syslog.closelog()


# Determine whether we are running on an Arista platform
def is_platform_arista():
    proc = subprocess.Popen(["sonic-cfggen", "-H", "-v", "DEVICE_METADATA.localhost.platform"],
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    (stdout, stderr) = proc.communicate()

    if proc.returncode != 0:
        log_error("Failed to retrieve platform string")
        return False

    return "arista" in stdout


def main():
    log_info("Starting up...")

    if not os.geteuid() == 0:
        log_error("Must be root to run this daemon")
        print "Error: Must be root to run this daemon"
        sys.exit(1)

    if not is_platform_arista():
        log_info("Platform is not an Arista platform. Exiting...")
        sys.exit(0)

    # Instantiate a ConfigUpdater object
    config_updater = ConfigUpdater()
    config_updater.run()

if __name__ == "__main__":
    main()