path: root/fbin.py

#!/usr/bin/env python2

import templates
import settings, db, os, random, datetime, shutil, mimetypes, cgi, tempfile, hashlib, Cookie

base62_alphabet = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'

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

class Application(object):
	def get_user(self, username, password):
		session = db.Session()
		try:
			user = session.query(db.User).filter(db.and_(db.User.username == username, db.User.password == password)).one()
		except db.NoResultFound:
			return None
		finally:
			session.close()

		return user

	def get_user_by_name(self, username):
		session = db.Session()
		try:
			return session.query(db.User).filter(db.User.username == username).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 add_user(self, username, password):
		session = db.Session()
		try:
			user = db.User(username, password)
			session.add(user)
			session.commit()
		except db.IntegrityError:
			return None
		finally:
			session.close()

		return user

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

		return os.path.join(settings.file_directory, hash + os.path.splitext(file.filename)[1])

	def add_file(self, path, filename, file_hash, user = 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])
		shutil.copyfile(path, new_path)

		session = db.Session()
		try:
			file = db.File(hash, file_hash, filename, datetime.datetime.utcnow(), user.id if user else None)
			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, cookie):
		if not cookie:
			return None

		identifier = cookie['identifier'].value

		if 'username' in cookie:
			user = self.get_user_by_name(cookie['username'].value)
			if not user:
				return None
			digest = hashlib.sha1(user.username + user.password).hexdigest()
			return user if (digest == identifier) else None

		user = self.get_user_by_id(cookie['uid'].value)
		if not user:
			return None

		digest = hashlib.sha1(str(user.id) + user.password).hexdigest()
		return user if (digest == identifier) else None

	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 file(self, environ, start_response, path):
		hash = path[1]
		if '.' in hash:
			hash = hash.split('.')[0]
		filename = self.get_file_path(hash)
		if 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>']

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

		start_response('200 OK', [('Content-Type', mime), ('Content-Length', str(os.path.getsize(filename)))])
		return open(filename, 'rb')

	def upload(self, environ, start_response, path):
		c = Cookie.SimpleCookie(environ['HTTP_COOKIE'] if 'HTTP_COOKIE' in environ else None)
		user = self.validate_cookie(c)
		form = cgi.FieldStorage(fp = environ['wsgi.input'], environ = environ)
		if environ['REQUEST_METHOD'] != 'POST' or not 'file' in form or not 'filename' in form:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return str(templates.upload(searchList = {'user': user}))

		filename = form.getvalue('filename')

		temp = tempfile.NamedTemporaryFile(mode = 'wb', prefix = 'fbin', delete = True)
		temp.write(form.getvalue('file'))
		temp.flush()

		m = hashlib.md5()
		with open(temp.name) as f:
			s = f.read(128)
			while len(s):
				m.update(s)
				s = f.read(128)
		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:
			hash = self.add_file(temp.name, filename, file_hash, user)

		temp.close()

		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 = {
				'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):
		c = Cookie.SimpleCookie(environ['HTTP_COOKIE'] if 'HTTP_COOKIE' in environ else None)
		user = self.validate_cookie(c)
		form = cgi.FieldStorage(fp = environ['wsgi.input'], environ = environ)
		if environ['REQUEST_METHOD'] != 'POST' or not 'username' in form or not 'password' in form:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return str(templates.login(searchList = {
				'user': user,
				'error': None,
			}))

		username = form.getvalue('username')
		password = hashlib.sha1(form.getvalue('password')).hexdigest()

		user = self.get_user(username, password)

		if user == None:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return str(templates.login(searchList = {
				'user': user,
				'error': 'Login failed',
			}))

		c = Cookie.SimpleCookie()
		c['uid'] = user.id
		c['identifier'] = hashlib.sha1(str(user.id) + password).hexdigest()

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

		start_response('302 Found', [('Location', '/u'), ('Set-Cookie', c['uid'].OutputString()), ('Set-Cookie', c['identifier'].OutputString())])
		return []

	def register(self, environ, start_response, path):
		c = Cookie.SimpleCookie(environ['HTTP_COOKIE'] if 'HTTP_COOKIE' in environ else None)
		user = self.validate_cookie(c)
		form = cgi.FieldStorage(fp = environ['wsgi.input'], environ = environ)
		if environ['REQUEST_METHOD'] != 'POST' or not 'username' in form or not 'password' in form or not 'password2' in form:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return str(templates.register(searchList = {
				'user': user,
				'error': None,
			}))

		username = form.getvalue('username')
		password = form.getvalue('password')
		password2 = form.getvalue('password2')
		if password != password2:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return str(templates.register(searchList = {
				'user': user,
				'error': 'Passwords doesn\'t match',
			}))

		user = self.add_user(username, hashlib.sha1(password).hexdigest())
		if not user:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return str(templates.register(searchList = {
				'user': None,
				'error': 'Username already taken.',
			}))

		start_response('302 Found', [('Location', '/l')])
		return []

	def logout(self, environ, start_response, path):
		c = Cookie.SimpleCookie()
		expires = datetime.datetime.utcfromtimestamp(0).strftime('%a, %d-%b-%y %H:%M:%S GMT')
		c['uid'] = 0
		c['uid']['expires'] = expires
		c['identifier'] = ''
		c['identifier']['expires'] = expires
		start_response('302 Found', [('Set-Cookie', c['uid'].OutputString()), ('Set-Cookie', c['identifier'].OutputString()), ('Location', '/')])
		return []

	def static(self, environ, start_response, path):
		filename = path[1]
		if not filename in ('style.css',):
			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('static/' + filename, 'rb')

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

	def my_files(self, environ, start_response, path):
		c = Cookie.SimpleCookie(environ['HTTP_COOKIE'] if 'HTTP_COOKIE' in environ else None)
		user = self.validate_cookie(c)
		if user == None:
			start_response('200 OK', [('Content-Type', 'text/html')])
			return ['Not logged in.']
		files = self.get_files(user)
		start_response('200 OK', [('Content-Type', 'text/html')])
		return str(templates.my(searchList = {
			'user': user,
			'files': files,
		}))

	f = file
	u = upload
	l = login
	s = static
	h = help
	m = my_files
	o = logout
	r = register

	def __call__(self, environ, start_response):
		path = environ['PATH_INFO'].split('/')[1:]
		module = path[0]
		if len(module) and module in 'fulshmor':
			return getattr(self, module)(environ, start_response, path)
		else:
			start_response('302 Found', [('Location', '/u')])
			return []

if __name__ == '__main__':
	from wsgiref.simple_server import make_server, WSGIServer
	# enable IPv6
	WSGIServer.address_family |= 10
	http = make_server('', 8000, Application())
	http.serve_forever()