#!/usr/bin/env python
"""Modify the packets in a pcap file, returning a new one with the modified packets."""

#V0.4

#Processes around 386 pps

from scapy.all import sniff, PcapWriter, IP, Ether, ARP, UDP, TCP

FIN = 0x01
SYN = 0x02
RST = 0x04
PSH = 0x08
ACK = 0x10
URG = 0x20
ECE = 0x40
CWR = 0x80

four_b = (2 ** 32) - 1		#Four billion; used in sequence numbers that wrap at 2^32 -1

tuple_stats = {}		#Dictionary: keys are port_tuple strings, values are lists: [packets_written, start_seq, payload_bytes_written] .  Start_seq used for tcp, bytes_written for udp.
base_timestamp = None		#Placed as global as we don't control the packet handling loop
bpfilter = ''


def tuple_string(c_proto, pkt, src_ip, dst_ip):
	"""Create a string desribing the connection as a key for the tuple_stats dictionary."""

	if c_proto == 'TCP6' or c_proto == 'TCP4':
		sport = str(pkt['TCP'].sport)
		dport = str(pkt['TCP'].dport)
	elif c_proto == 'UDP6' or c_proto == 'UDP4':
		udp_layer = pkt.getlayer(UDP)
		sport = str(udp_layer.sport)
		dport = str(udp_layer.dport)
	else:
		sport = ''
		dport = ''

	return str(c_proto) + '_' + str(src_ip) + '_' + str(sport) + '_' + str(dst_ip) + '_' + str(dport)


def should_write_tcp(p, port_tuple):
	"""Decides whether the packet - aready known to be IPv4 or IPv6 TCP - should be written."""

	global tuple_stats
	#We use, but do not modify, the global dictionary "args", so we don't declare it global.

	write_tcp = True

	F = p['TCP'].flags
	Seq = int(p['TCP'].seq)

	if port_tuple not in tuple_stats:
		tuple_stats[port_tuple] = [0, None, 0]

	if (F & (SYN + FIN + RST)) == SYN:						#Syn set, fin/rst clear
		#Remember the Sequence number for this SYN or SYN/ACK packet
		tuple_stats[port_tuple][1] = Seq + 1					#Since ACK packets start using Seq+1, we store this
	elif F & (SYN + FIN + RST + ACK) == ACK:					#Ack set, Syn/Fin/Rst all false
		if tuple_stats[port_tuple][1] is None:					#We never got an initial Seq from a Syn or syn/ack - we may be coming in the middle of a tcp session.  Just grab the seq number from this packet to get the next (user requested) kilobytes.
			tuple_stats[port_tuple][1] = Seq
		if args['ackcount'] is not None and tuple_stats[port_tuple][0] > args['ackcount']:
			write_tcp = False
		#or args['ackbytes'] is not None and (Seq > tuple_stats[port_tuple][1]) and (Seq - tuple_stats[port_tuple][1] < args['ackbytes']):
		#	#FIXME - handle case where Seq < tuple_stats[port_tuple][1] + 1 by adding 4B
		#	pass
		#else:
		#	write_tcp = False

	return write_tcp


def should_write_udp(p, port_tuple):
	"""Decides whether the packet - aready known to be IPv4 or IPv6 UDP - should be written."""

	global tuple_stats
	#We use, but do not modify, the global dictionary "args", so we don't declare it global.

	write_udp = True

	if port_tuple not in tuple_stats:
		tuple_stats[port_tuple] = [0, None, 0]

	if args['udpcount'] is not None and tuple_stats[port_tuple][0] > args['udpcount']:
		write_udp = False

	return write_udp


def should_write(p):
	"""Returns true if the supplied packet should be written (matches requirements), false otherwise."""

	global tuple_stats
	#We use, but do not modify, the global dictionary "args", so we don't declare it global.

	write_it = True
	port_tuple = 'Pkt____'

	if p.haslayer(Ether) and p[Ether].type == 0x86DD:					#IPv6
		sIP = str(p['IPv6'].src)
		dIP = str(p['IPv6'].dst)
		if p['IPv6'].nh == 6 and p.haslayer(TCP): 					#TCPv6
			if args['ackcount'] is not None or args['ackbytes'] is not None:
				port_tuple = tuple_string('TCP6', p, sIP, dIP)
				write_it = should_write_tcp(p, port_tuple)
		elif (p['IPv6'].nh == 17) and p.haslayer(UDP): 					#UDPv6
			if args['udpcount'] is not None or args['udpbytes'] is not None:
				port_tuple = tuple_string('UDP6', p, sIP, dIP)
				write_it = should_write_udp(p, port_tuple)
		else:
			#record IP Proto?
			port_tuple = tuple_string('IPv6', p, sIP, dIP)				#Mangled TCP/UDP or non-TCP/UDP packet, write unconditionally
	elif p.haslayer(IP) and isinstance(p[IP], IP):						#IPv4
		sIP = str(p['IP'].src)
		dIP = str(p['IP'].dst)
		if p['IP'].proto == 6 and p.haslayer(TCP) and isinstance(p[TCP], TCP):		#TCPv4
			if args['ackcount'] is not None or args['ackbytes'] is not None:
				port_tuple = tuple_string('TCP4', p, sIP, dIP)
				write_it = should_write_tcp(p, port_tuple)
		elif p['IP'].proto == 17 and p.haslayer(UDP) and isinstance(p[UDP], UDP):	#UDPv4
			if args['udpcount'] is not None or args['udpbytes'] is not None:
				port_tuple = tuple_string('UDP4', p, sIP, dIP)
				write_it = should_write_udp(p, port_tuple)
		else:
			#record IP Proto?
			port_tuple = tuple_string('IPv4', p, sIP, dIP)				#Mangled TCP/UDP or non-TCP/UDP packet, write unconditionally
	elif (p.haslayer(Ether) and p[Ether].type == 0x0806) and p.haslayer(ARP) and isinstance(p[ARP], ARP):		#ARP
		port_tuple = tuple_string('ARP', p, sIP, dIP)
	else:
		port_tuple = tuple_string('Other', p, sIP, dIP)
		#p.show()
		#quit()

	if port_tuple and port_tuple not in tuple_stats:
		tuple_stats[port_tuple] = [0, None, 0]

	if write_it:
		tuple_stats[port_tuple][0] += 1

	if args['verbose'] and port_tuple and port_tuple in tuple_stats:
		print(port_tuple + ":   " + str(tuple_stats[port_tuple]) + " " + str(write_it))

	return write_it


def process_a_packet(pkt):
	"""Handle a single packet read from the input pcap file."""

	global out_handle
	global base_timestamp
	#We use, but do not modify, "args" so we don't declare it as global.

	if args['timescale']:
		if not base_timestamp:
			base_timestamp = pkt.time
		delta_time = pkt.time - base_timestamp
		pkt.time = base_timestamp + (args['timescale'] * delta_time)
	if should_write(pkt):
		pmod = pkt
		out_handle.write(pmod)



#======== Main
if __name__ == "__main__":
	import argparse

	parser = argparse.ArgumentParser(description='Process a mail folder')
	parser.add_argument('-r', '--read', help='File from which to read packets', required=True)
	parser.add_argument('-w', '--write', help='File to which to write packets', required=True)
	parser.add_argument('-t', '--timescale', help='Multiplier for time scale', required=False, default=None)
	parser.add_argument('-v', '--verbose', help='Print info about each packet', dest='verbose', action='store_true', required=False, default=False)
	parser.add_argument('--ackcount', help='Maximum number of ACK packets to write on one side of a TCP connection', required=False, type=int, default=None)
	parser.add_argument('--ackbytes', help='ackbytes - not yet implemented', required=False, type=int, default=None)
	parser.add_argument('--udpcount', help='Maximum number of UDP packets to write on one side of a UDP conversation', required=False, type=int, default=None)
	parser.add_argument('--udpbytes', help='udpbytes - not yet implemented', required=False, type=int, default=None)
	args = vars(parser.parse_args())
	#FIXME - add --summary argument and show in and out stats


if args['write']:
	try:
		out_handle = PcapWriter(filename=args['write'])
	except:
		print("Unable to open " + str(args['write']) + " for write, exiting.")
		raise
		quit(1)
try:
	sniff(store=0, offline=args['read'], filter=bpfilter, prn=lambda x: process_a_packet(x))
except IOError:
	print("Unable to open " + str(args['read']) + ' , exiting.')
	raise
	quit(1)