#include "check_smart.h"

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

int getSmartAttrValue(string line)
{
  return stoi(line.substr(line.find_last_of(" ")+1));
}

int getSmartAttrID(string line)
{
  size_t first = line.find_first_not_of(' ');
  size_t last = line.find_first_of(' ',first);
  int id = -1;
  try {id = stoi(line.substr(first,last-first));}
  catch(...){}
  return id;
}

int checkDriveType(const char* disk)
{
  string output = "";
  string RRLine = "";
  string line= "";

  int rc = run_smartctl_cmd(string(SMARTCTL_CMD_INFO), disk,&output);
  if(rc)
  {
    cout << "Error reading SMART data from disk " << disk << endl;
    exit(3);
  }

  stringstream sstream;
  sstream.str(output);
  while(getline(sstream,line) && RRLine == "")
  {
    size_t found = line.find(ROTATION_INFO_STR);
    if(found != string::npos) {RRLine = line;}
  }
  if(RRLine != "")
  {
    size_t found = RRLine.find(SSD_DEV_STR);
    if(found != string::npos) {return SSD;}
  }
  return HDD;
}

map<int,SMARTAttr> prepareAttrMap(int driveType)
{
  map<int,SMARTAttr> map;

  SMARTAttr realloc;
  realloc.name = "Reallocated_Sector_Ct";
  realloc.value = -1;
  realloc.severity = WARN;
  map[REALLOC_SEC_COUNT_ID] = realloc;

  SMARTAttr pending;
  SMARTAttr off_uncorrect;
  SMARTAttr wear;
  SMARTAttr badblocks;
  SMARTAttr rep_uncorrect;

  switch(driveType)
  {
    case HDD:
      pending.name = "Current_Pending_Sector";
      pending.value = -1;
      pending.severity = CRIT;
      map[CURRENT_PENDING_SEC_ID] = pending;

      off_uncorrect.name = "Offline_Uncorrectable";
      off_uncorrect.value = -1;
      off_uncorrect.severity = CRIT;
      map[OFFLINE_UNCORRECT_ID] = off_uncorrect;

      break;
    case SSD:
      wear.name = "Wear_Leveling_Count";
      wear.value = -1;
      wear.severity = WARN;
      map[WEAR_COUNT_ID] = wear;

      badblocks.name = "Runtime_Bad_Block";
      badblocks.value = -1;
      badblocks.severity = CRIT;
      map[RUNTIME_BAD_BLOCKS_IDD] = badblocks;

      rep_uncorrect.name = "Reported_Uncorrect";
      rep_uncorrect.value = -1;
      rep_uncorrect.severity = CRIT;
      map[REP_UNCORRECT_ID] = rep_uncorrect;

      break;
  }
  return map;
}

int evalStatus(const char* disk, int driveType, string *status)
{
  string output = "";
  string line = "";

  int rc = run_smartctl_cmd(string(SMARTCTL_CMD_ATTRS), disk,&output);
  if(rc)
  {
    cout << "Error reading SMART data from disk " << disk << endl;
    exit(UNKN);
  }

  map<int,SMARTAttr> attrMap = prepareAttrMap(driveType);

  stringstream sstream;
  sstream.str(output);
  while(getline(sstream,line))
  {
    for(map<int,SMARTAttr>::iterator it = attrMap.begin(); it != attrMap.end(); ++it)
    {
      int id = it->first;
      if(getSmartAttrID(line) == id) {attrMap[id].value = getSmartAttrValue(line);}
    }
  }
  int ret = OK;
  *status = string(disk);
  for(map<int,SMARTAttr>::iterator it = attrMap.begin(); it != attrMap.end(); ++it)
  {
    int id = it->first;
    SMARTAttr attr = it->second;

    if(attr.value == -1)
    {
      *status = string(disk) + " status UNKNOWN";
      return UNKN;
    }
    int veredict = 0;
    if(attr.value) {veredict = attr.severity;}
    switch(veredict)
    {
      case OK:
        break;
      case WARN:
        if(ret == OK) {ret = WARN;}
        break;
      case CRIT:
        ret = CRIT;
        break;
    }
    *status += " " + attr.name + ":" + to_string(attr.value);
  }
  return ret;
}

int run_smartctl_cmd(const string command, const char* disk, string* output)
{
  int uid = getuid();
  setreuid(0,0);

  int rc = exec(command + disk,output);

  setreuid(uid,0);
  return rc;
}

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;
}

void set_timeout(unsigned int sec)
{
  struct itimerval timer;
  timer.it_value.tv_sec = sec;
  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);
<<<<<<< HEAD
<<<<<<< HEAD
=======

>>>>>>> f3a48aed56edc3f5053f68421dbc764d5930797f
=======

>>>>>>> f3a48aed56edc3f5053f68421dbc764d5930797f
}

int main(int argc, char **argv)
{
  set_timeout(10);
  int c;
  while ((c = getopt (argc, argv, "Vh")) != -1)
  {
    switch(c)
    {
      case 'h':
        printHelp(true);
	return OK;
      case 'V':
        printVersion();
        return OK;
      case '?':
        printHelp(false);
        return UNKN;
    }
  }

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

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

  for(int i = 1; i < argc; i++)
  {
    int driveType = checkDriveType(argv[i]);
    switch(driveType)
    {
      case HDD:
      case SSD:
        break;
      default:
        cout << "Unknown disk type" << endl;
        return UNKN;
        break;
    }

    switch(evalStatus(argv[i],driveType,&(results[i-1])))
    {
      case OK:
        break;
      case WARN:
        if(returnCode == OK) {returnCode = WARN;}
        break;
      case CRIT:
        returnCode = CRIT;
        break;
    }
  }

  cout << servicename;
  if(returnCode == OK) {cout << " OK - disk status: ";}
  else if(returnCode == WARN) {cout << " WARNING - disk status: ";}
  else if(returnCode == CRIT) {cout << " CRITICAL - disk status: ";}
  else if(returnCode == UNKN) {cout << " UNKNOWN - disk status: ";}
  for(int i = 0; i < (argc-1); i++)
  {
    cout << results[i];
    if(i != (argc-2)) {cout << ", ";}
  }
  cout << endl;
  return returnCode;
}