From e658eb9b2aa415480b7b3c26176cc376b3ae2d10 Mon Sep 17 00:00:00 2001
From: odysseusmax <royschristy@gmail.com>
Date: Thu, 10 Jun 2021 09:24:15 +0530
Subject: [PATCH] add redirection memory

---
 app/templates/login.html |  1 +
 app/views/login_view.py  | 14 ++++++++------
 app/views/middlewhere.py |  3 +++
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/app/templates/login.html b/app/templates/login.html
index 21e606a..685d576 100644
--- a/app/templates/login.html
+++ b/app/templates/login.html
@@ -16,6 +16,7 @@
         {% endif %}
         <form class="mt-8 space-y-6" action="/login" method="POST">
             <input type="hidden" name="remember" value="true">
+            <input type="hidden" name="redirect_to" value="{{redirect_to}}">
             <div class="rounded-md shadow-sm -space-y-px">
                 <div>
                     <label for="email-address" class="sr-only">Username</label>
diff --git a/app/views/login_view.py b/app/views/login_view.py
index dbd9983..1229259 100644
--- a/app/views/login_view.py
+++ b/app/views/login_view.py
@@ -13,25 +13,27 @@ class LoginView:
 
     async def login_post(self, req):
         post_data = await req.post()
+        redirect_to = post_data.get("redirect_to") or "/"
         location = req.app.router["login_page"].url_for()
+        if redirect_to != "/":
+            location = location.update_query({"redirect_to": redirect_to})
+
         if "username" not in post_data:
-            loc = location.with_query({"error": "Username missing"})
+            loc = location.update_query({"error": "Username missing"})
             raise web.HTTPFound(location=loc)
 
         if "password" not in post_data:
-            loc = location.with_query({"error": "Password missing"})
+            loc = location.update_query({"error": "Password missing"})
             raise web.HTTPFound(location=loc)
 
         authenticated = (post_data["username"] == req.app["username"]) and (
             post_data["password"] == req.app["password"]
         )
         if not authenticated:
-            loc = location.with_query({"error": "Wrong Username or Passowrd"})
+            loc = location.update_query({"error": "Wrong Username or Passowrd"})
             raise web.HTTPFound(location=loc)
 
-        resp = web.Response(
-            status=302, headers={"Location": str(req.app.router["home"].url_for())}
-        )
+        resp = web.Response(status=302, headers={"Location": redirect_to})
         now = time.time()
         resp.set_cookie(
             name="_tgindex_session",
diff --git a/app/views/middlewhere.py b/app/views/middlewhere.py
index 4e4d524..57a4dc0 100644
--- a/app/views/middlewhere.py
+++ b/app/views/middlewhere.py
@@ -18,6 +18,9 @@ def middleware_factory():
         ]:
             cookies = request.cookies
             url = request.app.router["login_page"].url_for()
+            if str(request.rel_url) != "/":
+                url = url.with_query(redirect_to=str(request.rel_url))
+
             if any(x not in cookies for x in ("_tgindex_session", "_tgindex_secret")):
                 raise HTTPFound(url)