#!/usr/bin/ruby -w
#
# acoc - Arbitrary Command Output Colourer
#
# $Id: acoc,v 1.38 2003/07/08 06:57:48 ianmacd Exp $
#
# Version : 0.4.1
# 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 - 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.1'

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

def colour(prog, *cmd_line)
  block = proc do |f|
    while ! f.eof?
      # don't know why the next line causes a Errno::EIO on ls(1), so
      # I'll just rescue it until I work it out. eof? seems to fail to
      # return 'true' when it should, so we'll just silently exit for now
      line = f.gets rescue exit
      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

      begin
	print line
	# catch broken pipes
      rescue Errno::EPIPE => reason
	$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 STOP))

  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
	Process.setsid
	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
run(ARGV) unless 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)
