# 
# 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.
# 

import xml.dom.minidom # for xml input/output
import time #for folder creation and sleep

####################################
class xmlfuzz:

	def __init__(self):
		self.file = "fuzz.xml"

	def create_document(self,localhost, localport, desthost, destport, proto):
		doc=xml.dom.minidom.Document()
		conversation_element = doc.createElementNS("", "conversation")
		conversation_element.setAttributeNS("","local_host", localhost)
		conversation_element.setAttributeNS("","local_port", localport)
		conversation_element.setAttributeNS("","dest_host", desthost)
		conversation_element.setAttributeNS("","dest_port", destport)
		conversation_element.setAttributeNS("","proto", proto)
		conversation_element.setAttributeNS("","time",time.strftime("%H:%M:%S"))
		doc.appendChild(conversation_element)
		return doc, conversation_element
	
	def get_document(self, name="fuzz.xml"):
		doc = xml.dom.minidom.parse(name)
		return doc,  doc.getElementsByTagName("conversation")
	
	def get_network_info(self, doc):
		conversation_element = doc.getElementsByTagName("conversation")
		return conversation_element[0].attributes["local_host"].value, conversation_element[0].attributes["local_port"].value,conversation_element[0].attributes["dest_host"].value, conversation_element[0].attributes["dest_port"].value, conversation_element[0].attributes["proto"].value
		
	def add_request(self, doc, conversation_element, request_ID):
		request_element = doc.createElementNS("", "request")
		request_element.setAttributeNS("","ID",str(request_ID))
		request_element.setAttributeNS("","time",time.strftime("%H:%M:%S"))
		request_element.setAttributeNS("","as_is","0")
		conversation_element.appendChild(request_element)
		return request_element
	
#	def set_request_attributes(self, request_element, request_ignore, request_fuzz):
#		request_element.setAttributeNS("","ignore",str(request_ignore))
#		request_element.setAttributeNS("","fuzzing",str(request_fuzz))
	
	def add_as_is(self, doc, request_element):
		request_element.setAttributeNS("", "as_is", "1")
		return request_element
		
	def del_as_is(self, doc, request_element):
		request_element.setAttributeNS("", "as_is", "0")
		return request_element
		
	def add_fuzz(self, doc, request_element, fuzz_from, fuzz_to, fuzz_bo, fuzz_fs, fuzz_obo, fuzz_dict, fuzz_len, fuzz_lfrom, fuzz_lto, fuzz_lvalue, fuzz_lformat):
		fuzz_element = doc.createElementNS("", "fuzz")
#		fuzz_element.setAttributeNS("", "ID", str(fuzz_ID))
		fuzz_element.setAttributeNS("", "from", str(fuzz_from))
		fuzz_element.setAttributeNS("", "to", str(fuzz_to))
		fuzz_element.setAttributeNS("", "bo", str(fuzz_bo))
		fuzz_element.setAttributeNS("", "fs", str(fuzz_fs))
		fuzz_element.setAttributeNS("", "obo", str(fuzz_obo))
		fuzz_element.setAttributeNS("", "dict", str(fuzz_dict))
		fuzz_element.setAttributeNS("", "checklen", str(fuzz_len))
		fuzz_element.setAttributeNS("", "lfrom", str(fuzz_lfrom))
		fuzz_element.setAttributeNS("", "lto", str(fuzz_lto))
		fuzz_element.setAttributeNS("", "lvalue", str(fuzz_lvalue))
		fuzz_element.setAttributeNS("", "lformat", str(fuzz_lformat))
		request_element.appendChild(fuzz_element)
		return fuzz_element

	# Deleted a fuzzing point of a given request
	def del_fuzz(self, doc, request_element, fuzzID):
		
		fuzz_elements = request_element.getElementsByTagName("fuzz")
		request_element.removeChild(fuzz_elements[fuzzID])
	
	def write_to_screen(self, doc):
		print doc.toxml()

	def write_to_file(self, doc):
		file_object = open(self.file, "w")
		doc.writexml(file_object)
		file_object.close()

	def get_requests(self, doc):
		myrequests = doc.getElementsByTagName('request')
		return myrequests

	def get_request(self, doc, id):
		myrequests = doc.getElementsByTagName('request')
		try:
			return myrequests[id]
		except:
			return False
			
	def get_number_requests(self, doc):
		myrequests = doc.getElementsByTagName('request')
		return len(myrequests)
	
	def ignore_request(self, doc, id):
		myrequests = doc.getElementsByTagName('request')
		return myrequests[id].attributes["ignore"].value
	
	def fuzz_bo(self, doc, idrequest, idfuzz):
		myrequests = doc.getElementsByTagName('request')
		myfuzz = myrequests[idrequest].getElementsByTagName("fuzz")
		return myfuzz[idfuzz].attributes["bo"].value
	
	def fuzz_fs(self, doc, idrequest, idfuzz):
		myrequests = doc.getElementsByTagName('request')
		myfuzz = myrequests[idrequest].getElementsByTagName("fuzz")
		return myfuzz[idfuzz].attributes["fs"].value

	def fuzz_int(self, doc, idrequest, idfuzz):
		myrequests = doc.getElementsByTagName('request')
		myfuzz = myrequests[idrequest].getElementsByTagName("fuzz")
		return myfuzz[idfuzz].attributes["obo"].value
	
	def fuzz_dict(self, doc, idrequest, idfuzz):
		myrequests = doc.getElementsByTagName('request')
		myfuzz = myrequests[idrequest].getElementsByTagName("fuzz")
		return myfuzz[idfuzz].attributes["dict"].value

	# Returns the "to" and "from " of a fuzzing point in a request	
	def get_from_to(self, doc, idrequest, idfuzz):
		myrequests = doc.getElementsByTagName('request')
		myfuzz = myrequests[idrequest].getElementsByTagName("fuzz")
		return myfuzz[idfuzz].attributes["from"].value, myfuzz[idfuzz].attributes["to"].value


	# Returns all the values needed for the variable lenght field	
	def get_len(self, doc, idrequest, idfuzz):
		myrequests = doc.getElementsByTagName('request')
		myfuzz = myrequests[idrequest].getElementsByTagName("fuzz")
		return myfuzz[idfuzz].attributes["checklen"].value, myfuzz[idfuzz].attributes["lfrom"].value, myfuzz[idfuzz].attributes["lto"].value, myfuzz[idfuzz].attributes["lvalue"].value, myfuzz[idfuzz].attributes["lformat"].value

	
	# Returns the number of fuzzing points of a request (id)
	def get_number_fuzz(self, doc, id):
		myrequests = doc.getElementsByTagName('request')
		return len (myrequests[id].getElementsByTagName("fuzz"))

	# Returns an array with the fuzzing points of a request (id)
	def get_fuzzs(self, doc, id): 
		myrequests = doc.getElementsByTagName('request')
		return myrequests[id].getElementsByTagName("fuzz")
