#include "check_smart.h"

char *servicename = (char*)"SMART";

int evalStatus(char* disk, string reallocatedCount, string currentPending, string offlineUncorrectable, string *status)
{
	int returnStatus = 0;
	*status = string(disk) + " reallocated_count:";
	string reallocatedNum = reallocatedCount.substr(reallocatedCount.find_last_of(" ")+1);
	if(reallocatedNum != "0")
	{
		returnStatus = 1;
	}
	*status += reallocatedNum + " current_pending:";
	string currentNum = currentPending.substr(currentPending.find_last_of(" ")+1);
	if(currentNum != "0")
	{
		returnStatus = 2;
	}
	*status += currentNum + " offline_uncorrectable:";
	string offlineNum = offlineUncorrectable.substr(offlineUncorrectable.find_last_of(" ")+1);
	if(offlineNum != "0")
	{
		returnStatus = 2;
	}
	*status += offlineNum;
	return returnStatus;
}

void printVersion()
{
	cout << "check_smart v" << VERSION << endl << endl;
}

void printHelp(bool longVersion)
{
	if(longVersion)
	{
		printVersion();
		cout << "Checks for pending, reallocated or uncorrectable sectors in disks using SMART" << endl << "WARNING: Requires the setuid bit to be set to run as root" << endl << endl;
		printHelp(false);
		cout << "Options:" << endl;
		cout << " -h" << endl;
		cout << "    Print detailed help screen" << endl;
		cout << " -V" << endl;
		cout << "    Print version information" << endl;
		cout << " DISKS" << endl;
		cout << "    Disks for which to retrieve S.M.A.R.T data" << endl << endl;
		return;
	}
	cout << "Usage: " << endl << "check_smart [-hV] DISKS..." << endl << endl;
}

int main(int argc, char **argv)
{
	struct itimerval timer;
	timer.it_value.tv_sec = 10;
	timer.it_value.tv_usec = 0;
	timer.it_interval.tv_sec = 0;
	timer.it_interval.tv_usec = 0;
	setitimer (ITIMER_VIRTUAL, &timer, 0);

	struct sigaction sa;
	memset (&sa, 0, sizeof (sa));
	sa.sa_handler = &timer_handler;
	sigaction (SIGVTALRM, &sa, 0);

	int c;
        while ((c = getopt (argc, argv, "Vh")) != -1)
        {
                switch(c)
                {
                        case 'h':
                                printHelp(true);
				return 0;
                        case 'V':
                                printVersion();
				return 0;
                        case '?':
                                printHelp(false);
                                return 3;
		}
	}

	if(argc == 1)
	{
		cout << "No disks checked" << endl;
		return 3;
	}

	string *results = new string[argc-1];
	int returnCode = 0;

	string command = string(SMARTCTL_CMD);
	stringstream sstream;
	string line = "";
	string reallocatedCount = "";
	string currentPending = "";
	string offlineUncorrectable = "";
	string status = "";
	for(int i = 1; i < argc; i++)
	{
		string output;
		int rc = exec(command + argv[i],&output);
		if(rc)
        	{
                	cout << "Error reading SMART data from disk" << endl;
                	exit(3);
        	}
		sstream.str(output);
		while(getline(sstream,line))
		{
			int j = 0;
			while(line[j] == ' ')
			{
				j++;
			}
			size_t spacePos = line.find_first_of(" ",j);
			if(line.substr(j,spacePos-j) == "5")
			{
				reallocatedCount = line;
			}
			else if(line.substr(j,spacePos-j) == "197")
			{
				currentPending = line;
			}
			else if(line.substr(j,spacePos-j) == "198")
                        {
                                offlineUncorrectable = line;
                        }
		}
		if((reallocatedCount == "") || (currentPending == "") || (offlineUncorrectable == ""))
		{
			returnCode = 3;
			status = string(argv[i]) + " status UNKNOWN";
		}
		else
		{
			switch(evalStatus(argv[i],reallocatedCount,currentPending,offlineUncorrectable,&status))
			{
				case 0:
					//everything is OK
					break;
				case 1:
					//Warning issued
					if(!returnCode)
					{
						returnCode = 1;
					}
					break;
				case 2:
					//Critical issued
					returnCode = 2;
					break;
			}
		}
		results[i-1] = status;
		reallocatedCount = "";
                currentPending = "";
                offlineUncorrectable = "";
                output = "";
		line = "";
		status = "";
                sstream.str("");
		sstream.clear();
	}
	cout << servicename;
	if(returnCode == 0)
	{
		cout << " OK - disk status: ";
	}
	else if(returnCode == 1)
	{
		cout << " WARNING - disk status: ";
	}
	else if(returnCode == 2)
	{
		cout << " CRITICAL - disk status: ";
	}
	else if(returnCode == 3)
	{
		cout << " UNKNOWN - disk status: ";
	}
	for(int i = 0; i < (argc-1); i++)
	{
		cout << results[i];
		if(i != (argc-2))
		{
			cout << ", ";
		}
	}
	cout << endl;
	return returnCode;
}