Saddle/saddle.py

#!/usr/bin/env python3
"""
A file hosting service similar to Pomf and Uguu but without the public nature.
"""
import os
import string
import random
import datetime

from aiohttp import web
import jinja2
import aiohttp_jinja2
from aiohttp_jinja2 import render_template
import asyncpg
import uvloop

import config
import buckler_aiohttp

uvloop.install()
routes = web.RouteTableDef()


@routes.get('/', name='index')
@routes.post('/', name='index')
async def index(request):
	"""The index page."""
	if request.method == 'GET':
		return render_template("index.html", request, locals())
	data = await request.post()
	rand_name = bool(data.get('rand_name'))
	files = []
	for filefield in data.getall('files'):
		files.append(handle_filefield(filefield, rand_name=rand_name))
	if data.get('delete_this'):
		delete_num = data.get('delete_num', '')
		delete_type = data.get('delete_type', '')

		try:
			delete_num = int(delete_num)
			assert delete_num >= 1 and delete_num <= 59
			assert delete_type in ['minutes', 'hours', 'days', 'weeks']
		except (ValueError, AssertionError):
			return 'ur ghey' # TODO: return error
		delta = datetime.timedelta(**{delete_type: delete_num})
		expiration_date = datetime.datetime.now() + delta
	else:
		expiration_date = None

	files_insert = []
	for file in files:
		t = (int(request.cookies.get('userid')), file[0], file[1], expiration_date)
		files_insert.append(t)
	async with request.app['pool'] as conn:
		await conn.executemany(
			"INSERT INTO upload (user_id, id, filename, expiration_date) "
			"VALUES ($1, $2, $3, $4)",
			files_insert)
	urls = [config.upload_url + f[1] for f in files]
	return render_template("result.html", request, locals())


def handle_filefield(filefield, rand_name=True):
	"""Handles a posted file."""
	filename = safe_filename(filefield.filename)
	if not filename:
		rand_name = True
	prefix = get_rand_chars()
	if rand_name:
		filename = prefix + os.path.splitext(filename)[1]
	else:
		filename = prefix + '_' + filename

	with open(os.path.join(config.upload_dir, filename), 'wb') as file:
		file.write(filefield.file.read())

	t = (prefix, filename)
	return t


def safe_filename(filename=''):
	"""Sanitizes the given filename."""
	safe_char = string.ascii_letters + string.digits + '._ '
	filename = ''.join([c for c in filename if c in safe_char])
	filename = filename.strip('._ ')
	return filename


def get_rand_chars(n=8):
	"""Returns `n` number of random ASCII characters."""
	chars = []
	for _ in range(n):
		char = random.choice(string.ascii_letters + string.digits)
		chars.append(char)
	return "".join(chars)


async def init_app():
	"""Initializes the application."""
	#app = web.Application(middlewares=[buckler_aiohttp.buckler_session])
	app = web.Application()
	aiohttp_jinja2.setup(app, loader=jinja2.FileSystemLoader('templates'))
	app['pool'] = await asyncpg.create_pool(**config.db)

	async with app['pool'].acquire() as conn:
		with open('saddle.sql', 'r') as file:
			await conn.execute(file.read())

	app.router.add_routes(routes)

	app_wrap = web.Application()
	app_wrap.add_subapp(config.url_prefix, app)
	return app_wrap


if __name__ == "__main__":
	app = init_app()
	web.run_app(app, host='0.0.0.0', port=5000)