#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 -1;
  }
  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;
  realloc.optional = false;
  map[REALLOC_SEC_COUNT_ID] = realloc;

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

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

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

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

      wearout.name = "Media_Wearout_Indicator";
      wearout.value = -1;
      wearout.severity = WARN;
      wearout.optional = true;
      map[WEAR_COUNT_ID] = wearout;

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

      rep_uncorrect.name = "Reported_Uncorrect";
      rep_uncorrect.value = -1;
      rep_uncorrect.severity = CRIT;
      rep_uncorrect.optional = false;
      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;
      int attrID = getSmartAttrID(line);
      if(attrID == -1) {
        continue;
      }
      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) {
      if(attr.optional) {
        continue;
      }
      *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);
}

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;
      case UNKN:
        returnCode = UNKN;
        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;
}