#!/usr/bin/env python

import sys, os, subprocess, argparse

def snmp_walk(command_line, oid):        
    command = command_line + ' ' + oid
    try:   
        result = subprocess.Popen(command.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    
    except Exception, details:
        raise Exception(command + ': ' + str(details))

    if result.wait():
        raise Exception(command + ': ' + result.stderr.read())

    output = result.stdout.read()
    
    if not output:
        raise Exception(command + ': ' + 'No response')

    result_stdout = output.strip().split('\n')

    return [ "".join( s.split()[-1:] ) for s in result_stdout ]      

def snmp_get_int(command_line, oid):
    return int(snmp_walk(command_line, oid)[0])

def check_malware(command_line):
    counters = [
        ('virus'     , snmp_get_int(command_line, 'DrWeb-Snmpd::knownVirus')),
        ('suspicious', snmp_get_int(command_line, 'DrWeb-Snmpd::suspicious')),        
        ('adware'    , snmp_get_int(command_line, 'DrWeb-Snmpd::adware'    )),
        ('dialers'   , snmp_get_int(command_line, 'DrWeb-Snmpd::dialers'   )), 
        ('joke'      , snmp_get_int(command_line, 'DrWeb-Snmpd::joke'      )),
        ('riskware'  , snmp_get_int(command_line, 'DrWeb-Snmpd::riskware'  )), 
        ('hacktool'  , snmp_get_int(command_line, 'DrWeb-Snmpd::hacktool'  ))
    ]

    total_threats = 0
    for counter in counters:
        total_threats += counter[1]
  
    percents = []
    for counter in counters:
        percent = 100*counter[1]/total_threats if total_threats > 0 else 0
        percents.append(counter[0] + '=' + str(percent) + '%')

    return 'SNMP OK - total malware: ' + str(total_threats) + '|' + ' '.join( percents )

def check_filecheck(command_line):
    counter = snmp_walk(command_line, 'DrWeb-Snmpd::fileCheckScannedBytes')[0]
    return 'SNMP OK - total scanned: ' + counter +'| scannedBytes=' + counter + 'c'


if __name__ == "__main__":
#    log = open('/tmp/check_drweb.log', 'a')
    commands = {
        'malware'   : check_malware,
        'filecheck' : check_filecheck
    }

    try:
        parser = argparse.ArgumentParser('check drweb state')
        parser.add_argument(
            '-t', '--target',
            choices=commands.keys(),
            help="get graph meta-info for specialized target", 
            default=commands.keys()[0]
        )

        parser.add_argument(
            '-s', '--snmp_walk_command',
            help="get graph meta-info by calling this command", 
            default='snmpwalk -c public -v 2c localhost:161'
        )
       
        args = parser.parse_args()
#        log.write('ARGS: ' + str(args) + '\n')

    except Exception, details:
        print('ERROR: '+str(details))
#        log.write("BAD ARGUMENTS: " + str(details) + '\n')
        parser.print_help()
        sys.exit(2)

    try:
        command_line = args.snmp_walk_command
        target       = args.target

        output = commands[target](command_line)
#        log.write('OUTPUT: ' + output + '\n')
        print(output)
        sys.exit(0)

    except Exception, details:
        print('ERROR: '+str(details))
#        log.write('ERROR: ' + str(details) + '\n')
        sys.exit(2)
