//
// $Id: InetListeningServersProbe.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 "InetListeningServersProbe.h"

//****************************************************************************************//
//						InetListeningServersProbe Class									  //	
//****************************************************************************************//

InetListeningServersProbe::InetListeningServersProbe()
{
	// -----------------------------------------------------------------------
	//	Abstract
	//
	//	Do nothing for now
	//
	// -----------------------------------------------------------------------

}

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

}

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  Public Members  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

pdVector InetListeningServersProbe::Run(ProbeData *probeDataIn)
{
  //------------------------------------------------------------------------------------//
  //  ABSTRACT
  //  
  //  Read the program_name passed in the data object and gather inet listening server
  //  information for all records that have a matching program name
  //  Support pattern matching on the program_name string.
  //  This probe will prepend a '^' and append a '$' to all litteral command strings.
  //  Then pattern matching can be used for all tests. 
  //------------------------------------------------------------------------------------//
  
  InetListeningServersData *dataIn = (InetListeningServersData*)probeDataIn;

  pdVector resultVector;

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

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

  try {
    // Gather new data and add to resultVector
    GetNetstat(dataIn, &resultVector);

  }catch(...) {
    InetListeningServersData *tmp = new InetListeningServersData();
    tmp->SetTestId(dataIn->GetTestId());
    tmp->msg = "Error: An unknown error has ocured in InetListeningServersProbe::GetNetstat().";
    resultVector.push_back(tmp);      
  }

  return resultVector;
}

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

void InetListeningServersProbe::GetNetstat(InetListeningServersData *dataIn, pdVector *resultVector)
{
  //------------------------------------------------------------------------------------//
  //  ABSTRACT
  //  
  //  Gather the listening servers data that matches dataIn's program name
  //  Add each match to the result vector with the test id from dataIn
  //
  //------------------------------------------------------------------------------------//

  InetListeningServersData *tmp = NULL;
  int fd1[2];
  int fd2[2];
  int pid = 0;
  //char *buf = NULL;
  //char *errbuf = NULL;
  NetstatResult *netstatData = NULL;

  // Open communication pipes between processes
  if (pipe(fd1) < 0 || pipe(fd2) < 0) {
    tmp = new InetListeningServersData();
    tmp->SetTestId(dataIn->GetTestId());
    tmp->msg = "Error: (InetListeningServersProbe) Could not open pipe.";
    resultVector->push_back(tmp);
    return;
  }

  if ((pid = fork()) < 0) {
    tmp = new InetListeningServersData();
    tmp->SetTestId(dataIn->GetTestId());
    tmp->msg = "Error: (InetListeningServersProbe) fork error before running netstat -tuwlnpe.";
    resultVector->push_back(tmp);
    return;

    // Child process
  }else if (pid == 0) {

    // Close unnecessary pipes
    close (fd1[0]);
    close (fd2[0]);  

    ChildExecNetstat(fd1[1], fd2[1]);
    
    // Parent process
  } else {
	
    // Close unnecessary pipes
    close (fd1[1]);
    close (fd2[1]);

    // Get the results of netstat
    netstatData = ParentGetChildResult(fd1[0], fd2[0], pid);
    
    // Check the results of running netstat
    if(netstatData->text.compare("") == 0) {
      tmp = new InetListeningServersData();
      tmp->SetTestId(dataIn->GetTestId());
      tmp->msg = "Error: netstat returned no data.";

      if(netstatData->errText.compare("") == 0)
	tmp->msg.append(" " + netstatData->errText);

      resultVector->push_back(tmp);
      return;
    }
	  
    // Walk through the netstat output lines, parsing them.
    string tmpSTR;
    string PIDandProgramName;

    int curDelim;
    int preDelim;
    int curColon;
    int preColon;
    unsigned int curPos;
    unsigned int tmpPos;

    string Protocol;
    string LocalFullAddress;
    string LocalAddress;
    string LocalPort;
    string ForeignFullAddress;
    string ForeignAddress;
    string ForeignPort;
    string UserID;
    string PID;
    string ProgramName;
    bool skip;

    // Toss out first two lines
    curPos = netstatData->text.find("\n") + 1;
    curPos = netstatData->text.find("\n", curPos) + 1;
   
    // Get the first line end 
    tmpPos = netstatData->text.find("\n", curPos);

    // Parse the remaining lines.
    while (tmpPos != string::npos) {
     
      // Get the next string
      tmpSTR = netstatData->text.substr(curPos, (tmpPos - curPos));
      curPos = tmpPos + 1;
   
      curDelim = 0;
      preDelim = 0;

      ////////////////////////////////////////////
      ///// Protocol:  TCP or UDP ////////////////
      ////////////////////////////////////////////

      // Find the first space and grab the characters between preDelim and that point
      curDelim = tmpSTR.find(" ",preDelim);
      Protocol = tmpSTR.substr(preDelim,(curDelim - preDelim));

      // Now find the end of that whitespace
      preDelim = tmpSTR.find_first_not_of(" ",curDelim);

      ////////////////////////////////////////////
      ///// Skip two fields //////////////////////
      ////////////////////////////////////////////
    
      // Find the next space and then skip past it
      curDelim = tmpSTR.find(" ",preDelim);
      preDelim = tmpSTR.find_first_not_of(" ",curDelim); 

      // Do that again.
      curDelim = tmpSTR.find(" ",preDelim);
      preDelim = tmpSTR.find_first_not_of(" ",curDelim); 

      ////////////////////////////////////////////
      ///// LocalFullAddress //// ////////////////
      ////////////////////////////////////////////

      curDelim = tmpSTR.find(" ",preDelim);
      LocalFullAddress = tmpSTR.substr(preDelim,(curDelim - preDelim));
      preDelim = tmpSTR.find_first_not_of(" ",curDelim);

      ////////////////////////////////////////////
      ///// LocalAddress /////////////////////////
      ////////////////////////////////////////////

      preColon = 0;
      curColon = LocalFullAddress.find(":",preColon);
      LocalAddress = LocalFullAddress.substr(preColon,(curColon - preColon));

      ////////////////////////////////////////////
      ///// LocalPort ////////////////////////////
      ////////////////////////////////////////////

      preColon = curColon + 1;
      curColon = LocalFullAddress.length();
      LocalPort = LocalFullAddress.substr(preColon,(curColon - preColon));

      ////////////////////////////////////////////
      ///// ForeignFullAddress ///////////////////
      ////////////////////////////////////////////

      curDelim = tmpSTR.find(" ",preDelim);
      ForeignFullAddress = tmpSTR.substr(preDelim,(curDelim - preDelim));
      preDelim = tmpSTR.find_first_not_of(" ",curDelim);

      if((preDelim - curDelim) > 20)
		skip = false;
      else
		skip = true;

      ////////////////////////////////////////////
      ///// ForeignAddress ///////////////////////
      ////////////////////////////////////////////

      preColon = 0;
      curColon = ForeignFullAddress.find(":",preColon);
      ForeignAddress = ForeignFullAddress.substr(preColon,(curColon - preColon));

      ////////////////////////////////////////////
      ///// ForeignPort //////////////////////////
      ////////////////////////////////////////////

      preColon = curColon + 1;
      curColon = ForeignFullAddress.length();
      ForeignPort = ForeignFullAddress.substr(preColon,(curColon - preColon));

      ///////////////////////////////////////////
      ///// Skip one field //////////////////////
      ///////////////////////////////////////////
      if(skip) {
		curDelim = tmpSTR.find(" ",preDelim);
		preDelim = tmpSTR.find_first_not_of(" ",curDelim);
      }
      ///////////////////////////////////////////
      ///// UserID //////////////////////////////
      ///////////////////////////////////////////

      curDelim = tmpSTR.find(" ",preDelim);
      UserID = tmpSTR.substr(preDelim,(curDelim - preDelim));
      preDelim = tmpSTR.find_first_not_of(" ",curDelim);

      ///////////////////////////////////////////
      ///// Skip one field //////////////////////
      ///////////////////////////////////////////
    
      curDelim = tmpSTR.find(" ",preDelim);
      preDelim = tmpSTR.find_first_not_of(" ",curDelim);

      ///////////////////////////////////////////
      ///// PID/ProgramName /////////////////////
      ///////////////////////////////////////////

      curDelim = tmpSTR.find(" ",preDelim);
      PIDandProgramName = tmpSTR.substr(preDelim,(curDelim - preDelim));
      preDelim = tmpSTR.find_first_not_of(" ",curDelim);
   
      ////////////////////////////////////////////
      ///// PID  /////////////////////////////////
      ////////////////////////////////////////////

      preColon = 0;
      curColon = PIDandProgramName.find("/",preColon);
      PID = PIDandProgramName.substr(preColon,(curColon - preColon));
        
      ////////////////////////////////////////////
      ///// ProgramName //////////////////////////
      ////////////////////////////////////////////

      preColon = curColon + 1;
      curColon = PIDandProgramName.length();
      ProgramName = PIDandProgramName.substr(preColon,(curColon - preColon));

      ////////////////////////////////////////////
      //////////////////////////////////////////// 


      // Check if the program name is a match
      if(myMatcher->IsMatch(dataIn->program_name->data.c_str(), ProgramName.c_str())) {
		tmp = new InetListeningServersData();
		tmp->SetTestId(dataIn->GetTestId());
		tmp->local_full_address = LocalFullAddress;
		tmp->local_address = LocalAddress;
		tmp->local_port = LocalPort;
		tmp->foreign_full_address = ForeignFullAddress;
		tmp->foreign_address = ForeignAddress;
		tmp->foreign_port = ForeignPort;
		tmp->user_id = UserID;
		tmp->pid = PID;
		tmp->program_name->data = ProgramName;
		tmp->protocol = Protocol;
		
		if(netstatData->errText.compare("") != 0)
		  tmp->msg = "Error: " + netstatData->errText;
		
		resultVector->push_back(tmp);
      }

      // Get the next line end 
      tmpPos = netstatData->text.find("\n", curPos);  
    }
    
    // Ensure that any messages are returned if no matches were found
    // Make sure that the result vector has at least one item
    if(resultVector->size() == 0) {
      InetListeningServersData *tmp = new InetListeningServersData();
      tmp->SetTestId(dataIn->GetTestId());
      tmp->msg = "No matching programs found";

      if(netstatData->errText.compare("") != 0)
		tmp->msg.append(" Error: " + netstatData->errText);   

      resultVector->push_back(tmp);
    }    
  }
}

void InetListeningServersProbe::ChildExecNetstat(int writeErrh, int writeh)
{
  //------------------------------------------------------------------------------------//
  //  ABSTRACT
  //  
  //  Redirect stdout and stderr to the provided pipes (writeh, and writeErrh). 
  //  Execute nestate with the correct options.
  //
  //------------------------------------------------------------------------------------//

  // Point STDOUT and STDERR of child at pipe.  When exec program, output and
  // all error messages will be sent down pipe instead of to screen.
  if (writeh != STDOUT_FILENO) {
    if (dup2(writeh, STDOUT_FILENO) != STDOUT_FILENO)
      exit(-1);
  }

  if (writeErrh != STDERR_FILENO) {
    if (dup2(writeErrh, STDERR_FILENO) != STDERR_FILENO)
      exit(-1);
  }

  // Output redirected (duplicated), no longer need pipe
  close (writeh);
  close (writeErrh);

  // Execute the command
  execl("/bin/netstat", "netstat", "-tuwlnpe", NULL);

  exit(0);
} 

NetstatResult* InetListeningServersProbe::ParentGetChildResult(int readErrh, int readh, int pid)//, char* buf, char* errbuf)
{
  //------------------------------------------------------------------------------------//
  //  ABSTRACT
  //  
  //  Read readErrh and readh until there is no more data to be read. Wait for the 
  //  child process to complere. Return a NetstatResult object with the data from
  //  netstat.
  //
  //------------------------------------------------------------------------------------//

  NetstatResult *result = new NetstatResult("", "");
  int bytes = 0;
  int maxFDS = 0;
  char *buf = NULL;
  fd_set readfds;
  bool errComplete = false;
  bool stdComplete = false;

  // Allocate memory for  buf
  buf = (char*)malloc(1024);
  if(buf == NULL) {

    // Wait for the child process to compelete
    waitpid (pid, NULL, 0);
      
    // Close the pipes
    close (readh);
    close (readErrh);

    // Set an error message
    result->errText.append("Error: unable to allocate memory to read netstat data into.");
    
    return result;
  }

  // Init the maxFDS value
  if(readh >= readErrh) {
    maxFDS = readh + 1;
  }else {
    maxFDS = readErrh + 1;
  }

  // Loop over the call to select without using a timmer
  // Only stop looping when select fails. select will
  // fail when the file descriptors are closed by the 
  // child process.
  while(!errComplete || !stdComplete) {

    // Reset the fd_set
    FD_ZERO(&readfds);
    FD_SET(readErrh, &readfds);
    FD_SET(readh, &readfds);

    if(select(maxFDS, &readfds, NULL, NULL, NULL) != -1) {
      if(FD_ISSET(readErrh, &readfds)) {
	// Read some error output from command.  
	memset(buf, '\0', 1024);
	bytes = read(readErrh, buf, 1023);
	result->errText.append(buf);

	if(bytes == 0) 
	  errComplete = true;
      }

      if(FD_ISSET(readh, &readfds)) { 
	// Read allsome std output from command. 
	memset(buf, '\0', 1024);
	bytes = read(readh, buf, 1023);
	result->text.append(buf);

	if(bytes == 0)
	  stdComplete = true;
      }
	  
    }else {	    
      break;
    }
  }
  
  // Wait for the child process to compelete
  if(waitpid (pid, NULL, 0) == -1) {
    result->errText.append("Execution of netstat in child process failed.");
    return result;
  }

  // Close the pipes
  close (readh);
  close (readErrh);

  return result;
}



  
