#!/usr/bin/python

# Joe Gillotti - 5/16/2014

import select
import subprocess
import pty
import os
import sys
from getpass import getuser
from termcolor import colored

from optparse import OptionParser

class tailer:
  
  def __init__(self):
    self.read_fds = {}
    self.avail_colors = ['red', 'green', 'yellow', 'blue', 'magenta', 'cyan']
    self.box_colors = {}
    self.last_box = None

  def getColor(self, box):
    if box in self.box_colors:
      return self.box_colors[box]
    
    if len(self.avail_colors) == 0:
      return None

    self.box_colors[box] = self.avail_colors.pop()
    return self.box_colors[box]

  def showLines(self, data, box):

    if self.last_box != None and box != self.last_box:
      print '--'

    color = self.getColor(box)

    if color:
      box = colored(box, color, attrs=['bold'])
    else:
      box = colored(box, attrs=['bold'])
   
    print '\n'.join(['[%s] %s' % (box, line.strip()) for line in data.split('\n') if len(line)])
      

  def addBox(self, box, user, port, path):
    master_fd, slave_fd = pty.openpty()

    cmd = ['ssh', '-4C', '-l', user, '-p', port, '-oStrictHostKeyChecking=no', '-oBatchMode=yes', box, 'tail', '-f', path]

    proc = subprocess.Popen(cmd, stdout=slave_fd, stderr=subprocess.STDOUT, shell=False)

    if proc.poll() is not None:
      print 'proc for %s died prematurely..' % box
      return

    self.read_fds[master_fd] = {
      'box': box,
      'proc': proc
    }
    
  def tail(self):
    last_box = None

    try:
      while True:
          ready, _, _ = select.select(self.read_fds.keys(), [], [], .04)

          if ready:
            for fd in ready:

              data = os.read(fd, 2048)
              if not data:
                  continue
              
              self.showLines(data, self.read_fds[fd]['box'])

              self.last_box = self.read_fds[fd]['box']

              if self.read_fds[fd]['proc'].poll() is not None:
                del(self.read_fds[fd])
    
    except KeyboardInterrupt:
      self.shutdown()
          

  def shutdown(self):
    for proc in self.read_fds.values():
      proc['proc'].kill()


parser = OptionParser()

parser.add_option("--hosts", dest="hosts",
  help="Comma separated list of hosts")

parser.add_option("--login", '-l', dest="login",
  default=getuser(),
  help="SSH user")

parser.add_option("--port", '-p', dest="port",
  default="22",
  help="SSH port")

options, args = parser.parse_args()

try:
  hosts = options.hosts.split(',')
except:
  print 'must specify hosts'
  sys.exit(1)

try:
  path = args.pop()
except:
  print 'must specify path'
  sys.exit(1)

boxes = tailer()

for host in hosts:
  boxes.addBox(host, options.login, options.port, path)

boxes.tail()