updates per neale's suggestions (more to come)

pollster/pollster.py

@@ -7,30 +7,71 @@ import time
 import socket
 import urllib.request
 
+# TODO:
+# special stops for http and tftp?
+# get start time and end time of poll, sleep(60-exectime)
+# what to do about exceptions
+# no nested dicts
+# scoring interface
+# config interface
+# html that uses the proper css
+
 DEBUG         = True
 POLL_INTERVAL = 2
 IP_DIR        = 'iptest/'
 REPORT_PATH   = 'iptest/pollster.html'
+SOCK_TIMEOUT  = 0.5
 
-def socket_poll(ip, port, msg):
-	''' Connect via socket to the specified (ip, port), send
-	the specified msg and return the response or None if something
-	went wrong. '''
-	sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+def socket_poll(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. '''
+
+	# create a socket
+	try:
+		sock = socket.socket(socket.AF_INET, prot)
+	except Exception as e:
+		print('pollster: create socket failed')
+		return None
 	
+	sock.settimeout(SOCK_TIMEOUT)
+
+	# connect
 	try:
 		sock.connect((ip, port))
-	except Exception as e:
+	except socket.timeout as e:
+		print('pollster: attempt to connect to %s:%d timed out' % (ip, port))
 		return None
-	
+	except Exception as e:
+		print('pollster: attempt to connect to %s:%d failed' % (ip, port))
+		return None
+
+	# send something
 	sock.send(msg)
-	resp = sock.recv(1024)
+
+	# get a response
+	resp = ''
+	try:
+		# first read
+		data = sock.recv(1024)
+		resp += data.decode('utf-8')
+		max_recv -= 1
+
+		# remaining reads as necessary until timeout or socket closes
+		while(len(data) > 0 and max_recv > 0):
+			data = sock.recv(1024)
+			resp += data.decode('utf-8')
+			max_recv -= 1
+		sock.close()
+	except socket.timeout as e:
+		print('pollster: timed out waiting for a response from %s:%d' % (ip, port))
+	except Exception as e:
+		print('pollster: receive from %s:%d failed' % (ip, port))
+	
 	if len(resp) == 0:
 		return None
-	
-	sock.close()
 
-	resp = resp.decode('utf-8')
 	return resp
 
 # PUT POLLS FUNCTIONS HERE
@@ -39,39 +80,51 @@ def socket_poll(ip, port, msg):
 
 def poll_fingerd(ip):
 	''' Poll the fingerd service. '''
-	resp = socket_poll(ip, 79, b'flag\n')
+	resp = socket_poll(ip, 79, b'flag\n', socket.SOCK_STREAM)
 	if resp is None:
 		return None
 	return resp.strip('\r\n')
 
 def poll_noted(ip):
 	''' Poll the noted service. '''
-	resp = socket_poll(ip, 4000, b'rflag\n')
+	resp = socket_poll(ip, 4000, b'rflag\n', socket.SOCK_STREAM)
 	if resp is None:
 		return None
 	return resp.strip('\r\n')
 
 def poll_catcgi(ip):
 	''' Poll the cat.cgi web service. '''
-	url = urllib.request.urlopen('http://%s/cat.cgi/flag' % ip)
-	data = url.read()
-	if len(data) == 0:
+	request = bytes('GET /cat.cgi/flag HTTP/1.1\r\nHost: %s\r\n\r\n' % ip, 'ascii')
+	resp = socket_poll(ip, 80, request, socket.SOCK_STREAM, 3)
+	if resp is None:
 		return None
-	return data.strip('\r\n')
+
+	content = resp.split('\r\n\r\n')
+	if len(content) < 3:
+		return None
+
+	content = content[1].split('\r\n')
+
+	try:
+		content_len = int(content[0])
+	except Exception as e:
+		return None
+	
+	if content_len <= 0:
+		return None
+	return content[1].strip('\r\n')
 
 def poll_tftpd(ip):
-	''' Poll the ftp service. '''
-	sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-	sock.connect((ip, 69))
-
-	sock.send(b'\x00\x01' + b'flag' + b'\x00' + b'octet' + b'\x00')
-	resp = sock.recv(1024)
+	''' Poll the tftp service. '''
+	resp = socket_poll(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
 	
-	sock.close()
-
-	return resp[4:].decode('utf-8').strip('\r\n')
+	resp = resp.split('\n')[0]
+	return resp[4:].strip('\r\n')
 	
 # PUT POLL FUNCTIONS IN HERE OR THEY WONT BE POLLED
 POLLS = {
@@ -86,13 +139,12 @@ ip_re = re.compile('(\d{1,3}\.){3}\d{1,3}')
 # loop forever
 while(True):
 
-	# check that IP_DIR is there, exit if it isn't
-	if not os.path.isdir(IP_DIR):
-		sys.stderr.write('directory %s does not exist or is not readable\n' % IP_DIR)
-		sys.exit(1)
-
 	# gather the list of IPs to poll
-	ips = os.listdir(IP_DIR)
+	try:
+		ips = os.listdir(IP_DIR)
+	except Exception as e:
+		print('pollster: could not list dir %s' % IP_DIR)
+
 	results = {}
 	for ip in ips:
 
@@ -100,7 +152,11 @@ while(True):
 		if ip_re.match(ip) is None:
 			continue
 
-		#os.remove(os.path.join(IP_DIR, ip))
+		# 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))
 
 		results[ip] = {}