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

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 auth
import tools
import config

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

@routes.get(config.url_prefix + '/', name='index')
@auth.auth_required
async def index(request):
	"""The index page."""
	async with request.app['pool'].acquire() as conn:
		avail_sites = await conn.fetch(
			"SELECT app_info.name, app_info.url "
			"FROM app_info LEFT JOIN app_user "
			"ON (app_info.id = app_user.app_id) "
			"WHERE app_user.user_id = $1",
			request['session']['id'])
		if request['session']['admin']:
			apps = await conn.fetch(
				"SELECT id, name FROM app_info")
			user_perms = await conn.fetch(
				"SELECT user_info.id, user_info.username, app_user.app_id "
				"FROM user_info LEFT JOIN app_user "
				"ON (user_info.id = app_user.user_id)")
		fido2_keys = await conn.fetch(
			"SELECT * FROM user_credential WHERE user_id = $1",
			request['session']['id'])
		active_sessions = await conn.fetch(
			"SELECT id, ip_address FROM user_session "
			"WHERE user_id = $1",
			request['session']['id'])

	if request['session']['admin']:
		users = defaultdict(lambda: {app['name']: False for app in apps})
		for user_perm in user_perms:
			index = tools.find_dict(apps, 'id', user_perm['app_id'])
			if index != -1:
				users[user_perm['username']][apps[index]['name']] = True
		users_json = json.dumps(users)
	return render_template("index.html", request, locals())


@routes.get(config.url_prefix + '/change_password', name='change_password')
@auth.auth_required
async def change_password(request):
	"""Allows a user to change their password."""
	pass


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

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

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

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

	if has_cred['exists'] and user_info['passwordless']:
		url_prefix = config.url_prefix
		resp = render_template("login_key.html", request, locals())
		resp.set_cookie('userid', user_info['id'])
		return resp

	if argon2.verify(password, user_info['password_hash']):
		if has_cred['exists']:
			url_prefix = config.url_prefix
			resp = render_template("login_key.html", request, locals())
			resp.set_cookie('userid', user_info['id'])
			return resp

		url = request.cookies.get('redirect')
		if not url:
			url = request.app.router['index'].url_for()
		resp = web.HTTPFound(location=url)
		resp.set_cookie('redirect', '', max_age=0)
		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:
		login_failed = True
		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."""
	confirm_token = request.query.get('confirm')
	if confirm_token:
		async with request.app['pool'].acquire() as conn:
			confirm = await conn.fetchrow(
				"SELECT * FROM email_confirmation WHERE token = $1",
				confirm_token)
		if confirm:
			async with request.app['pool'].acquire() as conn:
				await conn.execute(
					"UPDATE user_info SET active = TRUE WHERE id = $1",
					confirm['user_id'])
				await conn.execute(
					"DELETE FROM email_confirmation WHERE user_id = $1",
					confirm['user_id'])
			message = "Success!" # TODO: more informative message
		else:
			message = "Invalid confirmation token."
		return render_template("register_result.html", request, locals())

	invite_token = request.query.get('invite')
	if invite_token:
		async with request.app['pool'].acquire() as conn:
			invite = await conn.fetchrow(
				"SELECT * FROM invite WHERE token = $1",
				invite_token)
		if invite:
			if request.method == 'GET':
				errors = {}
				return render_template("register.html", request, locals())

			form = await request.post()
			errors = await tools.validate_register(request, form)
			if any(errors.values()):
				return render_template("register.html", request, locals())

			username = form.get('username')
			email = form.get('email')
			password = form.get('password')
			pw_hash = argon2.hash(password)
			async with request.app['pool'].acquire() as conn:
				user = await conn.fetchrow(
					"INSERT INTO user_info (username, email, password_hash) "
					"VALUES ($1, $2, $3) RETURNING id",
					username, email, pw_hash)
				await conn.execute(
					"DELETE FROM invite WHERE token = $1",
					invite_token)
			await tools.send_confirmation(request, user['id'], email)
			message = "An email has been sent." # TODO: more better
		else:
			message = "Invalid invitation token."
		return render_template("register_result.html", request, locals())

	message = "Secret club only."
	return render_template("register_result.html", request, locals())


@routes.get(config.url_prefix + '/add_key', name='add_key')
@auth.auth_required
async def add_key(request):
	"""Add a new security key."""
	url_prefix = config.url_prefix
	return render_template("register_key.html", request, locals())


@routes.get(config.url_prefix + '/get_session', name='get_session')
async def get_session(request):
	"""Returns a user's application session."""
	# TODO: only allow LAN IPs
	app_id = request.query.get('app_id')
	app_key = request.query.get('app_key')
	user_id = request.query.get('userid')
	user_sid = request.query.get('session')

	try:
		app_id = int(app_id)
		user_id = int(user_id)
		assert all((app_id, app_key, user_id, user_sid))
	except (ValueError, AssertionError):
		return web.json_response({'error': "Invalid credentials."})

	conn = await request.app['pool'].acquire()
	app = await conn.fetchrow("SELECT * FROM app_info WHERE id = $1", app_id)
	if app:
		if argon2.verify(app_key, app['key_hash']):
			sessions = await conn.fetch(
				"SELECT * FROM user_session "
				"WHERE user_id = $1 AND expires > NOW()",
				user_id)
			session = [s for s in sessions if s.get('id') == user_sid]
			if session:
				session = session[0]
				data = await conn.fetchrow(
					"SELECT user_info.username, app_user.session_data "
					"FROM user_info LEFT JOIN app_user "
					"ON (user_info.id = app_user.user_id) "
					"WHERE user_info.id = $1 AND app_user.app_id = $2",
					session['user_id'], app['id'])
				
				tz = session['last_used'].tzinfo
				delta = datetime.now(tz) - session['last_used']
				if delta.seconds > 600:
					await conn.execute(
						"UPDATE user_session SET last_used = NOW(), "
						"expires = NOW() + INTERVAL '30 DAYS' "
						"WHERE user_id = $1",
						user_id)
				await conn.close()

				data_meta = dict(data)
				data_meta['last_used'] = session['last_used'].isoformat()
				data_meta['user_sid'] = user_sid
				data_meta['user_id'] = user_id
				session_data = data_meta.pop('session_data')
				data = {
					'meta': data_meta,
					'session_data': json.loads(session_data)
				}

				return web.json_response(data)
			else:
				error = {'error': "User ID or Session ID invalid."}
		else:
			error = {'error': "App ID or Key invalid."}
	else:
		error = {'error': "App ID or Key invalid."}
	await conn.close()
	return web.json_response(error)


@routes.post(config.url_prefix + '/set_session', name='set_session')
async def set_session(request):
	"""Allows an application to set a user app session."""
	# TODO: only allow LAN IPs
	app_id = request.query.get('app_id')
	app_key = request.query.get('app_key')
	user_id = request.query.get('userid')
	user_sid = request.query.get('session')

	try:
		app_id = int(app_id)
		user_id = int(user_id)
		assert all((app_id, app_key, user_id, user_sid))
	except (ValueError, AssertionError):
		return web.json_response({'error': "Invalid credentials."})

	conn = await request.app['pool'].acquire()
	app = await conn.fetchrow("SELECT * FROM app_info WHERE id = $1", app_id)
	if app:
		if argon2.verify(app_key, app['key_hash']):
			session = await conn.fetchrow(
				"SELECT * FROM user_session "
				"WHERE user_id = $1 AND id = $2 AND expires > NOW()",
				user_id, user_sid)
			if session:
				session_data = await request.text()
				# TODO: error handling, verify json
				await conn.execute(
					"UPDATE app_user SET session_data = $1 "
					"WHERE user_id = $2 AND app_id = $3",
					session_data, user_id, app_id)
				await conn.close()
				return web.json_response({'success': True})
			else:
				error = {'error': "User ID or Session ID invalid."}
		else:
			error = {'error': "App ID or Key invalid."}
	else:
		error = {'error': "App ID or Key invalid."}
	await conn.close()
	return web.json_response(error)


@routes.post(config.url_prefix + '/delete_key', name='delete_key')
@auth.auth_required
async def delete_key(request):
	"""Allows a user to delete a security key."""
	data = await request.post()
	async with request.app['pool'].acquire() as conn:
		for key, value in data.items():
			key_id = key.replace('fido-', '')
			if not key_id:
				continue
			try:
				key_id = int(key_id)
			except ValueError:
				continue
			if value != 'on':
				continue
			await conn.execute(
				"DELETE FROM user_credential "
				"WHERE id = $1 AND user_id = $2",
				key_id, request['session']['id'])
	index_url = request.app.router['index'].url_for()
	raise web.HTTPFound(location=index_url)


@routes.post(config.url_prefix + '/delete_session', name='delete_session')
@auth.auth_required
async def delete_session(request):
	"""Allows a user to delete a session ."""
	data = await request.post()
	async with request.app['pool'].acquire() as conn:
		for key, value in data.items():
			session_id = key.replace('session-', '', 1)
			if not session_id:
				continue
			if value != 'on':
				continue
			await conn.execute(
				"DELETE FROM user_session "
				"WHERE id = $1 AND user_id = $2",
				session_id, request['session']['id'])
	index_url = request.app.router['index'].url_for()
	raise web.HTTPFound(location=index_url)


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)
	app.router.add_routes(auth.routes)
	return app

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