####
# 02/2006 Will Holcomb <wholcomb@gmail.com>
# 
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
# 
# This library 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
# Lesser General Public License for more details.
#
"""
Enables the use of multipart/form-data for posting forms
"""

'''
Inspirations:
  Upload files in python:
	http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/146306
  urllib2_file:
	Fabien Seisen: <fabien@seisen.org>

Example:
  import MultipartPostHandler, urllib2, cookielib

  cookies = cookielib.CookieJar()
  opener = urllib2.build_opener(urllib2.HTTPCookieProcessor(cookies),
								MultipartPostHandler.MultipartPostHandler)
  params = { "username" : "bob", "password" : "riviera",
			 "file" : open("filename", "rb") }
  opener.open("http://wwww.bobsite.com/upload/", params)

Further Example:
  The main function of this file is a sample which downloads a page and
  then uploads it to the W3C validator.
'''

import urllib
import urllib2
import mimetools, mimetypes
import os, stat, md5
from core.data.fuzzer.fuzzer import fileStr

class Callable:
	def __init__(self, anycallable):
		self.__call__ = anycallable

# Controls how sequences are uncoded. If true, elements may be given multiple values by
#  assigning a sequence.
doseq = 1

class MultipartPostHandler(urllib2.BaseHandler):
	handler_order = urllib2.HTTPHandler.handler_order - 10 # needs to run first

	def http_request(self, request):
		data = request.get_data()
		if data is not None and type(data) != str:
			v_files = []
			v_vars = []
			try:
				for(key, value) in data.items():
					if type(value) == file:
						if not value.closed:
							v_files.append((key, value))
						else:
							v_vars.append((key, ''))
					elif hasattr( value, 'isFile'):
						v_files.append((key, value))
					else:
						v_vars.append((key, value))
			except TypeError:
				systype, value, traceback = sys.exc_info()
				raise TypeError, "not a valid non-string sequence or mapping object", traceback

			if len(v_files) == 0:
				data = urllib.urlencode(v_vars, doseq)
			else:
				boundary, data = self.multipart_encode(v_vars, v_files)
				contenttype = 'multipart/form-data; boundary=%s' % boundary
				if(request.has_header('Content-Type') and request.get_header('Content-Type').find('multipart/form-data') != 0):
					print "Replacing %s with %s" % (request.get_header('content-type'), 'multipart/form-data')
				request.add_unredirected_header('Content-Type', contenttype)

			request.add_data(data)
		return request

	def multipart_encode(vars, files, boundary = None, buffer = None):
		if boundary is None:
			# Before :
			# boundary = mimetools.choose_boundary()
			# '127.0.0.1.1000.6267.1173556103.828.1'
			# This contains my IP address, I dont like that...
			# Now:
			boundary = md5.new(mimetools.choose_boundary()).hexdigest()
		if buffer is None:
			buffer = ''
		
		for(key, value) in vars:
			buffer += '--%s\r\n' % boundary
			buffer += 'Content-Disposition: form-data; name="%s"' % key
			buffer += '\r\n\r\n' + value + '\r\n'
		
		for(key, fd) in files:
			file_size = getFileSize( fd )
			filename = fd.name.split( os.path.sep )[-1]
			contenttype = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
			buffer += '--%s\r\n' % boundary
			buffer += 'Content-Disposition: form-data; name="%s"; filename="%s"\r\n' % (key, filename)
			buffer += 'Content-Type: %s\r\n' % contenttype
			# buffer += 'Content-Length: %s\r\n' % file_size
			fd.seek(0)
			buffer += '\r\n' + fd.read() + '\r\n'
		buffer += '--%s--\r\n\r\n' % boundary
		return boundary, buffer
	multipart_encode = Callable(multipart_encode)

	https_request = http_request

def getFileSize( file ):
	'''
	Aux function to get the file size. Needed if I want to use my modified string to fuzz file content.
	'''
	if type( file ) == fileStr:
		return len( file )
	else:
		return os.fstat(file.fileno())[stat.ST_SIZE]
	
def main():
	import tempfile, sys

	validatorURL = "http://validator.w3.org/check"
	opener = urllib2.build_opener(MultipartPostHandler)

	def validateFile(url):
		temp = tempfile.mkstemp(suffix=".html")
		os.write(temp[0], opener.open(url).read())
		params = { "ss" : "0",			# show source
				   "doctype" : "Inline",
				   "uploaded_file" : open(temp[1], "rb") }
		print opener.open(validatorURL, params).read()
		os.remove(temp[1])
	
	def uploadFile( file ):
		params = { "MAX_FILE_SIZE" : "10000",
				   "uploadedfile" : open( file , "rb") }
		print opener.open( 'http://localhost/w3af/fileUpload/uploader.php', params).read()
		
	if len(sys.argv[1:]) > 0:
		for arg in sys.argv[1:]:
			uploadFile(arg)
	else:
		validateFile("http://www.google.com")

if __name__=="__main__":
	main()
