#!/usr/bin/python
##########################################################################
# socketstatefollower-w32.py - Watch for tcp/udp socket changes
##########################################################################
# Copyright 2016 Todd Shadburn <tshadburn@opencomputingsolutions.com>
#
# Licensed under the GNU GPL version 2
##########################################################################


import ctypes
import socket
import struct
import datetime

TCP_STATE_CLOSED = 1
TCP_STATE_LISTEN = 2
TCP_STATE_SYN_SENT = 3
TCP_STATE_SYN_RCVD = 4
TCP_STATE_ESTAB = 5
TCP_STATE_FIN_WAIT1 = 6
TCP_STATE_FIN_WAIT2 = 7
TCP_STATE_CLOSE_WAIT = 8
TCP_STATE_CLOSING = 9
TCP_STATE_LAST_ACK = 10
TCP_STATE_TIME_WAIT = 11
TCP_STATE_DELETE_TCB = 12

class SocketConnection:
	def __init__(self, lip, lport, rip, rport, proto, state):
		self.lip = lip
		self.lport = lport
		self.rip = rip
		self.rport = rport
		self.proto = proto
		self.state = state
		self.last_seen = datetime.datetime.now()
		return

	def seen(self, time=datetime.datetime.now()):
		self.last_seen = time
		return


SOCK_CHANGE_TYPE_NEW_ESTAB = 1
SOCK_CHANGE_TYPE_NEW_LISTEN = 2
SOCK_CHANGE_TYPE_CLOSED = 3

class SocketStateFollower:
	def __init__(self, *args, **kwargs):
		self.socks_index = {}	# key=lip:lport:rip:rport:proto
					# UDP will have 0s for rip, rport & state
		self.last_changes = []  # list of 'changes' between prior and
					# last socket scans
		return

	def get_last_changes(self):
		return self.last_changes

	def get_listeners(self):
		l = []
		for k,v in self.socks_index.iteritems():
			if v.state == TCP_STATE_LISTEN:
				l.append(v)
		return l

	def get_established(self):
		l = []
		for k,v in self.socks_index.iteritems():
			if v.state != TCP_STATE_LISTEN:
				l.append(v)
		return l
	
	def _process_socket_table(self):
		# Read from the kernel and process sockets
		# update the class state members
		# This method would be called periodically to captures socket changes
		scantime = datetime.datetime.now()
	
		self.last_changes  = []
		DWORD = ctypes.c_ulong
		NO_ERROR = 0
		NULL = ""
		bOrder = 0
		
		# define some constants used to identify the state of a TCP port
		MIB_TCP_STATE_CLOSED = 1
		MIB_TCP_STATE_LISTEN = 2
		MIB_TCP_STATE_SYN_SENT = 3
		MIB_TCP_STATE_SYN_RCVD = 4
		MIB_TCP_STATE_ESTAB = 5
		MIB_TCP_STATE_FIN_WAIT1 = 6
		MIB_TCP_STATE_FIN_WAIT2 = 7
		MIB_TCP_STATE_CLOSE_WAIT = 8
		MIB_TCP_STATE_CLOSING = 9
		MIB_TCP_STATE_LAST_ACK = 10
		MIB_TCP_STATE_TIME_WAIT = 11
		MIB_TCP_STATE_DELETE_TCB = 12
		
		ANY_SIZE = 1		 
		
		# defing our row structures
		class MIB_TCPROW(ctypes.Structure):
			_fields_ = [('dwState', DWORD),
				('dwLocalAddr', DWORD),
				('dwLocalPort', DWORD),
				('dwRemoteAddr', DWORD),
				('dwRemotePort', DWORD)]
		
		class MIB_UDPROW(ctypes.Structure):
			_fields_ = [('dwLocalAddr', DWORD),
				('dwLocalPort', DWORD)]
	  
		dwSize = DWORD(0)
		
		# call once to get dwSize 
		ctypes.windll.iphlpapi.GetTcpTable(NULL, ctypes.byref(dwSize), bOrder)
		ANY_SIZE = dwSize.value
		
		class MIB_TCPTABLE(ctypes.Structure):
			_fields_ = [('dwNumEntries', DWORD),
				('table', MIB_TCPROW * ANY_SIZE)]
		
		tcpTable = MIB_TCPTABLE()
		tcpTable.dwNumEntries = 0 # define as 0 for our loops sake
	
		# now make the call to GetTcpTable to get the data
		if (ctypes.windll.iphlpapi.GetTcpTable(ctypes.byref(tcpTable), 
			ctypes.byref(dwSize), bOrder) == NO_ERROR):
			maxNum = tcpTable.dwNumEntries
	
			placeHolder = 0
			
			# loop through every connection
			while placeHolder < maxNum:
			
				item = tcpTable.table[placeHolder]
				placeHolder += 1
				
				lport = socket.ntohs(item.dwLocalPort)
				lip = socket.inet_ntoa(struct.pack('L', item.dwLocalAddr))
				rport = socket.ntohs(item.dwRemotePort)
				rip = socket.inet_ntoa(struct.pack('L', item.dwRemoteAddr))
				state = item.dwState
				index = '%s:%s:%s:%s:%s' % (lip, lport, rip, rport, 'tcp')
				
				# process the existing sockets
				if index in self.socks_index:
					self.socks_index[index].seen(scantime)
				else:
					self.socks_index[index] = SocketConnection(
							lip,
							lport,
							rip,
							rport,
							'tcp',
							state)
					if state == MIB_TCP_STATE_ESTAB:
						self.last_changes.append((SOCK_CHANGE_TYPE_NEW_ESTAB,
							self.socks_index[index]))
					elif state == MIB_TCP_STATE_LISTEN:
						self.last_changes.append((SOCK_CHANGE_TYPE_NEW_LISTEN,
							self.socks_index[index]))
					else:
						self.last_changes.append((SOCK_CHANGE_TYPE_CLOSED,
							self.socks_index[index]))
				
		else:
			print "Error occurred when trying to get TCP Table"
	
		dwSize = DWORD(0)
		# call once to get dwSize
		ctypes.windll.iphlpapi.GetUdpTable(NULL, ctypes.byref(dwSize), bOrder)
		ANY_SIZE = dwSize.value # again, used out of convention 
		
		class MIB_UDPTABLE(ctypes.Structure):
			_fields_ = [('dwNumEntries', DWORD),
						('table', MIB_UDPROW * ANY_SIZE)]  
						
		udpTable = MIB_UDPTABLE()
		udpTable.dwNumEntries = 0 # define as 0 for our loops sake
		
		# now make the call to GetUdpTable to get the data
		if (ctypes.windll.iphlpapi.GetUdpTable(ctypes.byref(udpTable), 
			ctypes.byref(dwSize), bOrder) == NO_ERROR):
		
			maxNum = udpTable.dwNumEntries
			placeHolder = 0
			while placeHolder < maxNum:
	
				item = udpTable.table[placeHolder]
				placeHolder += 1
				lport = socket.ntohs(item.dwLocalPort)
				lip = socket.inet_ntoa(struct.pack('L', item.dwLocalAddr))
				index = '%s:%s:%s:%s:%s' % (lip, lport, '0', '0', 'udp')
	
				# process the existing sockets
				if index in self.socks_index:
					self.socks_index[index].seen(scantime)
				else:
					self.socks_index[index] = SocketConnection(
							lip,
							lport,
							'0',
							'0',
							'udp',
							'0')
					self.last_changes.append((SOCK_CHANGE_TYPE_NEW_LISTEN,
						self.socks_index[index]))
		else:
			print "Error occurred when trying to get UDP Table"
		
		# Process 'unseen' sockets from the prior scan
		l = []
		for k,v in self.socks_index.iteritems():
			if v.last_seen < scantime:
				self.last_changes.append((SOCK_CHANGE_TYPE_CLOSED, v))
				l.append(k)
		for k in l:
			del self.socks_index[k]
			
		return


if __name__ == "__main__":
	import time
	f = SocketStateFollower()
	while 1:
		f._process_socket_table()
		cl = f.get_last_changes()
		print 'Socket changes since last run (30 seconds):'
		for t in cl:
			i = t[1]
			print '%s:%s %s:%s %s %s %i' % (i.lip, i.lport, i.rip, i.rport, i.proto, str(i.state), t[0])
		print ''
		time.sleep(30)