From d1ec0358f69e292fdaf5138e38744311d2d725a8 Mon Sep 17 00:00:00 2001
From: Kraiem Taha Yassine <tahayk2@gmail.com>
Date: Fri, 10 Nov 2023 22:51:27 +0100
Subject: [PATCH] Api v1.15.0 (#1650)

* fix(chalice): fixed SSO support of refresh-token
---
 ee/api/chalicelib/core/users.py | 14 ++++++++++----
 ee/api/routers/saml.py          |  9 +++++++--
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/ee/api/chalicelib/core/users.py b/ee/api/chalicelib/core/users.py
index aaa0a5b18..9f7f57b35 100644
--- a/ee/api/chalicelib/core/users.py
+++ b/ee/api/chalicelib/core/users.py
@@ -854,10 +854,16 @@ def authenticate_sso(email, internal_id, exp=None):
         if r["serviceAccount"]:
             raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED,
                                 detail="service account is not authorized to login")
-        jwt_iat = TimeUTC.datetime_to_timestamp(refresh_jwt_iat_jti(r['userId']))
-        return authorizers.generate_jwt(r['userId'], r['tenantId'],
-                                        iat=jwt_iat, aud=f"front:{helper.get_stage_name()}",
-                                        exp=(exp + jwt_iat // 1000) if exp is not None else None)
+        jwt_iat, jwt_r_jti, jwt_r_iat = refresh_jwt_iat_jti(user_id=r['userId'])
+        return {
+            "jwt": authorizers.generate_jwt(user_id=r['userId'], tenant_id=r['tenantId'], iat=jwt_iat,
+                                            aud=f"front:{helper.get_stage_name()}"),
+            "refreshToken": authorizers.generate_jwt_refresh(user_id=r['userId'], tenant_id=r['tenantId'],
+                                                             iat=jwt_r_iat, aud=f"front:{helper.get_stage_name()}",
+                                                             jwt_jti=jwt_r_jti),
+            "refreshTokenMaxAge": config("JWT_REFRESH_EXPIRATION", cast=int),
+        }
+
     return None
 
 
diff --git a/ee/api/routers/saml.py b/ee/api/routers/saml.py
index 9451f4484..b7a5cc97c 100644
--- a/ee/api/routers/saml.py
+++ b/ee/api/routers/saml.py
@@ -107,9 +107,14 @@ async def process_sso_assertion(request: Request):
     jwt = users.authenticate_sso(email=email, internal_id=internal_id, exp=expiration)
     if jwt is None:
         return {"errors": ["null JWT"]}
-    return Response(
+    refresh_token = jwt["refreshToken"]
+    refresh_token_max_age = jwt["refreshTokenMaxAge"]
+    response = Response(
         status_code=status.HTTP_302_FOUND,
-        headers={'Location': SAML2_helper.get_landing_URL(jwt)})
+        headers={'Location': SAML2_helper.get_landing_URL(jwt["jwt"])})
+    response.set_cookie(key="refreshToken", value=refresh_token, path="/api/refresh",
+                        max_age=refresh_token_max_age, secure=True, httponly=True)
+    return response
 
 
 @public_app.post('/sso/saml2/acs/{tenantKey}', tags=["saml2"])