Aberrant/aberrant.py

#!/usr/bin/env python3
"""
The primary module for serving the Aberrant application.
"""
import json
import asyncio

import jinja2
import aiohttp_jinja2
from aiohttp import web, WSMsgType
from aiohttp_jinja2 import render_template

import config
import events
import rtorrent
import buckler_aiohttp


routes = web.RouteTableDef()

@routes.get('/', name='index')
async def index(request):
	"""The index page."""
	torrents = rtorrent.get_active()
	tracker_stats = rtorrent.get_stats()
	return render_template("index.html", request, locals())


@routes.get('/ws', name='ws')
async def websocket_handler(request):
	"""The websocket endpoint."""
	ws = web.WebSocketResponse(heartbeat=30)
	ws_ready = ws.can_prepare(request)
	if not ws_ready.ok:
		return web.Response(text="Cannot start websocket.")
	await ws.prepare(request)

	async for msg in ws:
		if msg.type == WSMsgType.TEXT:
			try:
				data = json.loads(msg.data)
			except json.JSONDecodeError:
				continue

			event = data.get('event')
			if not event or event not in events.events.keys():
				continue

			await events.events[event](ws, data.get('data'))
		else: # TODO: handle differnt message types properly
			break
	await ws.close()
	return ws


async def watcher_loop():
	try:
		while True:
			await rtorrent.update_torrents()
			await asyncio.sleep(10)
	except asyncio.CancelledError:
		return


async def start_background_tasks(app):
	rtorrent.init()
	app['watcher'] = asyncio.create_task(watcher_loop())


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


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

	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=5250)