#!/usr/bin/env python

# Daemon that listens to updates from ConfigDB about the source IP prefixes from which
# SSH connections are allowed. In case of change, it will update the SSHD configuration
# file accordingly. SSHD will notice the file has changed next time a connection comes in.
# Future enhancement: if an entry it modified/removed, go through all existing ssh
# connections and recompute their permission, and in case one is now denied, kill it.
#
# 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 subprocess
import sys
import syslog
import time
from swsssdk import ConfigDBConnector

VERSION = "1.0"

SYSLOG_IDENTIFIER = "sshd-config-updater"


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

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

    ACL_TABLE = "ACL_TABLE"
    ACL_RULE = "ACL_RULE"

    ACL_TABLE_TYPE_CTRLPLANE = "CTRLPLANE"

    ACL_SERVICE_SSH = "SSH"

    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-SSH service ACLs
            if self.ACL_SERVICE_SSH 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
    #   bash# cat /etc/sshd.allow
    #   sshd: [fd7a:629f:52a4:b0c3:ec4:7aff:fe99:201e]/128
    #   sshd: 172.17.0.1/32
    #   sshd: 172.18.1.0/24
    # Note that any matches are 'permits', and the default action is 'denied'
    # We assume the database contains valid ip addresses/hostnames.
    #
    # 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.allow" % (self.CONFIG_FILE_PATH, self.SERVICE)

        if len(src_ip_allow_list) == 0:
            if os.path.exists(filename):
                os.remove(filename)
            return

        filename_tmp = filename + ".tmp"

        f = open(filename_tmp, "w")

        for value in src_ip_allow_list:
            f.write("%s: %s\n" % (self.SERVICE, value))
        f.close()

        os.rename(filename_tmp, filename)

        # some previously accepted sessions might no longer be allowed: clear them
        os.system("/usr/bin/sshd-clear-denied-sessions")

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