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

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()
	files = []
	for filefield in data.getall('files'):
		files.append(handle_filefield(filefield))

	files_insert = []
	for file in files:
		t = (request.cookies.get('userid'), file[0], file[1], 'DEFAUlT')
	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)