#!/usr/bin/ruby -w
#
# acoc - Arbitrary Command Output Colourer
#
# $Id: acoc,v 1.22 2003/06/21 10:27:01 ianmacd Exp $
#
# Version : 0.2.5
# 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.
= 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
* /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 home page - http://raa.ruby-lang.org/list.rhtml?name=ansicolor"|URL:http://raa.ruby-lang.org/list.rhtml?name=ansicolor>))
= 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'

include Term::ANSIColor

# 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(["/usr/local/etc/acoc.conf", "/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   = Hash.new
  @flags = Hash.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 name
	    progs = $1.scan(/[^,\s]+/)
	    progs.each do |prog|
	      name, flags  = prog.split(%r(/))
	      @cmd[name] ||= Array.new
	      @flags[name] ||= ""
	      @flags[name] += flags unless flags.nil?
	      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
	  regex = Regexp.new(regex)
	  progs.each { |prog| @cmd[prog] << { regex => [colours, flags] } }
	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)
  begin
    exec(*args)
  rescue Errno::ENOENT => reason
    # can't find the program we're supposed to run
    $stderr.puts reason
    exit Errno::ENOENT::Errno
  end
end

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

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 without reinvoking the symlink
end

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

# if there's no config section for this command and no 'default' section,
# simply execute it normally
run(*ARGV) unless (@cmd.include?(prog = File.basename(ARGV[0])) ||
		   @cmd.include?('default'))

# use default section if no program-specific section available
prog = 'default' unless @cmd.include? prog

# if STDOUT is not a tty and the 't' flag is not specified, simply
# execute the program normally
run(*ARGV) unless $stdout.tty? || @flags[prog].include?('t')

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

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

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

# execute command
IO.popen(ARGV.join(' ')) do |f|
  while line = f.gets
    matched = false

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

      # get a pattern and attribute set pairing for this command
      pair.each do |regex, attrs|

	# split attribute set into colours and flags
	colours, flags = attrs

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

	    # perform global substitution
	    line.gsub!(regex) do |match|
	      index = [matches, colours.size - 1].min
	      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, colours.size - 1].min
	      ansi_offset = 0
	      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
    end

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

  end
end
