moth/pollster/pollster.py

262 lines
6.8 KiB
Python
Raw Normal View History

2009-10-07 16:54:47 -06:00
#!/usr/bin/env python3
import os
import re
import io
2009-10-07 16:54:47 -06:00
import sys
import time
import socket
import traceback
import subprocess
import random
import http.client
2009-10-07 16:54:47 -06:00
from ctf import config
from ctf import pointscli
DEBUG = False
POLL_INTERVAL = config.get('pollster', 'poll_interval')
IP_DIR = config.get('pollster', 'heartbeat_dir')
REPORT_PATH = config.get('pollster', 'results')
SOCK_TIMEOUT = config.get('pollster', 'poll_timeout')
POLL_IFACE = config.get('pollster', 'poll_iface')
POLL_MAC_VENDOR = config.get('pollster', 'poll_mac_vendor')
class BoundHTTPConnection(http.client.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=socket._GLOBAL_DEFAULT_TIMEOUT):
http.client.HTTPConnection.__init__(self, host, port, strict, timeout)
self.bindip = bindip
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))
if self._tunnel_host:
self._tunnel()
def random_mac():
''' Set a random mac on the poll interface. '''
mac = ':'.join([POLL_MAC_VENDOR] + ['%02x' % random.randint(0,255) for i in range(3)])
retcode = subprocess.call(('ifconfig', POLL_IFACE, 'hw', 'ether', mac))
def dhcp_request():
''' Request a new IP on the poll interface. '''
retcode = subprocess.call(('dhclient', POLL_IFACE))
def get_ip():
''' Return the IP of the poll interface. '''
path = os.path.join(os.getcwd(), 'get_ip.sh')
p = subprocess.Popen((path, POLL_IFACE), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
(out, err) = p.communicate()
return out.strip(b'\r\n').decode('utf-8')
def socket_poll(srcip, ip, port, msg, prot, max_recv=1):
''' Connect via socket to the specified <ip>:<port> using the
2009-10-08 16:42:26 -06:00
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). '''
2009-10-07 16:54:47 -06:00
# create a socket
try:
sock = socket.socket(socket.AF_INET, prot)
except Exception as e:
print('pollster: create socket failed (%s)' % e)
traceback.print_exc()
return None
2009-10-08 16:42:26 -06:00
sock.bind((srcip, 0))
sock.settimeout(SOCK_TIMEOUT)
# connect
2009-10-07 16:54:47 -06:00
try:
sock.connect((ip, port))
except socket.timeout as e:
print('pollster: attempt to connect to %s:%d timed out (%s)' % (ip, port, e))
traceback.print_exc()
return None
2009-10-07 16:54:47 -06:00
except Exception as e:
print('pollster: attempt to connect to %s:%d failed (%s)' % (ip, port, e))
traceback.print_exc()
2009-10-07 16:54:47 -06:00
return None
# send something
2009-10-07 16:54:47 -06:00
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 as e:
print('pollster: timed out waiting for a response from %s:%d (%s)' % (ip, port, e))
traceback.print_exc()
except Exception as e:
print('pollster: receive from %s:%d failed (%s)' % (ip, port, e))
traceback.print_exc()
2009-10-08 16:42:26 -06:00
sock.close()
2009-10-07 16:54:47 -06:00
if len(resp) == 0:
return None
return b''.join(resp)
2009-10-07 16:54:47 -06:00
# 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, b'flag\n', socket.SOCK_STREAM)
2009-10-07 16:54:47 -06:00
if resp is None:
return None
return resp.strip(b'\r\n')
2009-10-07 16:54:47 -06:00
def poll_noted(srcip, ip):
''' Poll the noted service. Returns None or a team name. '''
resp = socket_poll(srcip, ip, 4000, b'rflag\n', socket.SOCK_STREAM)
2009-10-07 16:54:47 -06:00
if resp is None:
return None
return resp.strip(b'\r\n')
2009-10-07 16:54:47 -06:00
def poll_catcgi(srcip, ip):
''' Poll the cat.cgi web service. Returns None or a team name. '''
try:
conn = BoundHTTPConnection(srcip, ip, timeout=SOCK_TIMEOUT)
conn.request('GET', '/cat.cgi/flag')
except:
2009-10-07 16:54:47 -06:00
return None
resp = conn.getresponse()
if resp.status != 200:
conn.close()
return None
data = resp.read()
conn.close()
return data.strip(b'\r\n')
2009-10-07 16:54:47 -06:00
def poll_tftpd(srcip, ip):
''' Poll the tftp service. Returns None or a team name. '''
resp = socket_poll(srcip, ip, 69, b'\x00\x01' + b'flag' + b'\x00' + b'octet' + b'\x00', socket.SOCK_DGRAM)
if resp is None:
return None
2009-10-08 16:42:26 -06:00
if len(resp) <= 5:
return None
2009-10-08 16:42:26 -06:00
resp = resp.split(b'\n')[0]
# ack
2010-01-26 13:30:23 -07:00
_ = socket_poll(srcip, ip, 69, b'\x00\x04\x00\x01' + '\x00' * 14, socket.SOCK_DGRAM, 0)
return resp[4:].strip(b'\r\n')
2009-10-08 16:42:26 -06:00
2009-10-07 16:54:47 -06:00
# PUT POLL FUNCTIONS IN HERE OR THEY WONT BE POLLED
POLLS = {
'fingerd' : poll_fingerd,
'noted' : poll_noted,
2009-10-07 16:54:47 -06:00
'catcgi' : poll_catcgi,
'tftpd' : poll_tftpd,
2009-10-07 16:54:47 -06:00
}
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()
2009-10-07 16:54:47 -06:00
# gather the list of IPs to poll
2009-10-09 10:15:24 -06:00
ips = os.listdir(IP_DIR)
2009-10-09 10:15:24 -06:00
out = io.StringIO()
out.write(config.start_html('Team Service Availability'))
2009-10-07 16:54:47 -06:00
for ip in ips:
# check file name format is ip
if ip_re.match(ip) is None:
continue
# remove the file
2010-01-19 15:11:59 -07:00
try:
os.remove(os.path.join(IP_DIR, ip))
except Exception as e:
print('pollster: could not remove %s' % os.path.join(IP_DIR, ip))
traceback.print_exc()
2009-10-07 16:54:47 -06:00
results = {}
2009-10-07 16:54:47 -06:00
if DEBUG is True:
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')
2009-10-07 16:54:47 -06:00
# perform polls
for service,func in POLLS.items():
2010-01-19 15:11:59 -07:00
try:
team = func(srcip, ip).decode('utf-8')
if len(team) == 0:
team = 'dirtbags'
2010-01-19 15:11:59 -07:00
except:
2009-10-07 16:54:47 -06:00
team = 'dirtbags'
if DEBUG is True:
print('\t%s - %s' % (service, team))
2009-10-07 16:54:47 -06:00
if out is not None:
out.write('<tr><td>%s</td><td>%s</td>\n' % (service, team))
2009-10-08 17:24:05 -06:00
pointscli.submit('svc.' + service, team, 1)
if out is not None:
out.write('</table>\n')
2009-10-08 16:42:26 -06:00
2009-10-07 16:54:47 -06:00
if DEBUG is True:
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))
2010-01-19 15:11:59 -07:00
t_end = time.time()
exec_time = int(t_end - t_start)
sleep_time = POLL_INTERVAL - exec_time
2009-10-07 16:54:47 -06:00
if out is not None:
out.write(config.end_html())
2009-10-09 10:15:24 -06:00
open(REPORT_PATH, 'w').write(out.getvalue())
2009-10-07 16:54:47 -06:00
# sleep until its time to poll again
time.sleep(sleep_time)
2009-10-07 16:54:47 -06:00