From 46dc521733ce623d1d0d02f7eeba49bda1221b81 Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Thu, 30 Mar 2023 16:06:19 +0300
Subject: [PATCH] Island: Include refresh token into responses

Login and registration should respond with a token pair
---
 .../flask_resources/login.py                     | 11 +++++++----
 .../flask_resources/register.py                  | 11 +++++++----
 .../flask_resources/utils.py                     | 16 +++++++++++++++-
 3 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/monkey/monkey_island/cc/services/authentication_service/flask_resources/login.py b/monkey/monkey_island/cc/services/authentication_service/flask_resources/login.py
index 300de4e3119..3d5404e5e8e 100644
--- a/monkey/monkey_island/cc/services/authentication_service/flask_resources/login.py
+++ b/monkey/monkey_island/cc/services/authentication_service/flask_resources/login.py
@@ -9,7 +9,11 @@
 from monkey_island.cc.flask_utils import AbstractResource, responses
 
 from ..authentication_facade import AuthenticationFacade
-from .utils import get_username_password_from_request, include_auth_token
+from .utils import (
+    add_refresh_token_to_response,
+    get_username_password_from_request,
+    include_auth_token,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -38,9 +42,8 @@ def post(self):
         try:
             username, password = get_username_password_from_request(request)
             response: ResponseValue = login()
-            # TODO send these back
-            _tokens = self._authentication_facade.generate_refresh_token(current_user)
-            del _tokens
+            refresh_token = self._authentication_facade.generate_refresh_token(current_user)
+            response = add_refresh_token_to_response(response, refresh_token)
         except Exception:
             return responses.make_response_to_invalid_request()
 
diff --git a/monkey/monkey_island/cc/services/authentication_service/flask_resources/register.py b/monkey/monkey_island/cc/services/authentication_service/flask_resources/register.py
index dd82c182719..bd91935b0a0 100644
--- a/monkey/monkey_island/cc/services/authentication_service/flask_resources/register.py
+++ b/monkey/monkey_island/cc/services/authentication_service/flask_resources/register.py
@@ -9,7 +9,11 @@
 from monkey_island.cc.flask_utils import AbstractResource, responses
 
 from ..authentication_facade import AuthenticationFacade
-from .utils import get_username_password_from_request, include_auth_token
+from .utils import (
+    add_refresh_token_to_response,
+    get_username_password_from_request,
+    include_auth_token,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -37,9 +41,8 @@ def post(self):
                 }, HTTPStatus.CONFLICT
             username, password = get_username_password_from_request(request)
             response: ResponseValue = register()
-            # TODO send these back
-            _tokens = self._authentication_facade.generate_refresh_token(current_user)
-            del _tokens
+            refresh_token = self._authentication_facade.generate_refresh_token(current_user)
+            response = add_refresh_token_to_response(response, refresh_token)
         except Exception:
             return responses.make_response_to_invalid_request()
 
diff --git a/monkey/monkey_island/cc/services/authentication_service/flask_resources/utils.py b/monkey/monkey_island/cc/services/authentication_service/flask_resources/utils.py
index a0bb17228a1..ebff04b2ad8 100644
--- a/monkey/monkey_island/cc/services/authentication_service/flask_resources/utils.py
+++ b/monkey/monkey_island/cc/services/authentication_service/flask_resources/utils.py
@@ -1,10 +1,13 @@
 import json
+from copy import deepcopy
 from functools import wraps
 from typing import Tuple
 
-from flask import Request, request
+from flask import Request, Response, request
 from werkzeug.datastructures import ImmutableMultiDict
 
+from monkey_island.cc.services.authentication_service.refresh_token_manager import RefreshToken
+
 
 def get_username_password_from_request(_request: Request) -> Tuple[str, str]:
     """
@@ -35,3 +38,14 @@ def decorated_function(*args, **kwargs):
         return func(*args, **kwargs)
 
     return decorated_function
+
+
+def add_refresh_token_to_response(response: Response, refresh_token: RefreshToken) -> Response:
+    """
+    Adds a refresh token to the response
+    :param response: A Flask Response object
+    """
+    new_data = deepcopy(response.json)
+    new_data["response"]["user"]["refresh_token"] = refresh_token
+    response.data = json.dumps(new_data).encode()
+    return response