Saddle/saddle.py

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

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

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'))
	response_type = data.get('response_type', 'plain')

	files = []
	for filefield in data.getall('files'):
		if not filefield:
			continue
		files.append(handle_filefield(filefield, rand_name=rand_name))
	if data.get('url'):
		files.append(handle_url(data.get('url'), 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'].acquire() 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]
	if response_type == 'html':
		return render_template("result.html", request, locals())
	elif response_type == 'plain':
		return web.Response(body='\n'.join(urls))
	elif response_type == 'json':
		return web.json_response(urls)


@routes.get('/gallery/{user_id}', name='gallery')
async def gallery(request):
	"""A user's gallery page."""
	try:
		user_id = int(request.match_info['user_id'])
	except ValueError:
		raise web.HTTPNotFound

	async with request.app['pool'].acquire() as conn:
		uploads = await conn.fetch(
			"SELECT * FROM upload WHERE user_id = $1",
			user_id)
	upload_url = config.upload_url
	return render_template("gallery.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())

	tup = (prefix, filename)
	return tup


def handle_url(url, rand_name=True):
	"""Handles a posted URL."""
	try:
		filename, data = download_file(url)
	except ValueError:
		return None

	filename = safe_filename(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(data)

	tup = (prefix, filename)
	return tup


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)


def download_file(url, timeout=10, max_file_size=config.client_max_size):
	"""
	Downloads the file at the given url while observing file size and
	timeout limitations.
	"""
	requests_kwargs = {
		'stream': True,
		'headers': {'User-Agent': "Steelbea.me LTD needs YOUR files."},
		'timeout': timeout,
		'verify': True
	}
	temp = b''
	with requests.get(url, **requests_kwargs) as r:
		size = 0
		start_time = time.time()
		for chunk in r.iter_content(102400):
			if time.time() - start_time > timeout:
				raise ValueError('timeout reached')
			if len(temp) > max_file_size:
				raise ValueError('response too large')
			temp += chunk

		if r.headers.get('Content-Disposition'):
			fname = re.search(r'filename="(.+)"',
					r.headers['content-disposition'])
		else:
			fname = os.path.basename(url)
	return (fname, temp)


async def cleaner(app):
	"""Removes files marked for deletion."""
	async with app['pool'].acquire() as conn:
		expired = await conn.fetch(
			"SELECT * FROM upload WHERE expiration_date < NOW()")
		if not expired:
			return
		for record in expired:
			os.remove(os.path.join(config.upload_dir, record['filename']))
		await conn.executemany(
			"DELETE FROM upload WHERE id = $1",
			[(record['id'],) for record in expired])


async def cleaner_loop(app):
	"""Loops cleaner() continuously until shutdown."""
	try:
		while True:
			await cleaner(app)
			await asyncio.sleep(60)
	except asyncio.CancelledError:
		return


async def start_background_tasks(app):
	app['cleaner'] = asyncio.create_task(cleaner_loop(app))


async def cleanup_background_tasks(app):
	app['cleaner'].cancel()
	await app['cleaner']


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

	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(client_max_size=config.client_max_size)
	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)