Storage.py

#!/usr/bin/python3 -O
"""
 File: Storage.py
 Author: Leon McClatchey
 Date: 2025-09-23

 Description: Utilizes SNMP calls to find the storage devices on a host
			  Gets the name of the device, Total Capacity, and Amount Used
			  Then from the SNMP Data, additional information is calculated
			  The results are then displayed to the command line,
			  which can be read by a NMS service such as Nagios
			  
 Requirements: This is a Standalone script, so no user libraries are called
			   However, if a log is maintained, it will be in the NMS/Nagios log folder
			   This script does rely on netsnmp availability on the server
			   The Availability of the MIB is all required
			   Libraries argparse and IPy must be installed
			   Linux Hosts, /etc/snmp/snmpd.conf must contain the line "includeAllDisks"
"""
# System Libraries
import os, sys, socket, subprocess
import argparse 
import logging
from IPy import IP
from errno import ECONNREFUSED
from datetime import datetime
# Initialize Global Strings
Config = {"OID":{"Base":"1.3.6.1.2.1.25.2.3.1","Mibs":["Entry","Index","Type","Descr","AllocationUnits","Size","Used"]}}
MyScript = os.path.splitext(os.path.basename(sys.argv[0]))[0]
# Class MyParser used to get the command line Arguments
class MyParser(argparse.ArgumentParser):
	def Error(self, message):
		global Config
		WriteLog(f"Argument Parsing Error: {message}",Config['log'])
		self.print_help()
		sys.stderr.write(f"error: {message}\n")
		sys.exit(2)
# Function to Archive the Logfile
def ArchiveLog(dic):
	cdir = os.getcwd()
	tDate=datetime.today().strftime("%Y%m%d")
	pth = os.path.dirname(dic['path'])
	archive = f"{os.path.splitext(os.path.basename(dic['path']))[0]}{tDate}.zip"
	log=f"{os.path.splitext(os.path.basename(dic['path']))[0]}"
	try:
		os.chdir(pth)
		lsize = os.path.getsize(os.path.basename(dic['path']))
		msize = dic['max'] * 10e6
		if lsize >= msize:
			zip = zipfile.ZipFile(archive,'w',zipfile.ZIP_DEFLATED)
			archive.write(zip)
			archive.close()
			os.remove(os.path.basename(dic['path']))
			WriteLog(f"Successfully compressed {os.path.basename(dic['path'])} to {archive} and removed {os.path.basename(dic['path'])}",dic['path'])
	except:
		WriteLog(f"Error Compressing Log",dic['path'])
	finally:
		os.chdir(cdir)
# Function to Build the Return List
def BuildList(idx):
	global Config
	Args = Config['Args']
	Oids = Config['OID']
	Mibs = Oids['Mibs']
	base = Mibs[idx]
	list = []
	for i in range(len(Oids['Index'])):
		cmd = f"snmpget -v{Args['version']} -c{Args['community']} {Args['host']} oid {base}.{Oids['Index'][i]}"
		rc = LocalCommand(cmd)
		if rc['out'] is not None:
			item = rc['out'].decode('utf-8').rsplit('\n',1)[-2]
			match idx:
				case 1:
					box = item.rsplit('hrStorage',1)[-1]
				case 2:
					box = item.rsplit(':',1)[-1].strip()
				case 3:
					tmp = item.rsplit(':',1)[-1].strip()
					box = int(tmp.split()[0])
				case 4:
					box = int(item.rsplit(':',1)[-1].strip())
				case 5:
					box = int(item.rsplit(':',1)[-1].strip())
				case _:
					continue
			list.append(box)
	return list
# Function to Create a Folder if it does not exist
def CreateFolder(pth):
	if not pth:
		print("Path must be specified")
		return
	tst=os.path.splitext(pth)
	wpth=os.path.dirname(pth) if tst[1] else pth
	if not os.path.exists(wpth):
		os.makedirs(wpth)
# Function to return a table of indexes
def CreateIndexTable():
	global Config
	args = Config['Args']
	Oids = Config['OID']
	Mibs = Oids['Mibs']
	base = f"{Oids['Base']}.1"
	cmd = f"snmpwalk -v{args['version']} -c{args['community']} {args['host']} {base}"
	rc = LocalCommand(cmd)
	if rc['code'] == 0:
		tbl = rc['out'].decode('utf-8').split('\n')
		lst = []
		for i in range(len(tbl)):
			if len(tbl[i].strip())>0:
				tmp = tbl[i].split('=')[1].split(':')
				if len(tmp[1].strip()) > 0:
					nbr = tmp[1].strip()
					lst.append(int(nbr) if nbr.isnumeric() else nbr)
		Config['OID']['Index'] = lst
	else:
		dsp = f"No Results Returned from {args['host']}"
		sys.exit(2)
# Function to create Mib Table
def CreateMibTable():
	global Config
	Args = Config['Args']
	OID = Config['OID']
	Mibs = []
	cmd = f"snmpwalk -v{Args['version']} -c{Args['community']} {Args['host']} {OID['Base']}"
	WriteLog(f"Obtaining SNMP Mibs from {Args['host']}!",Args['log'])
	rc = LocalCommand(cmd)
	if rc['code'] == 0:
		tbl = rc['out'].decode('utf-8').split('\n')
		WriteLog(f"Query returned {len(tbl)} records from {Args['host']}",Args['log'])
		for i in range(len(tbl)):
			if len(tbl[i].strip()) > 0:
				tmp = tbl[i].strip().split("::")[1].split('.')[0].strip('if')
				if tmp not in Mibs:
					Mibs.append(tmp)
		Config['OID']['Mibs'] = Mibs
	else:
		dsp = f"Unable to Connect via SNMP with {Args['host']}"
		WriteLog(dsp,Args['log'])
		print(dsp)
		sys.exit(2)
# Function to Display the Results
def DisplayResults():
	global Config
	Oids = Config['OID']
	Mibs = Oids['Mibs']
	Args = Config['Args']
	WriteLog(f"Displaying {len(Oids)} OIDS for {Config['Args']['host']}",Config['Args']['log'])
	for i in range(len(Oids['Index'])):
		if 'dev' in Oids['Descr'][i] or 'run' in Oids['Descr'][i] or 'proc' in Oids['Descr'][i] or 'lib' in Oids['Descr'][i]: continue
		if Args['memory'] is False and 'memory' in Oids['Descr'][i].lower(): continue
		if Args['remote'] is False and 'NetworkDisk' in Oids['Type'][i]: continue
		p = Oids['Index'][i]
		Size = round(Oids['Size'][i] / 10e6,3)
		Used = round(Oids['Used'][i] / 10e3,3)
		Avail = round((Oids['Size'][i] - Oids['Used'][i]) / 10e6,3)
		try:
			Rate = round((Oids['Used'][i] / Oids['Size'][i]) * 100,2)
		except ZeroDivisionError:
			Rate = 0
		Type = f"{Oids['Type'][i]}: " if "Disk" in Oids['Type'][i] else ''
		dsp = f"{Type}{Oids['Descr'][i]}\tSize: {Size}G\tUsed: {Used}K\tAvail: {Avail}G\tUse%: {Rate}"
		print(dsp)
# Function to obtain the commandline arguments
def GetArguments():
	global MyScript
	parser = MyParser(description='SNMP Storage Manager Version 1.0.0', usage='%(prog)s [options]', formatter_class=argparse.RawTextHelpFormatter)
	parser.add_argument("-c", '--community', dest='community', type=str, default='public', help="The snmp community: defaults to 'public'")
	parser.add_argument("-e", '--engine',dest='engine',type=str,help="Engine ID if SNMP Version 3 has been specified")
	parser.add_argument("-f", '--flags',dest='flags',type=str,default="c 10:w 20",help="Colon deliminated string for critical and warning percentage settings, defaults to 'c 10:w 15'")
	parser.add_argument("-H", '--host', dest='host', type=str, required=True, help="Host (Required): Name of the host running the target service")
	parser.add_argument("-l", '--log', dest='log',default="path=/usr/local/logs:max=500", help="colon deliminated Location and Max size of Log file\nDefaults to 'path=/usr/local/logs:max=500'\nUnless Nagios is defined, then it will be in the Nagios Log directory")
	parser.add_argument("-m", "--memory",dest='memory',default=True,action=argparse.BooleanOptionalAction,help='Flag to display Memory Devices, Defaults to "True"')
	parser.add_argument("-n", '--nagios', dest='nagios', help='location of NMS config file, complete path and filename')
	parser.add_argument("-p", '--port', dest='port', type=int, default=161, help="The snmp port: Defaults to '161'")
	parser.add_argument("-P", '--password',dest='password',type=str,help="Password if SNMP Version 3 has been specified")
	parser.add_argument("-r", "--remote",dest='remote',default=False,action=argparse.BooleanOptionalAction,help='Flag to display Remote Devices, Defaults to "True"')
	parser.add_argument("-v", '--version', dest='version', type=str, default='2c', help="The snmp version: must be '1,2c, or 3': defaults to '2c'")
	parser.add_argument("-u", '--user',dest='user',type=str,help="User if SNMP Version 3 has been specified")
	args = vars(parser.parse_args())
	if args['version'].strip().lower() not in ['1','2c','3']:
		sys.exit(f"Invalid SNMP Version '{args['version']}' Specified\nAborting Mission!")
	(pth,mx) = args['log'].split(':')
	path=pth.split('=')
	max=mx.split('=')
	lpth = path[1] if args['nagios'] is None else getNagiosLogPath(args['nagios'])
	args['log']={'path':f"{lpth}/{MyScript}.log",'max':int(max[1])}
	return args
# Function to retrieve the Nagios Log Folder
def getNagiosLogPath(path):
	srch = "log_file"
	with open(path,'r') as file:
		for line_num, line in enumerate(file,1):
			if srch in line:
				log = os.path.dirname(line.strip().split('=')[1])
				break
	return log
# Function to execute the local query
def LocalCommand(cmd):
	my_env = os.environ.copy()
	try:
		proc = subprocess.Popen([cmd],env=my_env,stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
		(out,err)=proc.communicate()
		ecode=proc.wait()
		return {'out':out,'err':err,'code':ecode}
	except Exception as error:
		print >> sys.stderr, f"An {type(error).__name__} occurred in {inspect.stack()[3]} function!"
		print >> sys.stderr, f"Exception: {str(error)}"
		sys.exit(3)
# Main Function directs how the script is executed
def main():
	global Config
	Config["Args"] = GetArguments()
	log = Config['Args']['log']['path']
	logging.basicConfig(filename=log,filemode='a',level=logging.WARN)
	ArchiveLog(Config['Args']['log'])
	if ValidateHost(Config['Args']['host']) == True:
		CreateMibTable()
		CreateIndexTable()
		UpdateTable()
		DisplayResults()
	else:
		dsp = f"Unable to Connect to {Config['Args']['host']}"
		WriteLog(dsp,Config['Args']['log'])
		print(dsp)
		sys.exit(2)
# Function to Update the SNMP Table Contents
def UpdateTable():
	global Config
	Args = Config['Args']
	Mibs = Config['OID']['Mibs']
	WriteLog(f"Updating {len(Mibs)} entries for SNMP Query Table on {Args['host']}",Args['log'])
	for i in range(1, len(Mibs)):
		_,_,key = Mibs[i].partition("Storage")
		Config['OID'][key] = BuildList(i)
# Validate the Host
def ValidateHost(host):
	try:
		IP(socket.gethostbyname(host))
		rc = LocalCommand(f"ping -c 1 -w 5 {host}")
		return rc['code'] == 0
	except ValueError:
		print("Invalid IP Address Specified")
		return False
# Function to Write to a log file
def WriteLog(msg,log):
	tDate=datetime.today().strftime("%Y-%m-%d %H:%M:%S")
	logpath=log['path']
	CreateFolder(logpath)
	pth=f"{logpath}/{MyScript}.log" if not '.log' in logpath else logpath
	dsp=f"{tDate}; {msg}\n"
	lf=open(pth,"a")
	lf.write(dsp)
	lf.close()
# Executes the Script by calling the main function
if __name__ =="__main__":
	main()