#!/usr/bin/env python

from datetime import datetime
from seine.snmp.agent import SNMPAgent, Item
from ultimatum.logformats.auth import SSHViolationsDatabase

TREE_PREFIX = '1.3.6.1.3.14.2.74.22'

class SSHViolationsAgent(SNMPAgent):
    def __init__(self):
        SNMPAgent.__init__(self, TREE_PREFIX, reload_interval=60)

        self.database = SSHViolationsDatabase()

        self.indexes = self.register_tree('%s.3' % TREE_PREFIX)
        self.counters = self.register_tree('%s.4' % TREE_PREFIX)

        self.reload()

    def reload(self):
        self.clear()

        total = 0

        for i, entry in enumerate(self.database.source_address_counts()):
            self.indexes.add(Item(self.indexes.oid + [i+1], 'string', entry['address']))
            self.counters.add(Item(self.counters.oid + [i+1], 'integer', entry['count']))
            total += entry['count']

        self.tree.add(Item('%s.1' % TREE_PREFIX, 'integer', total))
        self.tree.add(Item('%s.2' % TREE_PREFIX, 'string', datetime.now().strftime('%Y-%m-%d %H:%M:%S')))

SSHViolationsAgent().run()
