#!/usr/bin/env python
# -----------------------------------------------------------------------
# Danbooru Upload script
# Written by Christophe 'CSCMEU' Nowicki
# -----------------------------------------------------------------------
# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------
# Danbooru upload script
# Copyright (C) 2009 Christophe Nowicki
#
# 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 MER-
# CHANTABILITY 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.,
# 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
# -----------------------------------------------------------------------

__author__           = "Christophe 'CSCMEU' Nowicki"
__author_email__     = 'cscm@csquad.org'
__maintainer__       = __author__
__maintainer_email__ = __author_email__
__version__          = '0.1a'

from mimetypes import guess_type
from os.path import walk,isfile,isdir,join,getsize,basename
from sys import exit,stderr,argv
from getopt import GetoptError, getopt
from xml.sax import parseString
from xml.sax.handler import ContentHandler
from hashlib import sha1, md5
from httplib import HTTPConnection, HTTPSConnection
import socket

def directory_tree_callback(arg, dirname, filenames):
	"""
	Callback fonction for os.path.walk
	Filter image files
	"""
	for filename in filenames:
		if isdir(join(dirname,filename)) == True:
			if arg[0] == False:
				filenames.remove(filename)
		if isfile(join(dirname,filename)) == True:
			type = guess_type(join(dirname,filename))[0]
			if type != None:
				if type.split('/')[0] == 'image':
					arg[1].append([ join(dirname,filename), type, getsize(join(dirname,filename)) ])

class DanbooruResponseSAXHandler(ContentHandler):
	def __init__(self):
		self.success = None
		self.reason = None
	def startElement(self, name, attrs):
		if name == 'response':
			for attr in attrs.getNames():
				if attr == 'success':
					self.success = attrs.getValue(attr)
				if attr == 'reason':
					self.reason = attrs.getValue(attr)

class DanbooruUpload:
	def __init__(self, usage):
		self.opt = {}
		self._argv = argv[1:]
		self._usage = usage
		self._cookies = {}
		return

	def usage(self):
		print >> stderr, self._usage
		exit(1)

	def getopt(self, shortopts, longopts=[]):
		"""
		Get options and merge into dict 'opt'"
		Check mandatory arguments
		"""
		try:
			options, self._argv = getopt(self._argv, shortopts, longopts)
		except GetoptError:
			self.usage()
		self.opt.update(dict(options))
		if self.opt.get('-u') == None or self.opt.get('-p') == None:
			self.usage()


	def shift(self):
		"""
        	pop first of remaining arguments (shift)"
		"""
		try:
			return self._argv.pop(0)
		except IndexError:
			self.usage()

	def walk(self, directory):
		"""
        	Walk directroy content (recursivly)"
		"""
		files = []
		if self.opt['-r'] == None:
			walk(directory, directory_tree_callback, (False, files))
		else:
			walk(directory, directory_tree_callback, (True, files))
		return files

	def post_multipart(self, selector, fields, file):
		"""
		Post fields and files to an http host as multipart/form-data.
		fields is a sequence of (name, value) elements for regular form fields.
		files is a sequence of (name, filename, value) elements for data to be uploaded as files
		Return the server's response page.
		"""
		content_type, body = self.encode_multipart_formdata(fields, file)
		if self.opt['-s'] == None:
			h = HTTPConnection(self.opt['-h'])
		else:
			h = HTTPSConnection(self.opt['-h'], key_file=self.opt['-k'], cert_file=self.opt['-c'])
		h.putrequest('POST', selector)
		h.putheader('Content-type', content_type)
		for (key, value) in self._cookies.items():
			h.putheader('Cookie', key + '=' + value)
		h.putheader('Content-length', str(len(body)))
		h.endheaders()
		h.send(body)
		response = h.getresponse()
		return(response.read())

	def encode_multipart_formdata(self, fields, files):
	    	"""
		fields is a sequence of (name, value) elements for regular form fields.
		files is a sequence of (name, filename, conent) elements for data to be uploaded as files
		Return (content_type, body) ready for httplib.HTTP instance
		"""
		BOUNDARY = '----------ThIs_Is_tHe_bouNdaRY_$'
		CRLF = '\r\n'
		L = []
		for (key, value) in fields:
			L.append('--' + BOUNDARY)
			L.append('Content-Disposition: form-data; name="%s"' % key)
			L.append('')
			L.append(value)
		for (key, filename, type, content) in files:
			L.append('--' + BOUNDARY)
			L.append('Content-Encoding: binary')
			L.append('Content-Disposition: form-data; name="%s"; filename="%s"' % (key, filename))
			L.append('Content-Type: %s' % type )
			L.append('')
			L.append(content)
			L.append('--' + BOUNDARY + '--')
			L.append('')
		body = CRLF.join(L)
		content_type = 'multipart/form-data; boundary=%s' % BOUNDARY
		return content_type, body

	def login(self):
		password_hash = sha1("choujin-steiner--" + self.opt['-p'] + "--").hexdigest()
		base_url = '/post/index.xml?login=test&password_hash=' + password_hash 
		if self.opt['-v'] == '':
			print "Login on %s with user %s..." % (self.opt['-h'], self.opt['-u'])
		if self.opt['-s'] == None:
			h = HTTPConnection(self.opt['-h'])
		else:
			h = HTTPSConnection(self.opt['-h'], key_file=self.opt['-k'], cert_file=self.opt['-c'])
		h.putrequest('GET', base_url)
		try:
			h.endheaders()
		except socket.error,e:
			return False
		h.send('')
		response = h.getresponse()
		if response.status != 200:
			return False
		header_cookie = response.getheader('Set-Cookie')
		cookies = header_cookie[header_cookie.rfind(self.opt['-h']):header_cookie.rfind(';')] 
	 	(key, value) = cookies.split('=')
		self._cookies[key] = value
		return True

	def upload(self, files):
		for file in files:
			content = open(file[0], mode='rb').read(file[2])
			response = self.post_multipart('/post/create.xml', [ 
				['post[tags]', 'tagme'], 
				['post[source]', ''],
				['md5', md5(content).hexdigest()] 
			],  [ ['post[file]', basename(file[0]), file[1], content] ] )
			handler = DanbooruResponseSAXHandler()
			parseString(response, handler)
			if handler.success == 'false':
				print >> stderr, "Failed to upload %s: %s" % (file[0],handler.reason)
			else:
				if self.opt['-v'] == '':
					print 'Uploading file %s with tag %s.' % (file[0], self.opt['-t'])
			

if __name__ == '__main__':
	uploader = DanbooruUpload("""Usage: 
	danbooru_upload.py [flags] [options] directory
	where valid flags are:
	 -v - verbose mode
         -r - recursive mode
	 -s - use https
	and valid options are:
	 -t keywords - image keywords (tagme)
	 -h host     - hostname (danbooru.donmai.us)
         -u user     - username 
         -p pass     - password 
         -k key_file - ssl key file
         -c cert_file- ssl cert file
""")
	uploader.opt['-v'] = None
	uploader.opt['-r'] = None
	uploader.opt['-s'] = None
	uploader.opt['-k'] = None
	uploader.opt['-c'] = None
	uploader.opt['-t'] = 'tagme'
	uploader.opt['-h'] = 'danbooru.donmai.us'
	uploader.getopt('vrsk:h:u:p:')
	directory = uploader.shift()
	files = uploader.walk(directory)
	if uploader.login() != True:
		print >> stderr, "Login failed"
		exit(1)
	uploader.upload(files)