Merge branch '3137-new-token-pair-endpoint' into develop

Issue #3137
PR #3181

CHANGELOG.md

@@ -12,6 +12,8 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
 - `GET /api/agent-otp`. #3076
 - `POST /api/agent-otp-login` endpoint. #3076
 - A smarter brute-forcing strategy for SMB exploiter. #3039
+- `POST /api/refresh-authentication-token` endpoint that allows refreshing of
+  the access token. #3181
 
 ### Changed
 - Migrated the hard-coded SMB exploiter to a plugin. #2952

monkey/monkey_island/cc/app.py

@@ -44,7 +44,7 @@
 from monkey_island.cc.services.authentication_service.configure_flask_security import (
     configure_flask_security,
 )
-from monkey_island.cc.services.authentication_service.token import TokenGenerator, TokenValidator
+from monkey_island.cc.services.authentication_service.token import TokenGenerator, TokenParser
 from monkey_island.cc.services.representations import output_json
 
 HOME_FILE = "index.html"
@@ -162,16 +162,18 @@ def init_app(
 def _build_authentication_facade(container: DIContainer, security: Security):
     repository_encryptor = container.resolve(ILockableEncryptor)
     island_event_queue = container.resolve(IIslandEventQueue)
+
     token_generator = TokenGenerator(security)
     refresh_token_expiration = (
         security.app.config["SECURITY_TOKEN_MAX_AGE"]
         + security.app.config["SECURITY_REFRESH_TOKEN_TIMEDELTA"]
     )
-    refresh_token_validator = TokenValidator(security, refresh_token_expiration)
+    token_parser = TokenParser(security, refresh_token_expiration)
+
     return AuthenticationFacade(
         repository_encryptor,
         island_event_queue,
         security.datastore,
         token_generator,
-        refresh_token_validator,
+        token_parser,
     )

monkey/monkey_island/cc/services/authentication_service/authentication_facade.py

@@ -1,11 +1,14 @@
+from typing import Tuple
+
 from flask_security import UserDatastore
 
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.models import IslandMode
 from monkey_island.cc.server_utils.encryption import ILockableEncryptor
 from monkey_island.cc.services.authentication_service.token.token_generator import TokenGenerator
 
-from .token import Token, TokenValidator
+from . import AccountRole
+from .token import ParsedToken, Token, TokenParser
 from .user import User
 
 
@@ -20,28 +23,53 @@ def __init__(
         island_event_queue: IIslandEventQueue,
         user_datastore: UserDatastore,
         token_generator: TokenGenerator,
-        refresh_token_validator: TokenValidator,
+        token_parser: TokenParser,
     ):
         self._repository_encryptor = repository_encryptor
         self._island_event_queue = island_event_queue
         self._datastore = user_datastore
         self._token_generator = token_generator
-        self._refresh_token_validator = refresh_token_validator
+        self._token_parser = token_parser
 
     def needs_registration(self) -> bool:
         """
         Checks if a user is already registered on the Island
 
         :return: Whether registration is required on the Island
         """
-        return not User.objects.first()
+        island_api_user_role = self._datastore.find_or_create_role(
+            name=AccountRole.ISLAND_INTERFACE.name
+        )
+        return not self._datastore.find_user(roles=[island_api_user_role])
 
     def revoke_all_tokens_for_user(self, user: User):
         """
         Revokes all tokens for a specific user
         """
         self._datastore.set_uniquifier(user)
 
+    def generate_new_token_pair(self, refresh_token: Token) -> Tuple[Token, Token]:
+        """
+        Generates a new access token and refresh, given a valid refresh token
+
+        :param refresh_token: Refresh token
+        :raise TokenValidationError: If the refresh token is invalid or expired
+        :return: Tuple of the new access token and refresh token
+        """
+        parsed_refresh_token = self._token_parser.parse(refresh_token)
+        user = self._get_refresh_token_owner(parsed_refresh_token)
+
+        new_access_token = user.get_auth_token()
+        new_refresh_token = self._token_generator.generate_token(user.fs_uniquifier)
+
+        return new_access_token, new_refresh_token
+
+    def _get_refresh_token_owner(self, refresh_token: ParsedToken) -> User:
+        user = self._datastore.find_user(fs_uniquifier=refresh_token.user_uniquifier)
+        if not user:
+            raise Exception("Invalid refresh token")
+        return user
+
     def generate_refresh_token(self, user: User) -> Token:
         """
         Generates a refresh token for a specific user

monkey/monkey_island/cc/services/authentication_service/flask_resources/__init__.py

@@ -1,7 +1 @@
-from .register import Register
-from .registration_status import RegistrationStatus
-from .login import Login
-from .logout import Logout
 from .register_resources import register_resources
-from .agent_otp import AgentOTP
-from .agent_otp_login import AgentOTPLogin

monkey/monkey_island/cc/services/authentication_service/flask_resources/login.py

@@ -37,7 +37,6 @@ def post(self):
         returns an access token
 
         :return: Access token in the response body
-        :raises IncorrectCredentialsError: If credentials are invalid
         """
         try:
             username, password = get_username_password_from_request(request)

monkey/monkey_island/cc/services/authentication_service/flask_resources/refresh_authentication_token.py

@@ -0,0 +1,48 @@
+import logging
+from http import HTTPStatus
+
+from flask import make_response, request
+
+from monkey_island.cc.flask_utils import AbstractResource, responses
+from monkey_island.cc.services.authentication_service.token import TokenValidationError
+
+from ..authentication_facade import AuthenticationFacade
+from .utils import ACCESS_TOKEN_KEY_NAME, REFRESH_TOKEN_KEY_NAME
+
+logger = logging.getLogger(__name__)
+
+
+class RefreshAuthenticationToken(AbstractResource):
+    """
+    A resource for refreshing tokens
+    """
+
+    urls = ["/api/refresh-authentication-token"]
+
+    def __init__(self, authentication_facade: AuthenticationFacade):
+        self._authentication_facade = authentication_facade
+
+    def post(self):
+        """
+        Accepts a refresh token and returns a new token pair
+
+        :return: Response with new token pair or an invalid request response
+        """
+        try:
+            old_refresh_token = request.json[REFRESH_TOKEN_KEY_NAME]
+            access_token, refresh_token = self._authentication_facade.generate_new_token_pair(
+                old_refresh_token
+            )
+            response = {
+                "response": {
+                    "user": {
+                        ACCESS_TOKEN_KEY_NAME: access_token,
+                        REFRESH_TOKEN_KEY_NAME: refresh_token,
+                    }
+                }
+            }
+            return response, HTTPStatus.OK
+        except TokenValidationError:
+            return make_response({}, HTTPStatus.UNAUTHORIZED)
+        except Exception:
+            return responses.make_response_to_invalid_request()

monkey/monkey_island/cc/services/authentication_service/flask_resources/register_resources.py

@@ -5,6 +5,7 @@
 from .agent_otp_login import AgentOTPLogin
 from .login import Login
 from .logout import Logout
+from .refresh_authentication_token import RefreshAuthenticationToken
 from .register import Register
 from .registration_status import RegistrationStatus
 
@@ -18,3 +19,8 @@ def register_resources(api: flask_restful.Api, authentication_facade: Authentica
     api.add_resource(Logout, *Logout.urls, resource_class_args=(authentication_facade,))
     api.add_resource(AgentOTP, *AgentOTP.urls)
     api.add_resource(AgentOTPLogin, *AgentOTPLogin.urls)
+    api.add_resource(
+        RefreshAuthenticationToken,
+        *RefreshAuthenticationToken.urls,
+        resource_class_args=(authentication_facade,),
+    )

monkey/monkey_island/cc/services/authentication_service/flask_resources/utils.py

@@ -8,6 +8,9 @@
 
 from monkey_island.cc.services.authentication_service.token import Token
 
+REFRESH_TOKEN_KEY_NAME = "refresh_token"
+ACCESS_TOKEN_KEY_NAME = "authentication_token"
+
 
 def get_username_password_from_request(_request: Request) -> Tuple[str, str]:
     """
@@ -49,6 +52,6 @@ def add_refresh_token_to_response(response: Response, refresh_token: Token) -> R
     :return: A Flask Response object
     """
     new_data = deepcopy(response.json)
-    new_data["response"]["user"]["refresh_token"] = refresh_token
+    new_data["response"]["user"][REFRESH_TOKEN_KEY_NAME] = refresh_token
     response.data = json.dumps(new_data).encode()
     return response

monkey/monkey_island/cc/services/authentication_service/token/__init__.py

@@ -1,3 +1,3 @@
 from .token_generator import TokenGenerator
-from .token_validator import TokenValidator
+from .token_parser import TokenParser, ParsedToken, TokenValidationError
 from .types import Token

monkey/monkey_island/cc/services/authentication_service/token/token_parser.py

@@ -0,0 +1,67 @@
+from flask_security import Security
+from itsdangerous import BadSignature, Serializer, SignatureExpired
+from pydantic import PrivateAttr
+
+from common.base_models import InfectionMonkeyBaseModel
+
+from .types import Token
+
+
+class TokenValidationError(Exception):
+    """Raise when an invalid token is encountered"""
+
+
+class InvalidTokenSignatureError(TokenValidationError):
+    """Raise when a token's signature is invalid"""
+
+
+class ExpiredTokenError(TokenValidationError):
+    """Raise when a token has expired"""
+
+
+class ParsedToken(InfectionMonkeyBaseModel):
+    raw_token: Token
+    user_uniquifier: str
+    expiration_time: int
+    _token_serializer: Serializer = PrivateAttr()
+
+    def __init__(self, token_serializer: Serializer, *, raw_token: Token, expiration_time: int):
+        self._token_serializer = token_serializer
+
+        user_uniquifier = self._token_serializer.loads(raw_token, max_age=expiration_time)
+        super().__init__(
+            raw_token=raw_token, user_uniquifier=user_uniquifier, expiration_time=expiration_time
+        )
+
+    def is_expired(self) -> bool:
+        try:
+            self._token_serializer.loads(self.raw_token, max_age=self.expiration_time)
+            return False
+        except SignatureExpired:
+            return True
+
+
+class TokenParser:
+    def __init__(self, security: Security, token_expiration: int):
+        self._token_serializer = security.remember_token_serializer
+        self._token_expiration = token_expiration  # in seconds
+
+    def parse(self, token: Token) -> ParsedToken:
+        """
+        Parses a token and returns a data structure with its components
+
+        :param token: The token to parse
+        :return: The parsed token
+        :raises TokenValidationError: If the token could not be parsed
+        """
+        try:
+            return ParsedToken(
+                token_serializer=self._token_serializer,
+                raw_token=token,
+                expiration_time=self._token_expiration,
+            )
+        except SignatureExpired:
+            # NOTE: SignatureExpired is a subclass of BadSignature; this clause must come first.
+            raise ExpiredTokenError("Token has expired")
+        except BadSignature:
+            raise InvalidTokenSignatureError("Invalid token signature")

monkey/tests/unit_tests/monkey_island/cc/services/authentication_service/conftest.py

@@ -3,17 +3,10 @@
 import pytest
 from tests.unit_tests.monkey_island.conftest import init_mock_security_app
 
+from monkey_island.cc.services.authentication_service import register_resources
 from monkey_island.cc.services.authentication_service.authentication_facade import (
     AuthenticationFacade,
 )
-from monkey_island.cc.services.authentication_service.flask_resources import (
-    AgentOTP,
-    AgentOTPLogin,
-    Login,
-    Logout,
-    Register,
-    RegistrationStatus,
-)
 
 REFRESH_TOKEN = "refresh_token"
 
@@ -37,15 +30,7 @@ def inner():
 
 def get_mock_auth_app(authentication_facade: AuthenticationFacade):
     app, api = init_mock_security_app()
-    api.add_resource(Register, *Register.urls, resource_class_args=(authentication_facade,))
-    api.add_resource(Login, *Login.urls, resource_class_args=(authentication_facade,))
-    api.add_resource(Logout, *Logout.urls, resource_class_args=(authentication_facade,))
-    api.add_resource(
-        RegistrationStatus, *RegistrationStatus.urls, resource_class_args=(authentication_facade,)
-    )
-    api.add_resource(AgentOTP, *AgentOTP.urls)
-    api.add_resource(AgentOTPLogin, *AgentOTPLogin.urls)
-
+    register_resources(api, authentication_facade)
     return app