moth/pollster/pollster.py

255 lines
6.5 KiB
Python
Executable File

#!/usr/bin/env python3
import os
import re
import io
import sys
import time
import socket
import traceback
import subprocess
import random
import http.client
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
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 as e:
print('pollster: create socket failed (%s)' % e)
traceback.print_exc()
return None
sock.bind((srcip, 0))
sock.settimeout(SOCK_TIMEOUT)
# connect
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
except Exception as 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 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()
sock.close()
if len(resp) == 0:
return None
return b''.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, b'flag\n', socket.SOCK_STREAM)
if resp is None:
return None
return resp.strip(b'\r\n')
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)
if resp is None:
return None
return resp.strip(b'\r\n')
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:
return None
resp = conn.getresponse()
if resp.status != 200:
conn.close()
return None
data = resp.read()
conn.close()
return data.strip(b'\r\n')
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
if len(resp) <= 5:
return None
resp = resp.split(b'\n')[0]
# ack
_ = socket_poll(srcip, ip, 69, b'\x00\x04' + resp[2:4], socket.SOCK_DGRAM, 0)
return resp[4:].strip(b'\r\n')
# PUT POLL FUNCTIONS IN HERE OR THEY WONT BE POLLED
POLLS = {
'fingerd' : poll_fingerd,
'noted' : poll_noted,
'catcgi' : poll_catcgi,
'tftpd' : poll_tftpd,
}
ip_re = re.compile('(\d{1,3}\.){3}\d{1,3}')
poll_no = 0
# 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(IP_DIR)
out = io.StringIO()
out.write(config.start_html('Team Service Availability'))
for ip in ips:
# check file name format is ip
if ip_re.match(ip) is None:
continue
# remove the file
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()
results = {}
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')
# 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 DEBUG is True:
print('\t%s - %s' % (service, team))
if out is not None:
out.write('<tr><td>%s</td><td>%s</td>\n' % (service, team))
pointscli.submit('svc.' + service, team, 1)
if out is not None:
out.write('</table>\n')
if DEBUG is True:
print('+-----------------------------------------+')
out.write('<p>Poll number: %d</p>' % poll_no)
poll_no += 1
t_end = time.time()
exec_time = int(t_end - t_start)
sleep_time = POLL_INTERVAL - exec_time
if out is not None:
out.write(config.end_html())
open(REPORT_PATH, 'w').write(out.getvalue())
# sleep until its time to poll again
time.sleep(sleep_time)