Juice/juice.py

#!/usr/bin/env python3
"""
A hub for controlling IOT devices.
"""
import os
import re
import copy
import json

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

import mqtt
import tools
import config
import events
import models
import buckler_aiohttp
from models import network

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

@routes.get('/', name='index')
async def index(request):
	"""The index page."""
	init_state = {}
	for device in network:
		if device.type != "RelayDevice":
			continue
		device_state = {}
		for sub_dev in device.sub_devices:
			device_state[sub_dev.id] = sub_dev.state
		init_state[device.id] = device_state
	d = {'network': network, 'init_state': init_state, 'request': request}
	return render_template('index.html', request, d)


@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)
	request.app['websockets'].append(ws)

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

		event = data.get('event')
		if not event or event not in events.events.keys():
			break
		await events.events[event](request, ws, data.get('data'))
	await ws.close()
	request.app['websockets'].remove(ws)
	return ws


async def send_json_all(self, d):
	"""Sends to all connected websockets."""
	for ws in self['websockets']:
		await ws.send_json(d)


async def start_background_tasks(app):
	app['mqtt'] = mqtt.init_client(app)


async def cleanup_background_tasks(app):
	app['mqtt'].loop_stop()


async def init_app():
	"""Initializes the application."""
	web.Application.send_json_all = send_json_all
	app = web.Application(middlewares=[buckler_aiohttp.buckler_session])
	aiohttp_jinja2.setup(
		app,
		trim_blocks=True,
		lstrip_blocks=True,
		undefined=jinja2.StrictUndefined,
		loader=jinja2.FileSystemLoader('templates'),
		filters={'html_color': tools.to_html_color},
	)
	app.on_startup.append(start_background_tasks)
	app.on_cleanup.append(cleanup_background_tasks)
	app['websockets'] = []

	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()
	aiohttp.web.run_app(app, host='0.0.0.0', port=5300)