From 5b3949e2ab3a4ab5e642aeb965b6fbca2cabaea6 Mon Sep 17 00:00:00 2001
From: Blallo <blallo@autistici.org>
Date: Thu, 1 Aug 2019 00:29:28 +0200
Subject: [PATCH] Session-based firefox instance.

---
 api/__init__.py |  5 +++++
 api/app.py      | 35 ++++++++++++++++++++++++++---------
 api/rest.py     | 30 ++++++++++++++++++++++++++++--
 3 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/api/__init__.py b/api/__init__.py
index e69de29..42971ab 100644
--- a/api/__init__.py
+++ b/api/__init__.py
@@ -0,0 +1,5 @@
+# -*- encoding: utf-8 -*-
+import os
+
+# TODO: create config module and put this constant there.
+BASE_URI = os.environ.get("BOTZ_BASE_URI", "http://localhost")
diff --git a/api/app.py b/api/app.py
index 3a0faa8..1912600 100644
--- a/api/app.py
+++ b/api/app.py
@@ -5,28 +5,45 @@ The application entrypoint.
 """
 
 from aiohttp import web
-import click
+import base64
+from cryptography import fernet
 import logging
 import os
 import typing as T
 
+import click
+from aiohttp_session import setup, session_middleware
+from aiohttp_session.cookie_storage import EncryptedCookieStorage
+
 from bot_z.async_operator import AsyncOperator
 from api.rest import routes
 
 
-alog = logging.getLogger("api")
-alog.setLevel(os.environ.get("BOTZ_LOGLEVEL", logging.INFO))
-h = logging.StreamHandler()
-f = logging.Formatter("%(levelname)s [%(name)s] -> %(message)s")
-h.setFormatter(f)
-alog.addHandler(h)
-BASE_URI = os.environ.get("BOTZ_BASE_URI")
+def setup_log() -> logging.Logger:
+    alog = logging.getLogger("api")
+    alog.setLevel(os.environ.get("BOTZ_LOGLEVEL", logging.INFO))
+    h = logging.StreamHandler()
+    f = logging.Formatter("%(levelname)s [%(name)s] -> %(message)s")
+    h.setFormatter(f)
+    alog.addHandler(h)
+    return alog
+
+
+def init_secret() -> bytes:
+    fernet_key = fernet.Fernet.generate_key()
+    return base64.urlsafe_b64decode(fernet_key)
+
+
+def setup_session(app: web.Application):
+    secret = init_secret()
+    setup(app, EncryptedCookieStorage(secret))
 
 
 def run(address: T.Text, port: int) -> None:
     """Application entrypoint."""
+    alog = setup_log()
     app = web.Application(logger=alog)
-    app["async_operator"] = AsyncOperator(BASE_URI, "test")
+    setup_session(app)
     app.add_routes(routes)
     web.run_app(app, host=address, port=port)
 
diff --git a/api/rest.py b/api/rest.py
index 06c93ed..e4caff9 100644
--- a/api/rest.py
+++ b/api/rest.py
@@ -6,12 +6,31 @@ The REST endpoints.
 
 from aiohttp import web
 import logging
+import typing as T
 
+from aiohttp_session import get_session
+
+from bot_z.async_operator import AsyncOperator
 from api.async_bot import login, logout, checkin, checkout
+from api import BASE_URI
 
 
 alog = logging.getLogger("api")
 routes = web.RouteTableDef()
+OPERATORS = {}
+
+
+async def get_set_operator(
+    request: web.Request, user: T.Text, password: T.Text
+) -> AsyncOperator:
+    session = await get_session(request)
+    if "async_operator" in session:
+        op = OPERATORS[session["async_operator"]]
+    else:
+        op = AsyncOperator(BASE_URI, name=user)
+        session["async_operator"] = user
+        OPERATORS[user] = op
+    return op, session
 
 
 @routes.post("/login")
@@ -21,17 +40,24 @@ async def login_handler(request: web.Request) -> web.Response:
     password = data.get("password")
     if not user or not password:
         return web.json_response({"error": "Missing username or password"}, status=404)
-    op = request.app["async_operator"]
+    op, session = await get_set_operator(request, user, password)
     alog.debug("login - user: %s, password: %s", user, password)
     res = await login(op, user, password)
     alog.debug("login result: %s", res)
+    if not res:
+        session.invalidate()
+        alog.info("Login failed; session invalidated.")
     return web.json_response({"logged_in": res}, status=200)
 
 
 @routes.post("/logout")
 async def logout_handler(request: web.Request) -> web.Response:
     alog.debug("logout")
-    op = request.app["async_operator"]
+    session = await get_session(request)
+    op = OPERATORS.get(session.get("async_operator"))
+    if not op:
+        return web.json_response({"error": "No session"}, status=404)
     res = await logout(op)
     alog.debug("logout result: %s", res)
+    # FIX: assess if better to invalidate session and dump the browser session.
     return web.json_response({"logged_in": res}, status=200)