#!/usr/bin/env python
#
#       streamExtractor.py
#       
#       Copyright 2009 Serge Gorbunov <sgorbunov@hotmail.com> 
#		<http://gserge.com>
#       
#       This program is free software; you can redistribute it and/or modify
#       it under the terms of the GNU General Public License as published by
#       the Free Software Foundation; either version 2 of the License, or
#       (at your option) any later version.
#       
#       This program is distributed in the hope that it will be useful,
#       but WITHOUT ANY WARRANTY; without even the implied warranty of
#       MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#       GNU General Public License for more details.
#       
#       You should have received a copy of the GNU General Public License
#       along with this program; if not, write to the Free Software
#       Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#       MA 02110-1301, USA.

# Scapy collection of classes is required to this program.
# Download: http://www.secdev.org/projects/scapy/

import sys, getopt, urlparse, string
import re, os, tempfile
import base64
from   scapy.all import *

class streamExtractor():
	# Function: extracts data streams from outputDirNamea file
	#
	# Input:    inFilename - pcap file name
	#           port       - port number
	#
	# Output:   A list of data streams extracted for the specified port number
	def extractStreams( self, inFileName, port ):
		if ( None == inFileName ):
			return False
		
		# List of packet objects in easy to use format
		streams    = []
		packetList = [] 
		# Flags
		synFlag = 02 
		
		packets = rdpcap( inFileName )
		
		index = 0
		while ( index < len(packets) ):
			packet = packets[index]
			# Every new syn packets identifies start of a new strem
			# We simply add all packets that belong to the stream
			# to packetList and then packetList is added to streams array. 
			# Only packets specified by the destanation port are extracted. 
			
			try:
				if ( packet[Ether][IP][TCP].flags == synFlag ):
					 
					#print "im here"
					#print "%s" % (packet[Ehter][IP][TCP].dport)
					src   = packet[Ether][IP].src
					dest  = packet[Ether][IP].dst
					sport = packet[Ether][IP][TCP].sport
					dport = packet[Ether][IP][TCP].dport
					
					if ( dport == port ):
						packetList.append(packet)
					else:
						index = index + 1
						next 
					# Continue adding all packets in this stream
					while ( True and index < len(packets) ):
						index = index + 1
						# n prefix for the next packet
						try:
							npacket = packets[index]
							nsrc    = npacket[Ether][IP].src
							ndest   = npacket[Ether][IP].dst
							nsport  = npacket[Ether][IP][TCP].sport
							ndport  = npacket[Ether][IP][TCP].dport
							if ( ( src == nsrc and  dest == ndest and sport == nsport and dport == ndport ) or
								 ( src == ndest and dest == nsrc and sport == ndport and dport == nsport ) ):
								packetList.append(npacket)
							else:
								index = index - 1
								break	
						except:
							pass
					streams.append(packetList)
					packetList = []			
		
				else:
					index = index + 1
					pass
			except:
				index = index + 1
				pass
				
		return streams