#! /usr/bin/python
import os,sys,nids,re,pefile,peutils,StringIO,gzip,hashlib
#import pyclamd
###########

#CONFIG OPTIONS
#extractFiles	- Save files with md5 hash names to the local directory
#scanPE		- Scans for PE files
#scanMSO	- Scans for MS Office Files
#scanPDF	- Scans for PDF Files
#clamScan	- Scans file with clamav

extractFiles	= "NO"
scanPE		= "YES"
scanMSO		= "YES"
scanPDF		= "YES"
#clamScan	= "YES"
############

end_states = (nids.NIDS_CLOSE, nids.NIDS_TIMEOUT, nids.NIDS_RESET)
signatures = peutils.SignatureDatabase('UserDB.TXT')
mssig = '\xD0\xCF\x11\xE0\xA1\xB1\x1a\xE1'
pdfsig = '\x25\x50\x44\x46'
mstypes = ['doc','ppt','xls']

def handleTcpStream(tcp):
	found = False
	#print "tcps -", str(tcp.addr), " state:", tcp.nids_state
	if tcp.nids_state == nids.NIDS_JUST_EST:
		((src, sport), (dst, dport)) = tcp.addr
		if dport in (80, 8000, 8080) or sport == 20:
			tcp.client.collect = 1
			tcp.server.collect = 1
	elif tcp.nids_state == nids.NIDS_DATA:
		tcp.discard(0)
	elif tcp.nids_state in end_states:
		toserver = tcp.server.data[:tcp.server.count]
                toclient = tcp.client.data[:tcp.client.count]
		((src, sport), (dst, dport)) = tcp.addr
		print "%s:%s -> %s:%s" % (src,sport,dst,dport)
		scan_tcp_body = None
		proto = None
		if dport in (80, 8000, 8080):
			proto = "HTTP"
			header_len = toclient.find('\r\n\r\n')
			headers = toserver[:header_len].split('\n')[:-2]
			for header in headers:
                                print "%s" % header
			try:
				gzip_data = toclient[header_len+4:]
				gzip_data = StringIO.StringIO(gzip_data)
				gzipper = gzip.GzipFile(fileobj=gzip_data)
				scan_tcp_body = gzipper.read()
			except:
				scan_tcp_body = toclient[header_len+4:]
		elif sport == 20:
			proto = "FTP"
			scan_tcp_body = toserver
		if scan_tcp_body.startswith('MZ') and scanPE == "YES":
			found = True
			if extractFiles == "YES":
				extract_file(scan_tcp_body,'exe')
			pe = pefile.PE(data=scan_tcp_body)
			matches = signatures.match_all(pe, ep_only = True)
			for match in matches:
				print "%s - %s" % (proto, match)
		elif scan_tcp_body.startswith(mssig) and scanMSO == "YES":
			found = True
			mstype = None
			for type in mstypes:
				if get_request.find(type) != -1:
					mstype = type
			print "Office file found - %s" % mstype
			if extractFiles == "YES":
				extract_file(scan_tcp_body, mstype)
				print "File extracted"
		elif scan_tcp_body.startswith(pdfsig) and scanPDF == "YES":
			found = True
			print "PDF found"
			if extractFiles == "YES":
				extract_file(scan_tcp_body, 'pdf')
		#if found and clamScan == "YES":
		#	ret = pyclamd.scan_stream(scan_tcp_body)
		#	print "ClamScan: %s" % ret

def extract_file(data, filetype):
	h = hashlib.md5()
	h.update(data)
	md5 = h.hexdigest()
	print "MD5: %s" % str(md5)
	f = open("%s.%s" % (md5,filetype), 'w')
	f.write(data)
	f.close()	

def main():
	#if clamScan == "YES":
	#	pyclamd.init_network_socket('localhost', 3310)
	nids.param("scan_num_hosts", 0)
	if len(sys.argv) > 1 and str(sys.argv[1]).endswith('.pcap'):
		nids.param("filename", sys.argv[1])
	else:
		nids.param("device", "eth0")
	nids.init()
	nids.register_tcp(handleTcpStream)
	try:
		nids.run()
	except nids.error, e:
		print "nids/pcap error:", e
	except Exception, e:
		print sys.exc_info()

if __name__ == '__main__':
	main()