Saddle/saddle.py
2019-11-06 13:09:07 -05:00

239 lines
6.1 KiB
Python

#!/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)