#!/usr/bin/env python2

import templates
import settings, db, os, random, datetime, mimetypes, cgi, tempfile, hashlib, Cookie, urllib, subprocess, json
from PIL import Image
import jab.client

base62_alphabet = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
rfc1123_format = '%a, %d %b %Y %H:%M:%S +0000'
rfc1123_format_tzname = '%a, %d %b %Y %H:%M:%S %Z'

if not os.path.isdir(settings.file_directory):
	os.mkdir(settings.file_directory)

if not os.path.isdir(settings.thumb_directory):
	os.mkdir(settings.thumb_directory)

try:
	# Throws OSError if mogrify doesn't exist.
	subprocess.call(['mogrify', '-quiet'])
except OSError:
	has_mogrify = False
else:
	has_mogrify = True

class BinError(Exception): pass
class InvalidCookieError(BinError): pass
class InactiveLoginError(BinError): pass

class FileUploadFieldStorage(cgi.FieldStorage):
	def make_file(self, binary = None):
		# Make a temporary file in the destination directory, which will be renamed on completion.
		return tempfile.NamedTemporaryFile(prefix = 'upload_', dir = settings.file_directory, delete = True)

class Application(object):
	def get_or_create_user(self, username, jab_id):
		session = db.Session()
		try:
			return session.query(db.User).filter(db.User.jab_id == jab_id).one()
		except db.NoResultFound:
			try:
				user = db.User(username, jab_id)
				session.add(user)
				session.commit()
				session.refresh(user)
				return user
			except db.IntegrityError:
				return None
		finally:
			session.close()

	def get_user_by_jab_id(self, jab_id):
		session = db.Session()
		try:
			return session.query(db.User).filter(db.User.jab_id == jab_id).one()
		except db.NoResultFound:
			return None
		finally:
			session.close()

	def get_user_by_id(self, uid):
		session = db.Session()
		try:
			return session.query(db.User).filter(db.User.id == uid).one()
		except db.NoResultFound:
			return None
		finally:
			session.close()

	def get_file(self, hash, update_accessed = False):
		session = db.Session()
		try:
			f = session.query(db.File).filter(db.File.hash == hash).one()
			if update_accessed:
				f.accessed = datetime.datetime.utcnow()
				session.add(f)
				session.commit()
				# Refresh after field update.
				session.refresh(f)
			return f
		except db.NoResultFound:
			return None
		finally:
			session.close()

	def add_file(self, path, filename, file_hash, user = None, ip = None):
		hash = ''.join(random.choice(base62_alphabet) for x in xrange(5))
		new_path = os.path.join(settings.file_directory, hash + os.path.splitext(filename)[1])
		os.rename(path, new_path)
		if hasattr(settings, 'destination_mode'):
			os.chmod(new_path, settings.destination_mode)

		session = db.Session()
		try:
			file = db.File(hash, file_hash, filename, datetime.datetime.utcnow(), user.id if user else None, ip)
			session.add(file)
			session.commit()
		finally:
			session.close()
		
		return hash

	def get_files(self, user):
		session = db.Session()
		try:
			session.add(user)
			files = user.files
		except db.NoResultFound:
			return []
		finally:
			session.close()

		return files

	def validate_cookie(self, environ):
		cookie = Cookie.SimpleCookie(environ['HTTP_COOKIE'] if 'HTTP_COOKIE' in environ else None)
		if not cookie or not 'token' in cookie:
			return

		token = cookie['token'].value
		try:
			user = self.jab.get_user_by_token(token, settings.jab_identifier, environ['REMOTE_ADDR'])
			user = self.get_user_by_id(user['token_data']['user_id'])
		except jab.client.InvalidCredentialsError:
			user = None
		return user

	def get_file_by_file_hash(self, file_hash):
		session = db.Session()
		try:
			return session.query(db.File).filter(db.File.file_hash == file_hash).one()
		except db.NoResultFound:
			return None
		finally:
			session.close()

	def delete_file(self, file):
		session = db.Session()
		try:
			session.delete(file)
			session.commit()
			os.unlink(file.get_path())
			thumbfile = os.path.join(settings.thumb_directory, file.hash + '.jpg')
			if os.path.exists(thumbfile):
				os.unlink(thumbfile)
		except:
			raise
		finally:
			session.close()

	def redirect(self, environ, start_response, dest):
		start_response('302 Found', [('Location', settings.virtual_root + dest)])
		return []

	def not_modified(self, environ, date):
		if not 'HTTP_IF_MODIFIED_SINCE' in environ:
			return False
		try:
			mod_since_date = datetime.datetime.strptime(environ['HTTP_IF_MODIFIED_SINCE'], rfc1123_format)
		except ValueError:
			# some clients use timezone names (eg. GMT) instead of numeric timezones
			mod_since_date = datetime.datetime.strptime(environ['HTTP_IF_MODIFIED_SINCE'], rfc1123_format_tzname)
		return date == mod_since_date

	def file(self, environ, start_response, path):
		hash = path[1]
		if '.' in hash:
			hash = hash.split('.')[0]
		file = self.get_file(hash, True)
		filename = file.get_path() if file else None
		if file == None or filename == None:
			start_response('404 Not Found', [('Content-Type', 'text/html')])
			return ['<h1>Not Found</h1><p>The file you requested does not exist.</p>']

		# strip microseconds
		if self.not_modified(environ, file.date - datetime.timedelta(microseconds = file.date.microsecond)):
			start_response('304 Not Modified', [('Last-Modified', file.date.strftime(rfc1123_format))])
			return []

		do_range = 'HTTP_RANGE' in environ
		if do_range:
			file_range = environ['HTTP_RANGE'].split('bytes=')[1]

		mime = mimetypes.guess_type(file.filename, strict = False)[0] or 'application/octet-stream'

		# X-Sendfile handling
		if settings.use_xsendfile:
			headers = [('Content-Type', mime), ('Last-Modified', file.date.strftime(rfc1123_format))]
			if do_range:
				headers.append(('X-Sendfile2', '{filename} {range}'.format(filename = urllib.quote(filename.encode('utf8')), range = file_range)))
				if file_range.endswith('-'):
					file_range += str(os.path.getsize(filename)-1)
				headers.append(('Content-Range', 'bytes {range}/{size}'.format(range = file_range, size = os.path.getsize(filename))))
				status = '206 Partial Content'
			else:
				headers.append(('X-Sendfile', filename.encode('utf8')))
				status = '200 OK'
			start_response(status, headers)
			return []

		# Range handling
		if do_range:
			start, end = [int(x or 0) for x in file_range.split('-')]
			size = os.path.getsize(filename)

			if end == 0:
				end = size-1

			write_out = start_response('206 Partial Content', [('Content-Type', mime),
				('Content-Range', 'bytes {start}-{end}/{size}'.format(start = start, end = end, size = size)),
				('Content-Length', str(end - start + 1)), ('Last-Modified', file.date.strftime(rfc1123_format))])

			f = open(filename, 'rb')
			f.seek(start)
			remaining = end-start+1
			s = f.read(min(remaining, 1024))
			while s:
				write_out(s)
				remaining -= len(s)
				s = f.read(min(remaining, 1024))
			return []

		start_response('200 OK', [('Content-Type', mime), ('Content-Length', str(os.path.getsize(filename))),
			('Last-Modified', file.date.strftime(rfc1123_format))])
		return open(filename, 'rb')

	def upload(self, environ, start_response, path):
		user = self.validate_cookie(environ)
		tempfile.tempdir = settings.file_directory
		form = FileUploadFieldStorage(fp = environ['wsgi.input'], environ = environ)
		if environ['REQUEST_METHOD'] != 'POST' or not 'file' in form or not 'filename' in form:
			if user or settings.allow_anonymous_uploads:
				start_response('200 OK', [('Content-Type', 'text/html')])
				return [str(templates.upload(searchList = {'settings': settings, 'user': user}))]
			else:
				return self.redirect(environ, start_response, 'l')

		if not user and not settings.allow_anonymous_uploads:
			form['file'].file.delete = True
			start_response('403 Forbidden', [('Content-Type', 'text/plain')])
			return ['Anonymous uploads are disabled by the administrator.']

		filename = form.getvalue('filename')
		temp = form['file'].file
		temp.delete = False

		# If the name attribute is missing, assume this is a StringIO object, then create a new temporary file and copy the contents.
		if not hasattr(temp, 'name'):
			new_temp = tempfile.NamedTemporaryFile(prefix = 'upload_', dir = settings.file_directory, delete = False)
			new_temp.write(temp.getvalue())
			new_temp.seek(0)
			temp = new_temp

		m = hashlib.md5()
		s = temp.read(128)
		while len(s):
			m.update(s)
			s = temp.read(128)
		temp.close()
		file_hash = m.hexdigest()

		f = self.get_file_by_file_hash(file_hash)
		# TODO: Currently users uploading existing files won't get their files added to their account.
		if f:
			hash = f.hash
		else:
			# temp.name will be moved to the destination filename
			hash = self.add_file(temp.name, filename, file_hash, user, environ['REMOTE_ADDR'])
			# This avoids silly "not bound to a Session" errors when trying to use a newly added file object.
			f = self.get_file(hash)

		# If temp.name still exists, we most likely uploaded a file whose file hash already exists, so just delete the file.
		if os.path.exists(temp.name):
			os.unlink(temp.name)

		mime = f.get_mime_type()
		# TODO: Apparently TIFF also supports EXIF, test this.
		if has_mogrify and mime == 'image/jpeg':
			# NOTE: PIL doesn't support lossless rotation, so we call mogrify to do this.
			# NOTE: This changes the file, so the file_hash applies to the ORIGINAL file contents only.
			# NOTE: The file hash is only used to detect duplicates when uploading, so this should not be a problem.
			subprocess.call(['mogrify', '-auto-orient', f.get_path()])

		if 'api' in form:
			start_response('200 OK', [('Content-Type', 'text/plain')])
			return ['OK {hash}'.format(hash = hash)]
		else:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return [str(templates.uploaded(searchList = {
				'settings': settings,
				'user': user,
				'hash': hash,
				'filename': filename,
				'ext': os.path.splitext(filename)[1],
				'scheme': environ['wsgi.url_scheme'],
				'host': environ['HTTP_HOST'],
			}))]

	def login(self, environ, start_response, path):
		user = self.validate_cookie(environ)
		form = cgi.FieldStorage(fp = environ['wsgi.input'], environ = environ)
		if not user and not 'request_token' in form and not 'verified' in form:
			from urlparse import urljoin
			request_token = self.jab.generate_request_token(settings.jab_identifier, settings.jab_name,
					urljoin('%s://%s' % (environ['wsgi.url_scheme'], environ['HTTP_HOST']), environ['PATH_INFO']))
			if isinstance(request_token, unicode):
				request_token = request_token.encode('utf-8')
			start_response('302 Found', [('Location', urljoin(settings.jab_web_url, 'verify/' + request_token))])
			return []
		if 'request_token' in form and 'verified' in form:
			if form.getvalue('verified') == '0':
				start_response('200 OK', [('Content-Type', 'text/html')])
				return [str(templates.login(searchList = {
					'settings': settings,
					'user': None,
					'error': 'Login was declined.',
					'loggedin': False,
					'next': form.getvalue('next'),
				}))]
			request_token = form.getvalue('request_token')
			try:
				token = self.jab.request_user_token(request_token, settings.jab_identifier, settings.jab_name)
				jab_user = self.jab.get_user_by_token(token, settings.jab_identifier, environ['REMOTE_ADDR'])
			except jab.client.InvalidCredentialsError:
				start_response('200 OK', [('Content-Type', 'text/html')])
				return [str(templates.login(searchList = {
					'settings': settings,
					'user': None,
					'error': 'Failed to request login: invalid token or user.',
					'loggedin': False,
					'next': form.getvalue('next'),
				}))]
			user = self.get_or_create_user(jab_user['username'], jab_user['_id'])
			if not user:
				start_response('500 Internal Server Error', [])
				return []
			self.jab.set_token_data(token, settings.jab_identifier, {'user_id': user.id})
			c = Cookie.SimpleCookie()
			c['token'] = token
			start_response('200 OK', [
				('Content-Type', 'text/html'),
				('Set-Cookie', c['token'].OutputString())
			])
			return [str(templates.login(searchList = {
				'settings': settings,
				'user': user,
				'error': None,
				'loggedin': True,
				'next': form.getvalue('next'),
			}))]

		if user:
			rememberme = 'rememberme' in form
			forever = 'forever' in form

			cookie = Cookie.SimpleCookie(environ['HTTP_COOKIE'])
			token = cookie['token'].value
			c = Cookie.SimpleCookie()
			c['token'] = token

			dt = datetime.datetime.utcnow() + datetime.timedelta(days = 30)
			expires = dt.strftime('%a, %d-%b-%y %H:%M:%S GMT')
			if rememberme:
				c['token']['expires'] = expires
				if forever:
					c['forever'] = 1
					c['forever']['expires'] = expires

			# FIXME: This field is lost when we redirect to jab.
			next = form.getvalue('next')
			headers = [
				('Location', next if next else (settings.virtual_root + 'u')),
				('Set-Cookie', c['token'].OutputString())]
			if 'forever' in c:
				headers.append(('Set-Cookie', c['forever'].OutputString()))
			start_response('302 Found', headers)
			return []

		start_response('404 Not Found', [])
		return []

	def register(self, environ, start_response, path):
		start_response('302 Found', [('Location', settings.jab_web_url + 'register')])
		return []

	def logout(self, environ, start_response, path):
		c = Cookie.SimpleCookie(environ['HTTP_COOKIE'] if 'HTTP_COOKIE' in environ else None)
		if c and 'token' in c:
			try:
				self.jab.expire_user_token(c['token'].value, settings.jab_identifier)
			except:
				pass

		c = Cookie.SimpleCookie()
		expires = datetime.datetime.utcfromtimestamp(0).strftime('%a, %d-%b-%y %H:%M:%S GMT')
		c['forever'] = 0
		c['forever']['expires'] = expires
		c['token'] = 0
		c['token']['expires'] = expires
		start_response('302 Found', [
			('Set-Cookie', c['forever'].OutputString()),
			('Set-Cookie', c['token'].OutputString()),
			('Location', settings.virtual_root)])
		return []

	def changepass(self, environ, start_response, path):
		start_response('302 Found', [('Location', settings.jab_web_url + 'changepass')])
		return []

	def static(self, environ, start_response, path):
		filename = path[1]
		if not filename in ('style.css', 'no-thumbnail.png', 'favicon.ico', 'jquery-2.1.0.min.js', 'jquery.lazy.min.js'):
			start_response('404 Not Found', [])
			return []

		filepath = os.path.join(settings.static_root, filename)
		if not os.path.exists(filepath):
			start_response('404 Not Found', [])
			return []
		mime = mimetypes.guess_type(filename, strict = False)[0] or 'application/octet-stream'
		start_response('200 OK', [('Content-Type', mime)])
		return open(filepath, 'rb')

	def help(self, environ, start_response, path):
		user = self.validate_cookie(environ)
		start_response('200 OK', [('Content-Type', 'text/html')])
		return [str(templates.help(searchList = {
			'settings': settings,
			'user': user,
			'scheme': environ['wsgi.url_scheme'],
			'host': environ['HTTP_HOST'],
		}))]

	def my_files(self, environ, start_response, path):
		user = self.validate_cookie(environ)
		if user == None:
			start_response('302 Found', [('Location', settings.virtual_root + 'l?' + urllib.urlencode({'next': settings.virtual_root + 'm'}))])
			return []
		files = self.get_files(user)
		start_response('200 OK', [('Content-Type', 'text/html')])
		return [str(templates.my(searchList = {
			'settings': settings,
			'user': user,
			'files': files,
			'total_size': db.File.pretty_size(sum([f.get_size() for f in files])),
		}))]

	def images(self, environ, start_response, path):
		user = self.validate_cookie(environ)
		if user == None:
			start_response('302 Found', [('Location', settings.virtual_root + 'l?' + urllib.urlencode({'next': settings.virtual_root + 'i'}))])
			return []
		files = [f for f in self.get_files(user) if f.is_image()]
		start_response('200 OK', [('Content-Type', 'text/html')])
		return [str(templates.images(searchList = {
			'settings': settings,
			'user': user,
			'files': files,
			'total_size': db.File.pretty_size(sum([f.get_size() for f in files])),
		}))]

	def thumb(self, environ, start_response, path):
		hash = path[1]
		thumbfile = os.path.join(settings.thumb_directory, hash + '.jpg')
		if not os.access(thumbfile, os.F_OK):
			file = self.get_file(hash)
			try:
				im = Image.open(file.get_path())
			except IOError:
				# We can't generate a thumbnail for this file, just say it doesn't exist.
				start_response('404 Not Found', [])
				return []
			# Check for valid JPEG modes.
			if im.mode not in ('1', 'L', 'RGB', 'RGBA', 'RGBX', 'CMYK', 'YCbCr'):
				im = im.convert('RGB')
			im.thumbnail(settings.thumb_size, Image.ANTIALIAS)
			im.save(thumbfile)

		date = datetime.datetime.utcfromtimestamp(os.path.getmtime(thumbfile))
		if self.not_modified(environ, date):
			start_response('304 Not Modified', [('Last-Modified', date.strftime(rfc1123_format))])
			return []

		start_response('200 OK', [('Content-Type', 'image/jpeg'), ('Last-Modified', date.strftime(rfc1123_format))])
		return open(thumbfile, 'rb')

	def delete(self, environ, start_response, path):
		user = self.validate_cookie(environ)
		if user == None:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return ['Not logged in.']
		hash = path[1]
		file = self.get_file(hash)
		if file == None:
			start_response('404 Not Found', [('Content-Type', 'text/html')])
			return ['<h1>Not Found</h1><p>The file you requested does not exist.</p>']
		if file.user_id != user.id:
			start_response('403 Forbidden', [('Content-Type', 'text/html')])
			return ['<h1>Forbidden</h1><p>You are not allowed to delete this file.</p>']
		if environ['REQUEST_METHOD'] == 'POST':
			try:
				self.delete_file(file)
			except Exception as e:
				start_response('500 Internal Error', [('Content-Type', 'text/html')])
				return ['Failed to delete file {filename} ({error}).'.format(filename = file.filename, error = str(e))]
			else:
				start_response('302 Found', [('Location', settings.virtual_root + 'u')])
				return []
		else:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return [str(templates.delete(searchList = {
				'settings': settings,
				'user': user,
				'hash': hash,
				'filename': file.filename,
			}))]

	def api(self, environ, start_response, path):
		def error(msg):
			start_response('200 OK', [('Content-Type', 'application/json')])
			return [json.dumps({'status': False, 'message': msg})]
		user = self.validate_cookie(environ)
		form = cgi.FieldStorage(fp = environ['wsgi.input'], environ = environ)
		method = form.getvalue('method')
		data = {'status': False, 'method': method, 'message': None}
		if method in ('list', 'images'):
			if user == None:
				return error('Not logged in')
			files = self.get_files(user)
			data['files'] = [
				{
					'name': f.filename,
					'hash': f.hash,
					'date': int(f.date.strftime('%s')),
					'size': f.get_size(),
				}
				for f in files if method == 'list' or (method == 'images' and f.is_image())
			]
			data['status'] = True
		elif method == 'get_token':
			try:
				token = self.jab.generate_user_token(form['username'].value, form['password'].value, settings.jab_identifier, '%s (API)' % settings.jab_name)
				jab_user = self.jab.get_user_by_token(token, settings.jab_identifier, environ['REMOTE_ADDR'])
				user = self.get_or_create_user(jab_user['username'], jab_user['_id'])
				if not user:
					return error('Error fetching user data')
				self.jab.set_token_data(token, settings.jab_identifier, {'user_id': user.id})
			except:
				return error('Invalid credentials')
			data['token'] = token
			data['status'] = True
		elif method == 'expire_token':
			try:
				self.jab.expire_token(form['token'].value, settings.jab_identifier)
			except:
				pass
			data['status'] = True
		elif method == 'test_token':
			try:
				user = self.jab.get_user_by_token(form['token'].value, settings.jab_identifier, environ['REMOTE_ADDR'])
			except jab.client.InvalidCredentialsError:
				return error('Invalid token')
			except:
				return error('Error fetching user data')
			data['status'] = True
		else:
			data['message'] = 'Unknown method "%s"' % method
		start_response('200 OK', [('Content-Type', 'application/json')])
		return [json.dumps(data)]

	f = file
	u = upload
	l = login
	s = static
	h = help
	m = my_files
	i = images
	t = thumb
	o = logout
	r = register
	c = changepass
	d = delete
	a = api

	def __init__(self):
		self.jab = jab.client.JabClient(settings.jab_socket)

	def __call__(self, environ, start_response):
		def never_forget_wrapper(status, headers, exc_info = None):
			c = Cookie.SimpleCookie(environ['HTTP_COOKIE'] if 'HTTP_COOKIE' in environ else None)
			if 'forever' in c and c['forever'].value == '1' and not any(x[0] == 'Set-Cookie' for x in headers):
				dt = datetime.datetime.utcnow() + datetime.timedelta(days = 30)
				expires = dt.strftime('%a, %d-%b-%y %H:%M:%S GMT')
				for k in c.keys():
					c[k]['path'] = settings.virtual_root
					c[k]['expires'] = expires
					headers.append(('Set-Cookie', c[k].OutputString()))
			return start_response(status, headers, exc_info)
		try:
			path = environ['PATH_INFO'].split('/')[1:]
			module = path[0] if len(path) else ''
			if len(module) and module in 'fulshmitorcda':
				return getattr(self, module)(environ, never_forget_wrapper, path)
			elif path == ['favicon.ico']:
				return self.static(environ, start_response, ['s'] + path)
			else:
				start_response('302 Found', [('Location', settings.virtual_root + 'u')])
				return []
		except (InactiveLoginError, InvalidCookieError):
			return self.logout(environ, start_response, path)

if __name__ == '__main__':
	import sys
	if len(sys.argv) == 3:
		from flup.server.fcgi import WSGIServer
		WSGIServer(Application(), bindAddress = (sys.argv[1], int(sys.argv[2]))).run()
	else:
		from wsgiref.simple_server import make_server, WSGIServer
		# enable IPv6
		WSGIServer.address_family |= 10
		http = make_server('', 8000, Application())
		http.serve_forever()

# vim: noet ts=4