from twisted.internet.protocol import DatagramProtocol
from twisted.internet import gtk2reactor
# 
# TAOF, the general purpose TCP fuzzer.
# Copyright (C) 2007 Rodrigo Marcos
# 
# 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.
# 

gtk2reactor.install()

from twisted.protocols import portforward
from twisted.internet import reactor

import fuzzutils


####################################
####################################

#UDP Forwarding

class UDPClient(DatagramProtocol):
	def __init__(self, server, host, port):
		self.server = server
		self.host = host
		self.port = port
    
	def startProtocol(self):
		self.transport.connect(fuzzutils.context["dest_host"], int(fuzzutils.context["dest_port"]))
    
	def datagramReceived(self, data, (host, port)):
		self.server.transport.write(data, (self.host, self.port))
 
class UDPServer(DatagramProtocol):
	client = None
    
	def datagramReceived(self, data, (host, port)):
		if not self.client or self.client.host != host or self.client.port != port:
			self.client = UDPClient(self, host, port)
			reactor.listenUDP(0, self.client)

		fuzzutils.context["REQUEST_COUNTER"] += 1	
		fuzzutils.context["logger"].log_request(fuzzutils.context["REQUEST_COUNTER"], data)				#log the actual raw request
		fuzzutils.context["xmllog"].add_request(fuzzutils.context["xmldoc"], fuzzutils.context["xmlconv"], str(fuzzutils.context["REQUEST_COUNTER"])) 	#log into xml file
		self.client.transport.write(data, (fuzzutils.context["dest_host"], int(fuzzutils.context["dest_port"])))
 

#TCP Forwarding

def server_dataReceived(self, data):
	#print 'server recieved data:', data
	fuzzutils.context["REQUEST_COUNTER"] += 1	
	fuzzutils.context["logger"].log_request(fuzzutils.context["REQUEST_COUNTER"], data)				#log the actual raw request

	fuzzutils.context["xmllog"].add_request(fuzzutils.context["xmldoc"], fuzzutils.context["xmlconv"], str(fuzzutils.context["REQUEST_COUNTER"])) 	#log into xml file
	
	portforward.Proxy.dataReceived(self, data)

portforward.ProxyServer.dataReceived = server_dataReceived

def client_dataReceived(self, data):
	#print 'client received data:', data
	#fuzzutils.context["logger"].log_response(fuzzutils.context["REQUEST_COUNTER"], data)
	portforward.Proxy.dataReceived(self, data)

portforward.ProxyClient.dataReceived = client_dataReceived


####################################
		