sender

#!/usr/bin/env python

import sys
import time
import os.path
import os
import glob
import socket
import optparse
import md5
import zipfile
import gzip
import tempfile
import shutil
import subprocess
import threading
import Queue
import signal
import traceback

DEFAULT_SIGNAGURE_LENGTH=256

filter_stopped = Queue.Queue()
stop_signal = Queue.Queue()
running_lock = threading.Lock()

def should_stop():
	return not stop_signal.empty()

def main():
	parser = optparse.OptionParser(usage="Usage: %prog [options] <pattern0> ... <patternN>", description=
		"Outputs the content of files resolved by the patterns passed as "
		"parameters and keep monitoring them for new content. "
		"Optionally the offsets can be writen to the <offsets> file. Which "
		"will be read when starting. "
		"IMPORTANT: only files with size >= signature-length (default %d) bytes will be processed. "
		"Zip files will be open recursively, and only once." % DEFAULT_SIGNAGURE_LENGTH)

	parser.add_option('-f', '--follow', action="store_true", dest="follow", help="Pools the file very second for changes in an infinite loop.", default=False)
	parser.add_option('-p', '--offsets', action="store", dest="offsets", help="File to persist the last offsets read for each file. If it doesn't exist, the files are read from beginning.", default='/dev/null')
	parser.add_option('-t', '--tcp', action="store", dest="host", help="Sends the output to the host:port via TCP.", metavar="HOST:PORT")
	parser.add_option('-s', '--signature-length', action="store", dest="signatureLength", help="First BYTES of the each file will be used as it's ID in offsets file, so the files can be renamed not losing the offset. Default is %d." % DEFAULT_SIGNAGURE_LENGTH, metavar="BYTES", default=DEFAULT_SIGNAGURE_LENGTH)
	parser.add_option('-g', '--greater-or-equal-than', action="store", dest="starting", help="Only copy lines which are >= than STARTING in lexicographical order. In lines starting with datetime, this can be used to ignore older events.")
	parser.add_option('-e', '--start-from-the-tail', action="store_true", dest="start_from_tail", help="Generate offsets file at the tail of every file. Useful if you want to start following the end of all files returned in pattern. .. Well, ask Fabio.")
	parser.add_option('-l', '--filter', action="store", dest="filter", help="Command or shell script that will filter the events using stdin and stdout.")
	parser.add_option('-r', '--retry-on-network-error', action="store_true", dest="retry_on_network_error", help="Keep trying to reconnect instead of exiting with an error.")
	parser.add_option('-i', '--dump-pid', action="store", dest="dump_pid", help="Write process id to this file. Disabled by default.", metavar="FILE")
	parser.add_option('-k', '--kill', action="store", dest="kill", help="Send SIGTERM to process id saved in FILE and exit", metavar="FILE")

	options, args = parser.parse_args()

	if options.dump_pid:
		f = open(options.dump_pid,'w')
		f.write("%d" % os.getpid())
		f.flush()
		f.close()

	if options.kill:
		pid = int(open(options.kill).readline())
		os.kill(pid, signal.SIGTERM)
		sys.exit(0)

	if len(args) == 0:
		parser.print_help()
	else:
		output = Stdout()
		if options.host:
			a = options.host.split(":")
			try:
				output = Netcat(a[0], int(a[1]), options.retry_on_network_error)
			except socket.error, e:
				log("Socket error: %s" % str(e))
				sys.exit(1)


		if options.filter:
			filter = subprocess.Popen(["%s/filter_wrapper" % os.path.dirname(sys.argv[0]), options.filter], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
			last_output = output
			output = filter.stdin
			input = filter.stdout
			def filter_pipe():
				while filter.poll() is None:
					line = input.readline()
					if len(line) == 0: # in some platforms readline returns zero-length strings if process has terminated and python still doesnt know it.. weird right?
						continue
					try:
						last_output.write(line)
						last_output.flush()
					except socket.error, e:
						log("Socket error in filer: %s" % e)
						break
					except:
						log("Some exception happened. Exiting.")
						break
				filter_stopped.put(True)

			t = threading.Thread(target=filter_pipe)
			t.daemon = True
			t.start()

			signal_names = dict((k, v) for v, k in reversed(sorted(signal.__dict__.items()))
 				if v.startswith('SIG') and not v.startswith('SIG_'))

			def close_filter(sig=signal.SIGINT):
				if not filter.stdin.closed:
					filter.stdin.close()
				try:
					count = 0
					while filter.poll() is None:
						time.sleep(0.05)
						count += 1
						if count == 8:
							log("Filter is taking some time to exit. Waiting a bit.")
						if count > 40:
							log("That does it. Sending SIGINT.")
							os.kill(filter.pid, sig)
							break
				except OSError, e:
					log("An exception happened while waiting for process to finish: %s" % e)


			def signal_handler(signal, frame):
				log("%s caught" % signal_names[signal])
				stop_signal.put(True)
				# if not filter.stdin.closed:
				# 	running_lock.acquire(True)
				time.sleep(0.1)
				close_filter(signal)

			signal.signal(signal.SIGINT, signal_handler)
			signal.signal(signal.SIGTERM, signal_handler)

		try:
			Tail(options, args, output).run()

			if options.filter:
				close_filter()


		except KeyboardInterrupt:
			pass

def dump_all_threads_stack_trace():
	for th in threading.enumerate():
		print(th)
		traceback.print_stack(sys._current_frames()[th.ident])
		print()


def log(msg):
	sys.stderr.write(str(msg))
	sys.stderr.write("\n")
	sys.stderr.flush()

def debug(msg):
	# sys.stderr.write(str(msg))
	# sys.stderr.write("\n")
	# sys.stderr.flush()
	pass


class Stdout(object):
	def write(self, msg):
		sys.stdout.write(msg)

	def flush(self):
		sys.stdout.flush()

	def fileno(self):
		return sys.stdout.fileno()


class Netcat(object):
	def __init__(self, host, port, retry_on_network_error):
		self.retry_on_network_error = retry_on_network_error
		self.addr = (host, port)
		self.show_error_message = True
		self.send_last = False
		self.connect()

	def connect(self):
		while True:
			try:
				self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
				self.s.connect(self.addr)
				self.show_error_message = True
				break;
			except socket.error, e:
				self.error("Socket error: %s" % e)
			except IOError, e:
				self.error("IOError: %s" % e)
			self.exit_or_retry()

	def error(self, msg):
		if self.show_error_message:
			log(msg)

	def exit_or_retry(self):
		if not self.retry_on_network_error:
			sys.exit(1)

		self.error("Retrying every 3 seconds.")
		self.show_error_message = False
		time.sleep(3)


	def write(self, msg):
		while True:
			try:
				if self.send_last:
					self.s.sendall(self.last_msg)
				self.s.sendall(msg)
				self.last_msg = msg
				self.send_last = False
				break
			except socket.error, e:
				self.error("Socket error: %s" % e)
			except IOError, e:
				self.error("IOError: %s" % e)
			self.send_last = True
			self.exit_or_retry()
			self.reconnect()


	def reconnect(self):
		if self.s:
			try:
				self.s.close()
			except:
				pass
		self.connect()

	def flush(self):
		pass

	def fileno(self):
		return self.s.fileno()


class Info(object):
	def __init__(self, line=None, sig=None, name='-', offset=0):
		if line:
			a = line.split(' ')
			self.sig = a[0]
			self.offset = int(a[1])
			self.name = ' '.join(a[2:])
		else:
			self.sig = sig
			self.name = name
			self.offset = offset


	def dump(self):
		return '%s %d %s' % (self.sig, self.offset, self.name)

class Tail(object):

	def __init__(self, options, fnpatterns, output):
		self.fnpatterns = fnpatterns
		self.offsetsfn = options.offsets
		self.signatureLength = int(options.signatureLength)
		self.follow = options.follow
		self.output = output
		self.starting = options.starting
		self.start_from_tail = options.start_from_tail
		self.readOffsets()

	def persistOffsets(self):
		if self.offsetsfn == '/dev/null':
			return

		temp = '%s-temp' % self.offsetsfn
		f = open(temp,'w')
		try:
			for sig, info in self.offsets.iteritems():
				f.write(info.dump())
				f.write('\n')
			f.flush()
		finally:
			f.close()

		shutil.move(temp, self.offsetsfn)

	def readOffsets(self):
		if not os.path.isfile(self.offsetsfn):
			self.offsets = {}
			return

		f = open(self.offsetsfn,'r')
		try:
			self.offsets = {}
			for line in f:
				line = line.strip()
				if len(line) == 0:
					continue
				info = Info(line.strip())
				self.offsets[info.sig] = info
		finally:
			f.close()

	def generateSignature(self, f):
		offset = f.tell();
		header = f.read(self.signatureLength)
		f.seek(offset)
		if len(header) == 0:
			return None
		else:
			return md5.new(header).hexdigest()

	def purgeOffsetsNotIn(self, existing):
		newOffsets = {}
		for sig in existing:
			newOffsets[sig] = self.offsets[sig]
		if len(self.offsets) != len(newOffsets):
			self.offsets = newOffsets
			self.persistOffsets()

	def copy(self, f):
		if should_stop():
			return
		for line in f:
			if not self.starting or line >= self.starting:
				self.output.write(line)
			if should_stop():
				break;
		self.output.flush()

	def isCompressed(self, fn):
		return fn.endswith('.zip') or fn.endswith('.gz')


	def visitzip(self, zip, visitor=None):
		zf = zipfile.ZipFile(zip)
		try:
			if visitor:
				for name in zf.namelist():
					visitor(name, zf)
			return zf.namelist()
		finally:
			zf.close()


	def unzip(self, to, f, fn):
		if sys.version_info[0] >= 2 and sys.version_info[1] >= 6:
			return self.visitzip(f, lambda n,z: z.extract(n, to))

		os.system('unzip -qq %s -d %s' % (fn, to))
		return self.visitzip(f)


	def processCompressedFile(self, f, fn):
		if fn.endswith('.gz'):
			self.processGzipFile(fn)
		else:
			self.processZipFile(f, fn)

	def processGzipFile(self, fn):
		debug("gz: %s" % fn)
		f = gzip.open(fn)
		try:
			self.processFile(f, '/var/tmp/fake.log', {})
		finally:
			f.close()

	def processZipFile(self, f, fn):
		path = tempfile.mkdtemp()
		for name in self.unzip(path, f, fn):
			self.processFileByName('%s/%s' % (path, name), {})
		shutil.rmtree(path)

	def processFileByName(self, fn, existing):
		debug("processFileByName: %s" % fn)
		f = open(fn, 'rb')
		try:
			self.processFile(f, fn, existing)
		finally:
			f.close()

	def processFile(self, f, fn, existing):
		debug("processFile: %s" % fn)
		sig = self.generateSignature(f)
		if not sig:
			return

		if sig in existing and os.path.getsize(fn) != os.path.getsize(existing[sig]):
			log("WARN Files '%s' and '%s' have same signature and different sizes" % (fn, existing[sig]))

		info = self.offsets.get(sig, Info(sig=sig, name=fn))

		lastOffset = info.offset

		info.name = fn

		if self.isCompressed(fn):
			debug("compressed: %s" % fn)
			if info.offset == 0:
				if not self.start_from_tail:
					self.processCompressedFile(f, fn)
				info.offset = -1
		else:
			if self.start_from_tail:
				info.offsets = os.path.getsize(fn)
			if os.path.exists(fn) and os.path.getsize(fn) < info.offset:
				log("WARN file %s was truncated" % fn)
				info.offset = os.path.getsize(fn)
			else:
				f.seek(info.offset)
				self.copy(f)
				info.offset = f.tell()

		existing[sig] = fn

		if lastOffset != info.offset:
			self.offsets[sig] = info
			self.persistOffsets()


	def run(self):
		running_lock.acquire(True)
		try:
			while not should_stop():
				if not filter_stopped.empty():
					sys.exit(1)
				existing = {}
				for fnpattern in self.fnpatterns:
					for fn in glob.glob(fnpattern):
						try:
							if not os.path.isfile(fn):
								log("File no longer exists: %s" % fn)
								continue
							if os.path.getsize(fn) < self.signatureLength and not self.isCompressed(fn):
								continue
							self.processFileByName(fn, existing)
						except Exception, e:
							log("Exception: %s" % e)
							exc_type, exc_value, exc_traceback = sys.exc_info()
							traceback.print_tb(exc_traceback)


				if len(existing) != len(self.offsets):
					self.purgeOffsetsNotIn(existing)
				if not self.follow:
					break
				time.sleep(1)
		finally:
			running_lock.release()



if __name__ == '__main__':
	main()

#import pdb ; pdb.set_trace()