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