#!/usr/bin/ruby -w
#
# acoc - Arbitrary Command Output Colourer
#
# $Id: acoc,v 1.52 2003/09/12 16:16:15 ianmacd Exp $
#
# Version : 0.4.6
# Author  : Ian Macdonald <ian@caliban.org>
# 
# Copyright (C) 2003 Ian Macdonald
# 
#   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 2, 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 should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software Foundation,
#   Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

=begin

= NAME
acoc - arbitrary command output colourer
= SYNOPSIS
 acoc command [arg1 .. argN]
 acoc -h|--help|-v|--version
= DESCRIPTION
((*acoc*)) is a regular expression based colour formatter for programs that
display output on the command-line. It works as a wrapper around the target
program, executing it and capturing the stdout stream. Optionally, stderr can
be redirected to stdout, so that it, too, can be manipulated.

((*acoc*)) then applies matching rules to patterns in the output and applies
colour sets to those matches. If the ((|$ACOC|)) environment variable is set
to 'none', ((*acoc*)) will not perform any colouring.
= OPTIONS
: -h or --help
  Display usage information.
: -v or --version
  Display version information.
= AUTHOR
Written by Ian Macdonald <ian@caliban.org>
= COPYRIGHT
 Copyright (C) 2003 Ian Macdonald

 This is free software; see the source for copying conditions.
 There is NO warranty; not even for MERCHANTABILITY or FITNESS
 FOR A PARTICULAR PURPOSE.
= FILES
* /usr/local/etc/acoc.conf /etc/acoc.conf ~/.acoc.conf
= CONTRIBUTING
acoc is only as good as the configuration file that it uses. If you compose
pattern-matching rules that you think would be useful to other people, please
send them to me for inclusion in a subsequent release.
= SEE ALSO
* acoc.conf(5)
* ((<"acoc home page - http://www.caliban.org/ruby/"|URL:http://www.caliban.org/ruby/>))
* ((<"Term::ANSIColor - http://raa.ruby-lang.org/list.rhtml?name=ansicolor"|URL:http://raa.ruby-lang.org/list.rhtml?name=ansicolor>))
* ((<"Ruby/TPty - http://www.tmtm.org/ruby/tpty/"|URL:http://www.tmtm.org/ruby/tpty/>))
= BUGS
* Nested regular expressions do not work well. Inner subexpressions need to use clustering (?:), not capturing (). In other words, they can be used for matching, but not for colouring.

=end

require 'English'
require 'term/ansicolor'
begin
  require 'tpty'
rescue LoadError
end

PROGRAM_NAME = File::basename($0)
PROGRAM_VERSION = '0.4.6'

include Term::ANSIColor

class Config < Hash; end

class Program
  attr_accessor :flags, :specs

  def initialize(flags)
    @flags = flags || ""
    @specs = Array.new
  end
end

class Rule
  attr_reader :regex, :flags, :colours

  def initialize(regex, flags, colours)
    @regex   = regex
    @flags   = flags
    @colours = colours
  end
end

# set things up
#
def initialise
  # Queen's or Dubya's English?
  if ENV['LANG'] == "en_US" || ENV['LC_ALL'] == "en_US"
    @colour = "color"
  else
    @colour = "colour"
  end

  if parse_config("/etc/acoc.conf", "/usr/local/etc/acoc.conf",
		  ENV['HOME'] + "/.acoc.conf") == 0
    $stderr.puts "No readable config files found."
    exit 1
  end
end

# display usage message and exit
#
def usage(code = 0)
  $stderr.puts <<EOF
Usage: #{PROGRAM_NAME} command [arg1 .. argN]
       #{PROGRAM_NAME} [-h|--help|-v|--version]
EOF

  exit code
end

# display version and copyright message, then exit
#
def version
  $stderr.puts <<EOF
#{PROGRAM_NAME} #{PROGRAM_VERSION}

Copyright 2003 Ian Macdonald <ian@caliban.org>
This is free software; see the source for copying conditions.
There is NO warranty; not even for MERCHANTABILITY or FITNESS
FOR A PARTICULAR PURPOSE, to the extent permitted by law.
EOF

  exit
end

# get configuration data
#
def parse_config(*files)
  @cmd   = Config.new
  parsed = 0

  files.each do |file|
    next unless FileTest::file?(file) && FileTest::readable?(file)

    begin
      f = File.open(file) do |f|
	while line = f.gets do
	  next if line =~ /^(#|$)/     # skip blank lines and comments

	  if line =~ /^\[([^\]]+)\]$/  # start of program section
	    # get program invocation
	    progs = $1.split(/\s*,\s*/)
	    progs.each do |prog|
	      invocation, flags = prog.split(%r(/))

	      if ! flags.nil? && flags.include?('r')
	        # remove matching entries for this program
		program = invocation.sub(/\s.*$/, '')
		@cmd.each_key do |key|
		  @cmd.delete(key) if key =~ /^#{program}\b/o
		end
		flags.delete 'r'
	      end

	      # create entry for this program
	      if @cmd.has_key?(invocation)
		@cmd[invocation].flags += flags unless flags.nil?
	      else
	        @cmd[invocation] = Program.new(flags)
	      end
	      prog.sub!(%r(/\w+), '')
	    end
	    next
	  end

	  begin
	    regex, flags, colours =
	      /^(.)([^\1]*)\1(g?)\s+(.*)/.match(line)[2..4]
	  rescue
	    $stderr.puts "Ignoring bad config line #{$NR}: #{line}"
	  end

	  colours = colours.split(/\s*,\s*/)
	  colours.join(' ').split(/[+\s]+/).each do |colour|
	    raise "#{colour} is not a supported #{@colour}" \
	      unless attributes.collect { |a| a.to_s }.include? colour
	  end

	  progs.each do |prog|
	    @cmd[prog].specs << Rule.new(Regexp.new(regex), flags, colours)
	  end
	end
      end
    rescue Errno::ENOENT
      $stderr.puts "Failed to open config file: #{$ERROR_INFO}"
      exit 1
    rescue
      $stderr.puts "Error while parsing config file #{file} @ line #{$NR}: #{$ERROR_INFO}"
      exit 2
    end

    parsed += 1
  end

  if $DEBUG
    $stderr.printf("Action data: %s\n", @cmd.inspect)
    $stderr.printf("Flag data: %s\n", @flags.inspect)
  end

  parsed
end

# make sure terminal is never left in a coloured state
#
def ignore_signal(signals)
  signals.each do |signal|
    trap(signal) { print reset }
  end
end

def run(args)
  exec(*args)
rescue Errno::ENOENT => reason
  # can't find the program we're supposed to run
  $stderr.puts reason
  exit Errno::ENOENT::Errno
end

# match and colour an individual line
#
def colour_line(prog, line)
  matched = false

  # act on only the first match unless the /a flag was given
  break if matched && ! @cmd[prog].flags.include?('a')

  # get a pattern and attribute set pairing for this command
  @cmd[prog].specs.each do |spec|

    if r = spec.regex.match(line)  # line matches this regex
      matched = true
      if spec.flags.include? 'g'   # global flag
	matches = 0

	# perform global substitution
	line.gsub!(spec.regex) do |match|
	  index = [matches, spec.colours.size - 1].min
	  spec.colours[index].split(/[+\s]+/).each do |colour|
	    match = match.send(colour)
	  end
	  matches += 1
	  match
	end

      else  # colour each match separately
	# work from right to left, bracketing each match
	(r.size - 1).downto(1) do |i|
	  start  = r.begin(i)
	  length = r.end(i) - start
	  index  = [i - 1, spec.colours.size - 1].min
	  ansi_offset = 0
	  spec.colours[index].split(/[+\s]+/).each do |colour|
	    line[start + ansi_offset, length] =
	      line[start + ansi_offset, length].send(colour)
	    # when applying multiple colours, we apply them one at a
	    # time, so we need to compensate for the start of the string
	    # moving to the right as the colour codes are applied
	    ansi_offset += send(colour).length
	  end
	end
      end
    end
  end

  line
end

# process program output, one line at a time
#
def colour(prog, *cmd_line)
  block = proc do |f|
    while ! f.eof?
      begin
	line = f.gets
      rescue  # why do we need rescue here?
	exit  # why the Errno::EIO when running ls(1)?
      end

      coloured_line = colour_line(prog, line)

      begin
	print coloured_line
      rescue Errno::EPIPE => reason   # catch broken pipes
	$stderr.puts reason
	exit Errno::EPIPE::Errno
      end

    end
  end

  # prepare command line: requote each argument for the shell
  cmd_line = "'" << cmd_line.join(%q(' ')) << "'"

  # redirect stderr to stdout if /e flag given
  cmd_line << " 2>&1" if @cmd[prog].flags.include? 'e'

  # make sure we don't buffer output when stdout is connected to a pipe
  $stdout.sync = true

  # install signal handler
  ignore_signal(%w(HUP INT QUIT))

  if @cmd[prog].flags.include?('p') && $LOADED_FEATURES.include?('tpty.so')
    # allocate program a pseudo-terminal and run through that
    pty = TPty.new do |s,|
      fork do
	# redirect child streams to slave
	STDIN.reopen(s)
	STDOUT.reopen(s)
	#STDERR.reopen(s)
	s.close
	run(cmd_line)
      end
    end

    # no buffering on pty
    # pty.master.sync = true
    block.call(pty.master)
  else
    # execute command
    IO.popen(cmd_line) { |f| block.call(f) }
  end
end

initialise

if File.lstat($0).symlink?  # we're being invoked via a symlink
  # remove symlink's directory from PATH
  ENV['PATH'] = ENV['PATH'].sub(/#{File.dirname($0)}:?/, '')

  # prefix command line with symlink's name
  ARGV.unshift PROGRAM_NAME # ARGV can now be either exec'ed or popen'ed
end

usage if ARGV.empty? || %w[-h --help].include?(ARGV[0])
version if %w[-v --version].include?(ARGV[0])

# sort the keys to ensure we find the longest (i.e. most specific) match for
# the command line, e.g. a config section for [ps ax] will match before one
# for [ps a]
prog = nil
@cmd.keys.sort { |a,b| b.length <=> a.length }.each do |key|
  if ARGV.join(' ').index(key) == 0
    prog = key
    break
  end
end

# if there's no config section for this command and no 'default' section,
# simply execute it normally. Do same if $ACOC set to 'none'.
run(ARGV) if ENV['ACOC'] == 'none' || ! (prog || @cmd.include?('default'))

# use default section if no program-specific section available
prog ||= 'default'

# if there's a config section for the command, but no rules to accompany it,
# simply execute it normally. Likewise if STDOUT is not a tty and the 't' flag
# is not specified
run(ARGV) if @cmd[prog].specs.empty? ||
	     ! ($stdout.tty? || @cmd[prog].flags.include?('t'))

# colour program output
colour(prog, ARGV)
