/* Copyright 2008 Peter Hercek, All rights reserved.
 * This program 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 3 of the License, or (at your option)
 * any later version. This program 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 can receive a copy of the GNU General
 * Public License at http://www.gnu.org/licenses/. */
#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>

#define bool int
#define false 0
#define true 1

/* Reads from fd, prints to stdout with optional bgnStr/endStr decoration.
 * If input descriptor fd is at EOF then returns false (otherwise true). */
bool in2out( int fd, char const* bgnStr, char const* endStr ) {
  static unsigned char buffer[512];
  ssize_t rv;

  while ( 0 < ( rv = read( fd , buffer, sizeof buffer ) ) ) {
    if ( bgnStr )
    if ( EOF == fputs( bgnStr, stdout ) ) {
        perror("colorize: fputs"); _exit(1); }
    if ( 1 != fwrite( buffer, rv, 1, stdout ) ) {
      perror("colorize: fwrite"); _exit(1); }
    if ( endStr )
    if ( EOF == fputs( endStr, stdout ) ) {
        perror("colorize: fputs"); _exit(1); }
  } /* while */
  if ( 0 == rv )
    return true; /*  we reached eof */
  if ( -1 == rv && EAGAIN != errno ) {
    perror("colorize: read"); _exit(1); }
  return false;
} /* in2out() */


int main( int ac, char* av[] ) {
  char const* const bgnStr = ac>1?av[1]:NULL;
  char const* const endStr = ac>2?av[2]:NULL;
  char const* const pipeName = ac>3?av[3]:NULL;
  int pipeFd = -1;
  int fdCount = 1;
  fd_set readFdSet;
  FD_ZERO( &readFdSet );

  /* Check of arguments. */
  if (ac != 3 && ac != 4) {
    fprintf( stderr, "Usage: colorize <bgnStr> <endStr> [ <pipeName> ]\n");
    fprintf( stderr, "  <bgnStr>   string to insert before any data read from stdin\n" );
    fprintf( stderr, "  <endStr>   string to insert after any data read from stdin\n" );
    fprintf( stderr, "  <pipeName> file name to use for creation of a named pipe from which data\n" );
    fprintf( stderr, "             will be read and copied to stdout without any decoration\n" );
    fprintf( stderr, "Note: The example at the end shows a typicall usage. When the last argument\n" );
    fprintf( stderr, "(<pipeName>) is added then the best synchronization between stdout and\n" );
    fprintf( stderr, "stderr can be achieved by redirecting the shell stdout (and that means also\n" );
    fprintf( stderr, "the stdout of all its child processes) to the named pipe. But it is not a\n" );
    fprintf( stderr, "good tradeoff in most cases since some programs check whether their stdout\n" );
    fprintf( stderr, "is a tty and if yes they use some nice colors, become interactive, ...\n" );
    fprintf( stderr, "If you have programs which interleave stderr and stdout a lot and the order\n" );
    fprintf( stderr, "gets messed up then use explicit redirection of just their stdout to the\n" );
    fprintf( stderr, "named pipe. Use e.g. shell aliases to achieve this easily. If your shell\n" );
    fprintf( stderr, "prompt is often shown before the last program stdout/stderr then add a\n" );
    fprintf( stderr, "short delay to your prompt string (e.g. PS1='`sleep 0.1`> '). Colorize\n" );
    fprintf( stderr, "also tries to make its priority higher so that the redirected stderr is\n" );
    fprintf( stderr, "quicker. This alone is often enough to achieve a good synchornization\n" );
    fprintf( stderr, "of stderr with stdout but to allow it you may need to change the program\n" );
    fprintf( stderr, "owner to root and allow setuid.\n" );
    fprintf( stderr, "Example: exec 2>> ( colorize `tput setaf 1` `tput sgr0` >/dev/tty & )\n" );
    _exit(2);
  }

  nice(-15); /* We do not want to delay program output unnecessarily. */
  if ( -1 == setuid( getuid() ) ) {
      perror("colorize: setuid"); _exit(1); }

  /* Setup local variables related to pipe handling provided that pipe creation is requested. */
  if ( pipeName ) {
    if ( -1 == mkfifo( pipeName, 0600 ) ) {
      perror("colorize: mkfifo"); _exit(1); }
    if ( -1 == ( pipeFd = open( pipeName, O_RDONLY|O_NDELAY ) ) ) {
      perror("colorize: open"); _exit(1); }
    fdCount = 1+pipeFd;
  }

  /* Set stdin to nonblocking mode. */
  { int rv;
    if ( -1 == ( rv = fcntl( 0, F_GETFL ) ) ) {
      perror("colorize: fcntl(GETFL)"); _exit(1); }
    if ( -1 == fcntl( 0, F_SETFL, rv|O_NDELAY ) ) {
      perror("colorize: fcntl(SETFL)"); _exit(1); }
  }

  /* Main reactor loop (dispatches on stdin, and pipe (if any)). */
  for (;;) {
    FD_SET( 0, &readFdSet );
    if ( 0 <= pipeFd ) {
      FD_SET( pipeFd, &readFdSet ); }

    if ( -1 == select( fdCount, &readFdSet, NULL, NULL, NULL ) ) {
      perror("colorize: select"); _exit(1); }

    if ( FD_ISSET( 0, &readFdSet ) )
      if ( in2out( 0, bgnStr, endStr ) )
        return 0; /* EOF reached on stdin -> exit */

    if ( 0 <= pipeFd && FD_ISSET( pipeFd, &readFdSet ) )
      if ( in2out( pipeFd, NULL, NULL ) ) {
        if ( -1 == close( pipeFd ) ) {
          perror("colorize: close"); _exit(1); }
        if ( -1 == ( pipeFd = open( pipeName, O_RDONLY|O_NDELAY ) ) ) {
          perror("colorize: (re)open"); _exit(1); }
        fdCount = 1+pipeFd;
      }

    fflush(stdout);
  } /* forever */

  return 0;
} /* main() */