/* Programmer: David Mullen
 * File: main.c
 * Course: CS 284
 * Section: A
 * Instructor: Matt Johnson
 * Due: 29 Jun 2005
 * Description: Allow execution of arbitrary commands. */

#define _GNU_SOURCE
#include <stdio.h>
#include <getopt.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <errno.h>
#include <ctype.h>

/* PreConditions:  The size is greater than zero.
 * PostConditions: If the function returns, the memory was successfully
 *                 allocated. Otherwise, an error message is printed
 *                 and the program exits. */
void *SafeMalloc(int size);

/* PreConditions:  The size is greater than zero.
 * PostConditions: If the function returns, the memory was successfully
 *                 reallocated. Otherwise, an error message is printed
 *                 and the program exits. */
void *SafeRealloc(void *p, int size);

/* PreConditions:  The string is a sequence of tokens separated by
 *                 various characters
 * PostConditions: The tokens are extracted from the string and 
 *                 collected into a vector.  The last element of the
 *                 vector is then set to NULL to indicate the end of
 *                 the tokens.  The function returns upon seeing either
 *                 a NUL character or one of the given terminators; in
 *                 either case, stringPtr is set to the address of the
 *                 stopping character. */
char **ParseTokens(const char **stringPtr, const char *separators,
                   const char *terminators);

/* PreConditions:  The tokens pointer is a vector of tokens, with the last
 *                 element set to NULL.
 * PostConditions: All storage used by the token vector has been freed. */
void FreeTokens(char **tokens);

/* PreConditions:  The src string is the pathname of the file to be copied.
 * PostConditions: A copy of the file has been written, the name of which
 *                 is src with ".cpy" appended to it. */
void CopyFile(const char *src);

/* PreConditions:  The command to execute may contain tokens, whitespace,
 *                 and pipes.  The inPipeFDs pointer may be NULL, in which
 *                 case this is the first program in the pipeline (i.e., 
 *                 its stdin is not remapped).
 * PostConditions: If the first token in the command names an existing
 *                 program, then it is ran with the arguments following
 *                 the command name.  If the first command is followed
 *                 by a pipe, then the output if the first command is piped
 *                 to the input of the second, and so on.  This function
 *                 is also responsible for closing inPipeFDs. */
void RunProgram(const char *command, int *inPipeFDs);

static struct option long_options[] =
{
   {"pwd", 0, NULL, 'p'},
   {"lastchar", 1, NULL, 'l'},
   {"help", 0, NULL, 'h'},
   {NULL, 0, NULL, 0}
};

int main(int argc, char **argv)
{
   char *line = NULL;
   size_t lineSize = 0;
   char *hostName;
   char wd[256];
   char *p;
   int lastChar = '>';
   int pwd = 0;
   int c;

   while ((c = getopt_long(argc, argv, "pl:h", long_options, NULL)) != -1)
   {
      switch (c)
      {
         case 'p':
            pwd = 1;
            break;

         case 'l':
            lastChar = optarg[0];
            break;

         case 'h':
            puts ("Synopsis: David's Shell, version 4.");
            puts ("Usage: driver [[option] [option_arg]]");
            puts ("   -h, --help        prints this help message");
            puts ("   -l, --lastchar    requires a one character argument"
                  " which becomes the last");
            puts ("                     character of the prompt");
            puts ("   -p, --pwd         tells the shell to include the"
                  " current working directory in");
            puts ("                     the prompt");
            return 0;
      }
   }

   /* Only bother fetching the working directory if we're going to
    * be displaying it. */
   if (pwd)
   {
      getcwd(wd, sizeof(wd));
   }

   hostName = getenv("HOSTNAME");
   if (!hostName)
   {
      hostName = "localhost";
   }

   while (1)
   {
      /* Print the prompt, which includes the machine name and (optionally)
       * the current working directory. */
      printf("%s", hostName);
      if (pwd)
      {
         printf(":%s", wd);
      }
      printf(":%c", lastChar);

      /* Exit on end-of-file (ctrl-D). */
      if (getline(&line, &lineSize, stdin) == -1)
      {
         printf ("\n\nGoodbye.\n\n");
         break;
      }

      /* Remove newline character. */
      p = strchr(line, '\n');
      if (p)
      {
         *p = '\0';
      }

      /* Dispatch on first word of command line. */
      if (strncmp(line, "exit", 4) == 0)
      {
         printf ("\nGoodbye.\n\n");
         break;
      }
      else if (strncmp(line, "cpy ", 4) == 0)
      {
         CopyFile(line + 4);
      }
      else if (strncmp(line, "cd ", 3) == 0)
      {
         if (chdir(line + 3) != 0)
         {
            fprintf(stderr, "%s: %s\n", line + 3, strerror(errno));
         }
         if (pwd)
         {
            getcwd(wd, sizeof(wd));
            setenv("PWD", wd, 1);
         }
      }
      else if (strncmp(line, "md ", 3) == 0)
      {
         if (mkdir(line + 3, 0777) != 0)
         {
            fprintf(stderr, "%s: %s\n", line + 3, strerror(errno));
         }
      }
      else if (strncmp(line, "rd ", 3) == 0)
      {
         if (rmdir(line + 3) != 0)
         {
            fprintf(stderr, "%s: %s\n", line + 3, strerror(errno));
         }
      }
      else
      {
         RunProgram(line, NULL);
      }
   }

   /* Free storage used by getline. */
   if (line)
   {
      free(line);
   }

   return 0;
}

void *SafeMalloc(int size)
{
   void *p = malloc(size);
   if (!p)
   {
      fprintf(stderr, "SafeMalloc: %s\n", strerror(errno));
      exit(1);
   }
   return p;
}

void *SafeRealloc(void *p, int size)
{
   if (p)
   {
      p = realloc(p, size);
   }
   else
   {
      p = malloc(size);
   }
   if (!p)
   {
      fprintf(stderr, "SafeRealloc: %s\n", strerror(errno));
      exit(1);
   }
   return p;
}

char **ParseTokens(const char **stringPtr, const char *separators,
                   const char *terminators)
{
   const char *p = *stringPtr;
   const char *q;
   int numTokens = 0;
   int numTokensMax = 1;
   char **tokens = NULL;
   int len;

   while (*p && !strchr(terminators, *p))
   {
      if (!strchr(separators, *p))
      {
         /* Scan to the end of the token. */
         for (q = p;
              *q && !strchr(separators, *q) && !strchr(terminators, *q);
              q++);

         /* Expand the token vector if necessary.  We compare against
          * one less than numTokensMax so that there will be room
          * for the NULL at the end of the vector. */
         if (!tokens || (numTokens == numTokensMax - 1))
         {
            numTokensMax *= 2;
            tokens = SafeRealloc(tokens, sizeof(char **) * numTokensMax);
         }

         /* Copy the token and append it to the vector. */
         len = q - p;
         tokens[numTokens] = SafeMalloc(len + 1);
         memcpy(tokens[numTokens], p, len);
         tokens[numTokens][len] = '\0';
         numTokens++;
         p = q;
      }
      else
      {
         p++;
      }
   }

   *stringPtr = p;
   if (tokens)
   {
      tokens[numTokens] = NULL;
   }
   return tokens;
}

void FreeTokens(char **tokens)
{
   int i;
   for (i = 0; tokens[i]; i++)
   {
      free(tokens[i]);
   }
   free(tokens);
}

void CopyFile(const char *src)
{
   char dest[256];
   int inFD, outFD;
   char buf[BUFSIZ];
   int n;

   /* Open file to be copied for reading. */
   if ((inFD = open(src, O_RDONLY)) < 0)
   {
      fprintf(stderr, "%s: %s\n", src, strerror(errno));
      return;
   }

   /* Create a new file with ".cpy" appended to the name of the
    * original file. */
   snprintf(dest, sizeof(dest), "%s.cpy", src);
   if ((outFD = open(dest, O_WRONLY | O_CREAT, 0666)) < 0)
   {
      fprintf(stderr, "%s: %s\n", dest, strerror(errno));
      close(inFD);
      return;
   }

   /* Read until end-of-file or error. */
   while ((n = read(inFD, buf, sizeof(buf))) > 0)
   {
      if (write(outFD, buf, n) < 0)
      {
         fprintf(stderr, "%s: %s\n", dest, strerror(errno));
      }
   }

   /* Report any problems in reading. */
   if (n < 0)
   {
      fprintf(stderr, "%s: %s\n", src, strerror(errno));
   }

   close(outFD);
   close(inFD);
}

void RunProgram(const char *command, int *inPipeFDs)
{
   char **argv;
   pid_t pid;
   int outPipeFDs[2];

   /* Get the arguments (including the name of the program), stopping
    * at any special character. */
   if (!(argv = ParseTokens(&command, " \t\n", "|&;<>!@")))
   {
      if (inPipeFDs)
      {
         fprintf(stderr, "Invalid pipeline\n");
         exit(1);
      }
      else
      {
         return;
      }
   }

   /* If the command is followed by a pipe, we need to create one for
    * standard output and continue the pipeline. */
   if (*command == '|')
   {
      if (pipe(outPipeFDs) != 0)
      {
         fprintf(stderr, "pipe: %s\n", strerror(errno));
         exit(1);
      }
   }
   
   /* Fork a new process that will execute the command. */
   pid = fork();
   if (pid < 0)
   {
      fprintf(stderr, "fork: %s\n", strerror(errno));
      exit(1);
   }
   else if (pid > 0)
   {
      /* Free arguments since they are needed only by the child. */
      FreeTokens(argv);

      /* As per the postconditions, we are responsible for closing the
       * input pipe FDs. */
      if (inPipeFDs)
      {
         close(inPipeFDs[0]);
         close(inPipeFDs[1]);
      }

      /* If this is a pipeline, recurse on the remainder, passing the current
       * output pipe as the input to the next program in the chain. */
      if (*command == '|')
      {
         RunProgram(command + 1, outPipeFDs);
      }

      /* Wait for the child to finish executing the command. */
      wait(NULL);
   }
   else
   {
      /* Connect the given pipe to standard input. */
      if (inPipeFDs)
      {
         dup2(inPipeFDs[0], 0);
         close(inPipeFDs[0]);
         close(inPipeFDs[1]);
      }

      /* If the pipeline continues, connect standard output to the pipe
       * we created prior to forking. */
      if (*command == '|')
      {
         dup2(outPipeFDs[1], 1);
         close(outPipeFDs[0]);
         close(outPipeFDs[1]);
      }

      /* Try to execute the program with the collected arguments. */
      execvp(argv[0], argv);
      fprintf(stderr, "%s: %s\n", argv[0], strerror(errno)); 
      _exit((errno == ENOENT) ? 127 : 126);
   }
}