//
// $Id: ProcessProbe.cpp,v 1.1 2004/06/01 17:05:21 bakerj Exp $
//
//************************** Property of the MITRE Corporation ***************************//
//
// Copyright (c) 2003 - The MITRE Corporation
//
// This file is part of the Query-based Network Assessment project.
//
// The Query-based Network Assessment is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License as published by the Free
// Software Foundation; either version 2 of the License.
//
// The Query-based Network Assessment is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along with the
// Query-based Network Assessment; if not, write to the Free Software Foundation, Inc.,
// 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
//
//****************************************************************************************//

#include "ProcessProbe.h"

//****************************************************************************************//
//								ProcessProbe Class								  //	
//****************************************************************************************//
ProcessProbe::ProcessProbe()
{
	// -----------------------------------------------------------------------
	//	Abstract
	//
	//	Do nothing for now
	// -----------------------------------------------------------------------
}

ProcessProbe::~ProcessProbe()
{
	// -----------------------------------------------------------------------
	//	Abstract
	//
	//	Do nothing for now
	// -----------------------------------------------------------------------
}

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  Public Members  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
pdVector ProcessProbe::Run(ProbeData *probeDataIn)
{
  //------------------------------------------------------------------------------------//
  //
  //  ABSTRACT
  //  
  //  Read the command passed in the data object and get all process information for 
  //  any process that matches the input command string.
  //  Support pattern matching on the command string.
  //  This probe will prepend a '^' and append a '$' to all litteral command strings.
  //  Then pattern matching can be used for all tests. This is different than the 
  //  rest of the probes.
  //
  //------------------------------------------------------------------------------------//
  ProcessData *dataIn = (ProcessData*)probeDataIn;
  pdVector resultVector;

  // set up the litteral stirng as a regular expression
  if(dataIn->command->type == LITTERAL_TYPE)
    {
      // escape any regex chars 
      dataIn->command->data = myMatcher->EscapeRegexChars(dataIn->command->data);

      // add begin and end string designators
      dataIn->command->data = "^" + dataIn->command->data + "$";
    }

  // Gather new data and add to resultVector
  GetPSInfo(dataIn, &resultVector);

  // Make sure that the result vector has at least one item
  if(resultVector.size() == 0)
    {
      ProcessData *tmp = new ProcessData();
      tmp->SetTestId(dataIn->GetTestId());
      tmp->msg = "No matching commands found";
      resultVector.push_back(tmp);
    }

  return resultVector;
}

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  Private Members  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

// Main method for gathering process information. 
//
void ProcessProbe::GetPSInfo(ProcessData *dataIn, pdVector *resultVector)
{
  string errMsg = "";
  ProcessData *tmp = NULL;

  // Time parameters
  time_t currentTime = 0;
  unsigned long adjustedStartTime, execTime = 0;

  // TTY String
  char ttyName[TTY_LEN + 1];

  // Process Parameters
  char cmdline[CMDLINE_LEN + 1];
  char schedulingClass[SCHED_CLASS_LEN + 1];

  int uid, pid, ppid;
  long priority = 0;
  unsigned long starttime = 0;

  int status = 0;


  // Grab the current time and uptime(Linux only) to calculate start and exec times later
  time(&currentTime);

  unsigned long uptime = 0;
  status = RetrieveUptime(&uptime, &errMsg);
  if(status < 0)
    {
      tmp = new ProcessData();
      tmp->SetTestId(dataIn->GetTestId());
      tmp->msg = errMsg;
      resultVector->push_back(tmp);
      return; 
    }

  DIR *proc;
  struct dirent *readProc;

  // Step into the /proc directory
  if((proc = opendir("/proc")) == NULL) {
    errMsg.append("ProcessProbe: Could not open /proc");
    tmp = new ProcessData();
    tmp->SetTestId(dataIn->GetTestId());
    tmp->msg = errMsg;
    resultVector->push_back(tmp);
    return;

  } else {

    // Loop through all of the files - we're only concerned with those that
    // start with a digit
    while((readProc = readdir(proc)) != NULL) {
      if(isdigit(readProc->d_name[0])) {

	// Clear the ps values
	memset(cmdline, 0, CMDLINE_LEN + 1);
	memset(schedulingClass, 0, SCHED_CLASS_LEN + 1);
	memset(ttyName, 0, TTY_LEN + 1);
	uid = pid = ppid = priority = starttime = 0;
	adjustedStartTime = execTime = 0;
	errMsg = "";

	// Retrieve the command line with arguments
	status = RetrieveCommandLine(readProc->d_name, cmdline, &errMsg);
	if(status < 0) { 
	  tmp = new ProcessData();
	  tmp->SetTestId(dataIn->GetTestId());
	  tmp->msg = errMsg;
	  resultVector->push_back(tmp); 
	  continue;

	  // If the command line matches the input command line get the remaining 
	  // data and add a new data object tot he resultVector
	} else if(status == 0 && myMatcher->IsMatch(dataIn->command->data.c_str(), cmdline)) {

	  // Read the 'stat' file for the remaining process parameters
	  status = RetrieveStatFile(readProc->d_name, &uid, &pid, &ppid, &priority, &starttime, &errMsg);
	  if(status < 0)
	    {
	      tmp = new ProcessData();
	      tmp->SetTestId(dataIn->GetTestId());
	      tmp->command->data = cmdline;
	      tmp->scheduling_class = "-";
	      tmp->msg = errMsg;
	      resultVector->push_back(tmp); 
	      continue;
	    }

	  // We can retrieve a value for the tty from the 'stat' file, but it's unclear
	  // how you convert that to a device name.  Therefore, we ignore that value
	  // and grab the device stdout(fd/0) is attached to.
	  RetrieveTTY(readProc->d_name, ttyName);

	  // The Linux start time is represented as the number of jiffies (1/100 sec)
	  // that the application was started after the last system reboot.  To get an
	  // actual timestamp, we have to do some gymnastics.  We then use the calculated
	  // start time to determine the exec time.
	  //
	  if(uptime > 0) {
	    adjustedStartTime = currentTime - (uptime - (starttime/100));
	    execTime = currentTime - adjustedStartTime;
	  }

	  // Add the data to a new data object and add th resultVector
	  tmp = new ProcessData();
	  tmp->SetTestId(dataIn->GetTestId());
	  tmp->command->data = cmdline;
	  tmp->scheduling_class = "-";
	  tmp->user_id = Common::ToString(uid);
	  tmp->pid = Common::ToString(pid);
	  tmp->ppid = Common::ToString(ppid);
	  tmp->priority = Common::ToString(priority);
	  tmp->start_time = Common::ToString(adjustedStartTime);
	  tmp->tty = ttyName;
	  tmp->exec_time = Common::ToString(execTime);
	  tmp->msg = errMsg;
	  resultVector->push_back(tmp); 
	}
      }
    }
    closedir(proc);
  }
}


// Read /proc/<pid>/cmdline to gather the application name and startup arguments
//
int ProcessProbe::RetrieveCommandLine(char *process, char *cmdline, string *errMsg) {
  int i = 0;
  int bytes = 0;
  FILE *cmdlineFile = NULL;

  // Build the absolute path to the command line file
  string cmdlinePath = "/proc/";
  cmdlinePath.append(process);
  cmdlinePath.append("/cmdline");

  // Open the file for reading
  if((cmdlineFile = fopen(cmdlinePath.c_str(), "r")) == NULL) {
    errMsg->append("ProcessProbe: Unable to obtain command line for pid: ");
    errMsg->append(process);
    return(-1);
    
  } else {
    
    // Read the contents of the file, and convert all of the NULL's
    // separating the args to spaces.
    //
    bytes = fread(cmdline, sizeof(char), CMDLINE_LEN, cmdlineFile);
    
    // Remember to leave the trailing NULL (bytes - 1).
    for(i = 0; i < bytes - 1; i++) {
      if(cmdline[i] == '\0') {
	cmdline[i] = ' ';
      }
    }
  }

  return(0);
}


// Read the stat file for a specific process
//
int ProcessProbe::RetrieveStatFile(char *process, int *uid, int *pid, int *ppid, long *priority, unsigned long *starttime, string *errMsg) {

  // Stat File parameters.  While we're really only concerned with gathering the parameters
  // that are passed in, these variables are good placeholders in case we decide to collect
  // something else in the future.
  //
  int pgrp, session, tty, tpgid, exit_signal, processor = 0;
  long cutime, cstime, nice, placeholder, itrealvalue, rss = 0;
  unsigned long flags, minflt, cminflt, majflt, cmajflt, utime, stime, vsize, rlim = 0;
  unsigned long startcode, endcode, startstack, kstkesp, kstkeip, signal, blocked, sigignore = 0;
  unsigned long sigcatch, wchan, nswap, cnswap = 0;
  char comm[PATH_MAX];
  char state;

  FILE *statFile = NULL;

  // Generate the absolute path name for the 'stat' file
  string statPath = "/proc/";
  statPath.append(process);
  statPath.append("/stat");

  // While we're here, stat() the 'stat' file to get the uid for the process.  If
  // we want to convert this to a user name, feed the uid to getpwuid().
  //
  struct stat statBuf;

  if((stat(statPath.c_str(), &statBuf)) < 0) {
    errMsg->append("ProcessProbe: Unable to obtain uid information for pid:  ");
    errMsg->append(process);
    return(-1);
  } else {
    *uid = statBuf.st_uid;
  }
    
  // Open the 'stat' file and read the contents
  if((statFile = fopen(statPath.c_str(), "r")) == NULL) {
    errMsg->append("ProcessProbe: Unable to obtain process information for pid: ");
    errMsg->append(process);
    return(-1);

  } else {

    // Linux gives us a nicely formatted file for fscanf to pull in
    fscanf(statFile, "%d %s %c %d %d %d %d %d %lu %lu %lu %lu %lu %lu %lu %ld %ld %ld %ld %ld %ld %lu %lu %ld %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu %d %d", pid, comm, &state, ppid, &pgrp, &session, &tty, &tpgid, &flags, &minflt, &cminflt, &majflt, &cmajflt, &utime, &stime, &cutime, &cstime, priority, &nice, &placeholder, &itrealvalue, starttime, &vsize, &rss, &rlim, &startcode, &endcode, &startstack, &kstkesp, &kstkeip, &signal, &blocked, &sigignore, &sigcatch, &wchan, &nswap, &cnswap, &exit_signal, &processor);

  }
  fclose(statFile);

  return(0);
}


// Since there appears to be no simple way to convert the 'tty' value contained in
// '/proc/<pid>/stat' into a device name, we instead use '/proc/<pid>/fd/0', which is
// normally linked to a device.  Note, the 'fd' directory is set read-only user, so
// if this probe is not run as root, many of these reads will fail.  In that case, we
// return '?' as the tty value.
//
void ProcessProbe::RetrieveTTY(char *process, char *ttyName) {
  int bytes = 0;

  // Generate the absolute path name for the stdout(0) file in 'fd'
  string ttyPath = "/proc/";
  ttyPath.append(process);
  ttyPath.append("/fd/0");

  // Attempt to read the name of the file linked to '0'
  bytes = readlink(ttyPath.c_str(), ttyName, TTY_LEN);

  // If there is an error, set the tty string to '?'
  if(bytes < 0 || strncmp(ttyName, "/dev", 4) != 0) {
    strncpy(ttyName, "?\0", 2);
  }
}


// Read the value contained in '/proc/uptime/' so that we can calculate
// the start time and exec time of the running processes.
//
int  ProcessProbe::RetrieveUptime(unsigned long *uptime, string *errMsg) {
  // The second value in this file represents idle time - we're not concerned with this.
  unsigned long idle = 0;
  FILE *uptimeHandle = NULL;

  // Open and parse the 'uptime' file
  if((uptimeHandle = fopen("/proc/uptime", "r")) == NULL) {
    errMsg->append("ProcessProbe: Could not open /proc/uptime");
    uptime = 0;
    return(-1);
  } else {
    fscanf(uptimeHandle, "%lu %lu", uptime, &idle);
  }
  fclose(uptimeHandle);

  return(0);
}
