#!/usr/bin/python
#
# Copyright (c) 2003 Antti Kuntsi
#
# This software is released under the GNU General Public License (GPL)
#
#

"""TTY Multiplexer
Options:

-h
--help		This help
-t <port>
--ttyPort=<port>	Set the port number for TTY feeds (default 4444)
-c <port>
--clientPort=<port>	Set the port number for clients (default 5555)
-g <message>
--greeting <message>	Set the greeting message diplayed to connecting clients
-n
--no-DNS-lookup		Do not do DNS lookups for TTY feeds (done by default)
-m
--max-clients		Maximum number of feed clients (default 100)
-v
--verbose		Announce connection at master terminal
"""

import os, sys, thread, time, select, mutex, threading, operator
from socket import *
from getopt import getopt

welcome="Welcome to the TTY Multiplexer, pick your TTY from the list below"

DNSlookup = 1
MAXTTYS=20
verbose=0

def announce(text):
	if verbose: print text

class TTYFeed:
	def __init__(self, sock, addr):
		self.remoteAddr, self.remotePort = addr
		if DNSlookup:
			try: self.remoteAddr = gethostbyaddr(self.remoteAddr)[0]
			except: pass
		self.sock=sock
		self.startTime = time.ctime(time.time())
		self.clients = []
		self.clientMutex = mutex.mutex()
		self.running=1
		self.lastBuffer=""
		self.initBuffer=""
		announce("New feed from %s:%d"%(self.remoteAddr, self.remotePort))

	def __del__(self):
		announce("Feed %s:%d exited"%(self.remoteAddr, self.remotePort))

	def addClient(self, client):
		self.clientMutex.lock(self._addClient, client)
	
	def _addClient(self, client):
		self.clients.append(client)
		if self.initBuffer:
			client.feedData(self.initBuffer)
		if self.lastBuffer:
			client.feedData(self.lastBuffer)
		self.clientMutex.unlock()
	
	def delClient(self, client):
		self.clientMutex.lock(self._delClient, client)

	def _delClient(self, client):
		try:
			self.clients.remove(client)
		except: # Nobody cares about errors anyways, right? =)
			pass
		self.clientMutex.unlock()

	def feedClients(self, data):
		for client in self.clients:
			client.feedData(data)
		self.clientMutex.unlock()
		if len(self.initBuffer)<1024:
			self.initBuffer+=data
		else:
			self.lastBuffer+=data
		if len(self.lastBuffer)>1024:
			self.lastBuffer=self.lastBuffer[-1024:]

	def stopClients(self, none):
		for client in self.clients:
			client.close()
		self.clients=[]
		self.clientMutex.unlock()

	def serve(self):
		data=self.sock.recv(10)
		while data:
			self.clientMutex.lock(self.feedClients, data)
			data=self.sock.recv(1000)
		self.clientMutex.lock(self.stopClients, None)
		self.running=0
		thread.exit()

class ClientFeed:
	def __init__(self, sock, addr):
		self.sock=sock
		self.feeds=[]
		self.buffer=[]
		self.running=1
		self.datafeed=threading.Semaphore(0)
		announce("New client from %s:%d"%(addr[0], addr[1]))

	def __del__(self):
		try:
			self.sock.close()
		except: pass
		announce("Client exited")

	def close(self):
		self.sock.close()
		self.running=0
		self.datafeed.release()
	
	def addFeeds(self, feeds):
		self.feeds.extend(feeds)

	def feedData(self, data):
		self.buffer.append(data)
		self.datafeed.release()

	def serve(self):
		feed=self.selectFeed()
		self.sock.setblocking(0)
		if feed == None:
			self.sock.close()
			thread.exit()
		feed.addClient(self)
		while self.running:
			self.datafeed.acquire()
			if self.buffer:
				try:
					self.sock.send(self.buffer.pop(0))
				except:
					feed.delClient(self)
					self.running=0
		thread.exit()

	def selectFeed(self):
		maxSelection=len(self.feeds)
		if maxSelection==0:
			feeds="No feeds available right now, try again later."
			maxSelection=None
		elif maxSelection==1:
			return self.feeds[0]
		else:
			feeds="\n".join(["%d) %s:%d (%s)"%(i, f.remoteAddr, f.remotePort, f.startTime) for i, f in zip(range(MAXTTYS), self.feeds)])
		feedstring="%s\n%s\n"%(welcome, feeds)
		self.sock.send(feedstring)
		if maxSelection==None:
			return None
		selection=-1
		while selection>maxSelection or selection<0:
			data=self.sock.recv(10)
			if not data: # client closed connection
				return None
			try:
				selection=int(data[:2])
			except:
				selection=MAXTTYS
				self.sock.send("Invalid selection, numbers 0 - %d only"%(maxSelection-1))
		return self.feeds[selection]

class TTYServer:
	def __init__(self, ttyPort=4444, clientPort=5555):
		self.ttyPort=ttyPort
		self.clientPort=clientPort
		self.ttys=[]
		self.ttySocket=self._bindSocket("", self.ttyPort)
		self.clientSocket=self._bindSocket("", self.clientPort)

	def serve(self):
		ilist=[self.ttySocket, self.clientSocket]
		iact, iout, iext = [], [], []
		while 1:
			iact, iout, iext = select.select(ilist, [], [])
			self.ttys=filter(lambda tty: tty.running, self.ttys)
			for sock in iact:
				if sock == self.ttySocket:
					if len(self.ttys)>MAXTTYS:
						s,a=sock.accept()
						s.close()
						del s, a
					else:
						s,a=sock.accept()
						tty = TTYFeed(s, a)
						self.ttys.append(tty)
						thread.start_new_thread(tty.serve, ())

				elif sock == self.clientSocket:
					s,a=sock.accept()
					if reduce(operator.add,
						[len(f.clients) for f in self.ttys], 0) >= self.clientLimit:
						s.close()
					else:
						client = ClientFeed(s, a)
						client.addFeeds(self.ttys)
						thread.start_new_thread(client.serve, ())

	def _bindSocket(self, addr, port):
		sock=socket(AF_INET, SOCK_STREAM)
		sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
		sock.bind(("", port))
		sock.listen(5)
		return sock


if __name__ == "__main__":

	ttyPort=4444
	clientPort=5555
	maxFeeds=100

	opts, args = getopt(sys.argv[1:], "hvt:c:g:m:", ["help", "tty-port=", "client-port=", "greeting=", "max-feed=", "verbose"])

	for opt, value in opts:
		if opt in ("-h", "--help"):
			print __doc__
			sys.exit()
		elif opt in ("-t", "--tty-port"):
			ttyPort=int(value)
		elif opt in ("-c", "--client-port"):
			clientPort=int(value)
		elif opt in ("-g", "--greeting"):
			welcome=value
		elif opt in ("-m", "--max-clients"):
			maxFeeds=int(value)
		elif opt in ("-v", "--verbose"):
			verbose = 1
		elif opt in ("-n", "--no-DNS-lookup"):
			DNSlookup = 0

	server=TTYServer(ttyPort, clientPort)
	TTYServer.clientLimit=maxFeeds
	server.serve()