From 2816b84befbee572c432480bae8cee8cc1b8501b Mon Sep 17 00:00:00 2001
From: mishaschwartz <4380924+mishaschwartz@users.noreply.github.com>
Date: Tue, 25 Jul 2023 11:12:55 -0400
Subject: [PATCH] login/logout of magpie as well

---
 .../jupyterhub_magpie_authenticator.py        | 41 ++++++++++++++++---
 1 file changed, 35 insertions(+), 6 deletions(-)

diff --git a/jupyterhub_magpie_authenticator/jupyterhub_magpie_authenticator.py b/jupyterhub_magpie_authenticator/jupyterhub_magpie_authenticator.py
index d9921b0..3d6bb23 100644
--- a/jupyterhub_magpie_authenticator/jupyterhub_magpie_authenticator.py
+++ b/jupyterhub_magpie_authenticator/jupyterhub_magpie_authenticator.py
@@ -1,15 +1,29 @@
 from traitlets import Unicode
 from jupyterhub.auth import Authenticator
-from tornado import gen
+from jupyterhub.handlers.login import LogoutHandler
 import requests
 
 
+class MagpieLogoutHandler(LogoutHandler):
+    """
+    Logout Handler that also logs the user out of magpie when logging out of jupyterhub.
+    """
+
+    async def handle_logout(self):
+        cookies = {key: morsel.coded_value for key, morsel in self.request.cookies.items()}
+        signout_url = self.authenticator.magpie_url.rstrip("/") + "/signout"
+        response = requests.get(signout_url, cookies=cookies, headers={"Host": self.authenticator.public_fqdn})
+        if response.ok and 'Set-Cookie' in response.headers:
+            self.set_header("Set-Cookie", response.headers["Set-Cookie"])
+
+
 class MagpieAuthenticator(Authenticator):
     """Authenticate to JupyterHub using Magpie.
 
     To use this authenticator, set the following parameters in the `jupyterhub_config.py` file:
      - c.JupyterHub.authenticator_class = 'jupyterhub_magpie_authenticator.MagpieAuthenticator'
-     - c.MagpieAuthenticator.magpie_url = "https://www.example.com/magpie"
+     - c.MagpieAuthenticator.magpie_url = "magpie:2000" # url where magpie is running (does not need to be public)
+     - c.MagpieAuthenticator.public_fqdn = "www.example.com"  # fqdn of server where magpie is running
     """
     default_provider = "ziggurat"
     magpie_url = Unicode(
@@ -17,17 +31,32 @@ class MagpieAuthenticator(Authenticator):
         config=True,
         help="Magpie endpoint to signin to"
     )
+    public_fqdn = Unicode(
+        config=True,
+        help="Public fully qualified domain name. Used to set the magpie login cookie."
+    )
+
+    def get_handlers(self, app):
+        return [
+            ('/logout', MagpieLogoutHandler)
+        ]
 
-    @gen.coroutine
-    def authenticate(self, handler, data):
+    async def authenticate(self, handler, data):
         signin_url = self.magpie_url.rstrip('/') + '/signin'
-        
+
         post_data = {
             "user_name": data["username"],
             "password": data["password"],
             "provider_name": self.default_provider,
         }
         response = requests.post(signin_url, data=post_data)
-        
+
         if response.ok:
+            for cookie in response.cookies:
+                handler.set_cookie(name=cookie.name,
+                                   value=cookie.value,
+                                   domain=self.public_fqdn,
+                                   expires=cookie.expires,
+                                   path=cookie.path,
+                                   secure=cookie.secure)
             return data['username']