#!/usr/bin/python

import os
import re
import sys
import time
import socket
import traceback
import subprocess
import random
import httplib
import optparse
import cStringIO as io

from ctf import pointscli, html

ifconfig = '/sbin/ifconfig'
udhcpc = '/sbin/udhcpc'

class BoundHTTPConnection(httplib.HTTPConnection):
	''' http.client.HTTPConnection doesn't support binding to a particular
	address, which is something we need. '''

	def __init__(self, bindip, host, port=None, strict=None, timeout=None):
		httplib.HTTPConnection.__init__(self, host, port, strict)
		self.bindip = bindip
		self.timeout = timeout

	def connect(self):
		''' Connect to the host and port specified in __init__, but
		also bind first. '''
		self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		self.sock.bind((self.bindip, 0))
		self.sock.settimeout(self.timeout)
		self.sock.connect((self.host, self.port))

def random_mac():
	''' Set a random mac on the poll interface. '''
	retcode = subprocess.call((ifconfig, opts.iface, 'down'))
	mac = ':'.join([opts.mac_vendor] + ['%02x' % random.randint(0,255) for i in range(3)])
	retcode = subprocess.call((ifconfig, opts.iface, 'hw', 'ether', mac, 'up'))

def dhcp_request():
	''' Request a new IP on the poll interface. '''
	retcode = subprocess.call((udhcpc, '-i', opts.iface, '-q'))

def get_ip():
	''' Return the IP of the poll interface. '''
	ip_match = re.compile(r'inet addr:([0-9.]+)')
	p = subprocess.Popen((ifconfig, opts.iface), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
	(out, err) = p.communicate()
	for line in out.splitlines():
		m = ip_match.search(line)
		if m is not None:
			return m.group(1).decode('utf-8')
	return '10.1.1.1'

def socket_poll(srcip, ip, port, msg, prot, max_recv=1):
	''' Connect via socket to the specified <ip>:<port> using the
	specified <prot>, send the specified <msg> and return the
	response or None if something went wrong. <max_recvs> specifies
	how many times to read from the socket (defaults to once). '''

	# create a socket
	try:
		sock = socket.socket(socket.AF_INET, prot)
	except Exception, e:
		print('pollster: create socket failed (%s)' % e)
		traceback.print_exc()
		return None

	sock.bind((srcip, 0))
	sock.settimeout(opts.timeout)

	# connect
	try:
		sock.connect((ip, port))
	except socket.timeout, e:
		print('pollster: attempt to connect to %s:%d timed out (%s)' % (ip, port, e))
		traceback.print_exc()
		return None
	except Exception, e:
		print('pollster: attempt to connect to %s:%d failed (%s)' % (ip, port, e))
		traceback.print_exc()
		return None

	# send something
	sock.send(msg)

	# get a response
	resp = []
	try:
		# read from the socket until <max_recv> responses or read,
		# a timeout occurs, the socket closes, or some other exception
		# is raised
		for i in range(max_recv):
			data = sock.recv(1024)
			if len(data) == 0:
				break
			resp.append(data)

	except socket.timeout, e:
		print('pollster: timed out waiting for a response from %s:%d (%s)' % (ip, port, e))
		traceback.print_exc()
	except Exception, e:
		print('pollster: receive from %s:%d failed (%s)' % (ip, port, e))
		traceback.print_exc()

	sock.close()

	if len(resp) == 0:
		return None

	return ''.join(resp)

# PUT POLLS FUNCTIONS HERE
#  Each function should take an IP address and return a team name or None
#  if (a) the service is not up, (b) it doesn't return a valid team name.

def poll_fingerd(srcip, ip):
	''' Poll the fingerd service. Returns None or a team name. '''
	resp = socket_poll(srcip, ip, 79, 'flag\n', socket.SOCK_STREAM)
	if resp is None:
		return None
	return resp.split('\n')[0]

def poll_noted(srcip, ip):
	''' Poll the noted service. Returns None or a team name. '''
	resp = socket_poll(srcip, ip, 4000, 'rflag\n', socket.SOCK_STREAM)
	if resp is None:
		return None
	return resp.split('\n')[0]

def poll_catcgi(srcip, ip):
	''' Poll the cat.cgi web service. Returns None or a team name. '''

	try:
		conn = BoundHTTPConnection(srcip, ip, timeout=opts.timeout)
		conn.request('GET', '/var/www/flag')
	except Exception, e:
		traceback.print_exc()
		return None

	resp = conn.getresponse()
	if resp.status != 200:
		conn.close()
		return None

	data = resp.read()
	conn.close()
	return data.split('\n')[0]

def poll_tftpd(srcip, ip):
	''' Poll the tftp service. Returns None or a team name. '''
	resp = socket_poll(srcip, ip, 69, '\x00\x01' + 'flag' + '\x00' + 'octet' + '\x00', socket.SOCK_DGRAM)
	if resp is None:
		return None

	if len(resp) <= 5:
		return None

	resp = resp.split('\n')[0]

	# ack
	_ = socket_poll(srcip, ip, 69, '\x00\x04\x00\x01' + '\x00' * 14, socket.SOCK_DGRAM, 0)

	return resp[4:].split('\n')[0]

# PUT POLL FUNCTIONS IN HERE OR THEY WONT BE POLLED
POLLS = {
	'fingerd'  : poll_fingerd,
	'noted'	   : poll_noted,
	'catcgi'   : poll_catcgi,
	'tftpd'	   : poll_tftpd,
}



p = optparse.OptionParser()
p.add_option('-d', '--debug', action='store_true', dest='debug',
             default=False,
             help='Turn on debugging output')
p.add_option('-s', '--interval', type='float', dest='interval',
             default=60,
             help='Time between polls, in seconds (default: %default)')
p.add_option('-t', '--timeout', type='float', dest='timeout',
             default=0.5,
             help='Poll timeout, in seconds (default: %default)')
p.add_option('-b', '--heartbeat-dir', dest='heartbeat_dir',
             default='/var/lib/ctf/heartbeat',
             help='Where in.heartbeatd writes its files')
p.add_option('-r', '--results-page', dest='results',
             default='/tmp/services.html',
             help='Where to write results file')
p.add_option('-i', '--interface', dest='iface',
             default='eth1',
             help='Interface to bind to')
p.add_option('-m', '--mac-vendor', dest='mac_vendor',
             default='00:01:c0',
             help='MAC vendor to look for (default: %default)')
opts, args = p.parse_args()

socket.setdefaulttimeout(opts.timeout)

ip_re = re.compile('(\d{1,3}\.){3}\d{1,3}')
# loop forever
while True:

	random_mac()
	dhcp_request()
	
	srcip = get_ip()

	t_start = time.time()

	# gather the list of IPs to poll
	ips = os.listdir(opts.heartbeat_dir)

	out = io.StringIO()
	for ip in ips:
		# check file name format is ip
		if ip_re.match(ip) is None:
			continue

		# remove the file
                fn = os.path.join(opts.heartbeat_dir, ip)
		try:
			os.remove(fn)
		except Exception, e:
			print('pollster: could not remove %s' % fn)
			traceback.print_exc()

		results = {}

		if opts.debug:
			print('ip: %s' % ip)

		if out is not None:
			out.write('<h2>%s</h2>\n' % ip)
			out.write('<table class="pollster">\n<thead><tr><td>Service Name</td></td>')
			out.write('<td>Flag Holder</td></tr></thead>\n')

		# perform polls
		for service,func in POLLS.items():
			try:
				team = func(srcip, ip).decode('utf-8')
				if len(team) == 0:
					team = 'dirtbags'
			except:
				team = 'dirtbags'

			if opts.debug:
				print('\t%s - %s' % (service, team))

			if out is not None:
				out.write('<tr><td>%s</td><td>%s</td>\n' % (service, team))

			pointscli.award('svc.' + service, team, 1)

		if out is not None:
			out.write('</table>\n')

	if opts.debug:
		print('+-----------------------------------------+')

	time_str = time.strftime('%a, %d %b %Y %H:%M:%S %Z')
	out.write('''
		<p>This page was generated on %s. That was <span id="diff">?</span> seconds ago.</p>
		<script type="text/javascript">
			var gen_time = new Date(%f);
			var cur_time = new Date();
			var diff = (cur_time.getTime() - gen_time.getTime())/1000;
			document.getElementById("diff").innerHTML = diff;
		</script>
		''' % (time_str, time.time()*1000))

	t_end = time.time()
	exec_time = int(t_end - t_start)
	sleep_time = opts.interval - exec_time

        html.write(opts.results,'Team Service Availability', out.getvalue())

	# sleep until its time to poll again
	time.sleep(sleep_time)