Buckler/buckler.py

#!/usr/bin/env python3
"""
A security shield for protecting a number of small web applications.
"""
import secrets
import functools
from datetime import datetime

from aiohttp import web
import jinja2
import aiohttp_jinja2
from aiohttp_jinja2 import render_template
from passlib.hash import argon2
import asyncpg
import uvloop

import config

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

def auth_required(func):
	"""
	Wrapper for views with should be protected by authtication.
	"""
	@functools.wraps(func)
	async def wrapper(request, *args, **kwargs):
		login_url = request.app.router['login'].url_for()
		sid = request.cookies.get('session')
		try:
			userid = int(request.cookies.get('userid'))
		except (ValueError, TypeError):
			userid = None
		if not sid or not userid:
			raise web.HTTPFound(location=login_url)
		async with request.app['pool'].acquire() as conn:
			sessions = await conn.fetch(
				"SELECT * FROM user_session "
				"WHERE user_id = $1 AND expiry > NOW() ",
				userid)
		session = [s for s in sessions if s.get('id') == sid]
		if session:
			session = session[0]
			tz = session['last_used'].tzinfo
			delta = datetime.now(tz) - session['last_used']
			if delta.seconds > 600:
				async with request.app['pool'].acquire() as conn:
					await conn.execute(
						"UPDATE user_session SET last_used = NOW(),"
						"expiry = NOW() + INTERVAL '30 DAYS'"
						"WHERE user_id = $1",
						userid)
			resp = await func(request, *args, **kwargs)
			if delta.seconds > 600:
				resp.set_cookie(
					'userid',
					userid,
					max_age=30*24*60*60,
					secure=True,
					httponly=True)
				resp.set_cookie(
					'session',
					sid,
					max_age=30*24*60*60,
					secure=True,
					httponly=True)
			return resp
		else:
			raise web.HTTPFound(location=login_url)
	return wrapper


@routes.get(config.url_prefix + '/', name='index')
@auth_required
async def index(request):
	"""The index page."""
	return render_template("index.html", request, locals())


@routes.get(config.url_prefix + '/login', name='login')
@routes.post(config.url_prefix + '/login', name='login')
async def login(request):
	"""Handle login."""
	if request.method == 'GET':
		return render_template("login.html", request, locals())

	data = await request.post()
	username = data.get('username')
	password = data.get('password')

	async with request.app['pool'].acquire() as conn:
		user_info = await conn.fetchrow(
			"SELECT * FROM user_info WHERE username = $1",
			username)

	if not user_info:
		return render_template("login.html", request, locals())

	if argon2.verify(password, user_info['password_hash']):
		index_url = request.app.router['index'].url_for()
		resp = web.HTTPFound(location=index_url)
		resp.set_cookie(
			'userid',
			user_info['id'],
			max_age=30*24*60*60,
			secure=True,
			httponly=True)
		sid = secrets.token_urlsafe(64)
		ip_address = request.headers['X-Real-IP']
		async with request.app['pool'].acquire() as conn:
			await conn.execute(
				"INSERT INTO user_session (user_id, id, ip_address)"
				"VALUES ($1, $2, $3)",
				user_info['id'],
				sid,
				ip_address)
		resp.set_cookie(
			'session',sid,
			max_age=30*24*60*60,
			secure=True,
			httponly=True)
		raise resp
	else:
		return render_template("login.html", request, locals())


@routes.get(config.url_prefix + '/register', name='register')
@routes.post(config.url_prefix + '/register', name='register')
async def register(request):
	"""Register new accounts."""
	if request.method == 'POST':
		data = await request.post()
		username = data.get('username')
		email = data.get('email')
		password = data.get('password')
		password_verify = data.get('password_verify')

		errors = {}
		if password != password_verify:
			errors[''] = ''

		pw_hash = argon2.hash(password)
		async with request.app['pool'].acquire() as conn:
			await conn.execute(
				"INSERT INTO user_info (username, email, password_hash) "
				"VALUES ($1, $2, $3)",
				username, email, pw_hash)
		login_url = request.app.router['login'].url_for()
		raise web.HTTPFound(location=login_url)
	return render_template("register.html", request, locals())


async def init_app():
	"""Initializes the application."""
	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('buckler.sql', 'r') as file:
			await conn.execute(file.read())

	app.router.add_routes(routes)
	return app

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