moth/bin/pollster

286 lines
7.7 KiB
Plaintext
Raw Normal View History

#!/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)