#!/usr/bin/python
# -*- coding: utf-8 -*-
# 
# check_dcmfind.py
# Nagios plugin to check a DICOM find service which is the first stage of a query/retrieve operation (find/c-move) .
# 
# The script is a wrapper for dcmtk's (http://dicom.offis.de/dcmtk.php.en) findscu
# to monitor FIND(Query) SCP.  dcmtk will have to be installed (or the binary built) on the
# nagios system. Easiest way on an ubuntu system is to use apt:
# sudo apt-get install dcmtk
# dcmtk version 3.6.0 has been used successfully
# The findscu command tries to find studies with the study date of 1950-01-01.  The query should be successful and return no studies.
# this work is based on check_dcm.py
#  Dale F. Hawley   July 2013

# 
# Usage:
#
#define command {
#         command_name    check_dcmfind
#          command_line    $USER1$/check_dcmfind.py -H $HOSTADDRESS$ -p $ARG1$ -aec $ARG2$ -aet $ARG3$ -v
#       
#         }
#
#define service {
#        host_name                       VARXcelera
#        service_description             DICOM Query VAR Xcelera
#        check_command                   check_dcmfind!7000!QR_SCP!TVCARDARCH 
#        initial_state                   o
#        max_check_attempts              3
#        check_interval                  10
#        retry_interval                  1
#        active_checks_enabled           1
#        check_period                    24x7
#        notification_interval           10
#        first_notification_delay        0
#        notification_period             24x7
#        notification_options            w,u,r,c,f,s
#        notifications_enabled           1
#        register                        1
#}

import subprocess
import sys
import argparse
import time
# Logging is enabled if DEBUG is True
DEBUG = False
#DEBUG = True
# Version
VERSION = '0.1.1'

# location of dcmtk findscu executable
cmd = ['/usr/local/bin/findscu']

# Exit Codes
STATE_OK = 0
STATE_WARNING = 1
STATE_CRITICAL = 2
STATE_UNKNOWN = 3

# Arguments
parser = argparse.ArgumentParser()
parser.add_argument("-V", "--version", action="store_true", help="display plugin version")
parser.add_argument("-v", "--verbosity", action="count", default=1, help="increase output verbosity")
parser.add_argument("-t", "--timeout", default='10', help="seconds before request times out")
parser.add_argument("-aet", "--aetitle", help="calling (this the client) Appplication Entity (AE) Title (default: FINDSCU)")
parser.add_argument("-aec", "--call", help="Appplication Entity (AE) Title of the DICOM service being checked (default: ANY-SCP)")
parser.add_argument("-H", "--hostname", help="hostname or IP address of the server to be checked")
parser.add_argument("-p", "--port", default='104', help="tcp/ip port number the server to be checked is listening on")

try:
	args = parser.parse_args()
except SystemExit, e:
	# bad args, so service state unknown
	sys.exit(STATE_UNKNOWN)
	raise e

if args.version:
	# service state unknown, but here's your version number.
	print VERSION
	sys.exit(STATE_UNKNOWN)

# build command line arguments
if args.verbosity >= 3:
	cmd.append('--debug')
elif args.verbosity == 2:
	cmd.append('--verbose')

cmd.append('--key')
cmd.append('(0008,0052)=STUDY')
cmd.append('--key')
cmd.append('(0008,0020)=19500101')
cmd.append('-P')

if args.timeout:
	cmd.append('-to') 
	cmd.append(args.timeout)

if args.aetitle:
	cmd.append('-aet')
	cmd.append(args.aetitle)

if args.call:
	cmd.append('-aec')
	cmd.append(args.call)

cmd.append(args.hostname)
cmd.append(args.port)


if DEBUG:
  # log the command to be executed
  # get local time
  localtime = time.asctime( time.localtime(time.time()) )
  #print "Local current time :", localtime
  # Open a file
  fo = open("/usr/local/nagios/chk_dcmfind.log", "a")
  fo.write( localtime )
  fo.write( " command: ")
  for item in cmd:
    fo.write("%s " % item)
  fo.write("\n")
  

# run the dcmtk command with arguments and return the result of the command execution
#  if DEBUG is TRUE stdout  ( the find result) will be written to the log file
try:
	p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
except Exception, e:
	if DEBUG:
            fo.write( "                         Exception when executing command\n")
            fo.close()
        sys.exit(STATE_UNKNOWN)
        
else:
	stdout = p.communicate()[0]

        if DEBUG:
            for line in stdout.splitlines():
		fo.write("                         %s \n" % line)
            fo.close()
		
	if p.returncode:
		if args.verbosity:
			print "DICOM CRITICAL - Association Request Failed (TCP Initialization Error: Connection refused)"
			if args.verbosity > 1:
				for line in stdout.splitlines():
					print line
		else:
			print "DICOM CRITICAL - Association Request Failed"

		sys.exit(STATE_CRITICAL)

	else:
		if args.verbosity:
			print "DICOM OK - Association Accepted (Received Response (Status: Success))"
			if args.verbosity > 1:
				for line in stdout.splitlines():
					print line
		else:
			print "DICOM OK - Association Accepted"

		sys.exit(STATE_OK)