Merge pull request #23 from stax-labs/feat/auth_errors

feat(auth) convert boto error's into staxapp errors

Author: ZXYmania

staxapp/auth.py

@@ -8,9 +8,11 @@
 from aws_requests_auth.aws_auth import AWSRequestsAuth
 from botocore import UNSIGNED
 from botocore.client import Config as BotoConfig
+from botocore.exceptions import ClientError
 from warrant import AWSSRP, Cognito
 
 from staxapp.config import Config as JumaConfig
+from staxapp.exceptions import InvalidCredentialsException
 
 
 class StaxAuth:
@@ -23,6 +25,15 @@ def __init__(self, config_branch):
         self.aws_region = config.get(config_branch).get("region")
 
     def requests_auth(self, username, password):
+        if username is None:
+            raise InvalidCredentialsException(
+                "Please provide an Access Key to your config"
+            )
+        if password is None:
+            raise InvalidCredentialsException(
+                "Please provide a Secret Key to your config"
+            )
+
         id_token = self.id_token_from_cognito(username, password)
         id_creds = self.sts_from_cognito_identity_pool(id_token)
         auth = self.sigv4_signed_auth_headers(id_creds)
@@ -32,26 +43,39 @@ def requests_auth(self, username, password):
 
         return JumaConfig.auth
 
-    def id_token_from_cognito(self, username=None, password=None):
+    def id_token_from_cognito(self, username=None, password=None, client=None):
         token = None
         if username and password:
-            client = boto3.client(
-                "cognito-idp",
-                region_name=self.aws_region,
-                config=BotoConfig(signature_version=UNSIGNED),
-            )
+            if not client:
+                client = boto3.client(
+                    "cognito-idp",
+                    region_name=self.aws_region,
+                    config=BotoConfig(signature_version=UNSIGNED),
+                )
             aws = AWSSRP(
                 username=username,
                 password=password,
                 pool_id=self.user_pool,
                 client_id=self.client_id,
                 client=client,
             )
-            tokens = aws.authenticate_user()
-            # logging.debug(f"TOKEN: {tokens}")
+            try:
+                tokens = aws.authenticate_user()
+            except ClientError as e:
+                if e.response["Error"]["Code"] == "UserNotFoundException":
+                    raise InvalidCredentialsException(
+                        message=str(e), detail="Please check your Secret Key is correct"
+                    )
+                elif e.response["Error"]["Code"] == "NotAuthorizedException":
+                    raise InvalidCredentialsException(
+                        message=str(e),
+                        detail="Please check your Access Key, that you have created your Api Token and that you are using the right STAX REGION",
+                    )
+                else:
+                    raise InvalidCredentialsException(
+                        f"Unexpected Client Error. Error details: {e}"
+                    )
             token = tokens["AuthenticationResult"]["IdToken"]
-        else:
-            token = jwt.encode({"sub": "unittest"}, "secret", algorithm="HS256")
         return token
 
     def sts_from_cognito_identity_pool(self, token, cognito_client=None):

staxapp/exceptions.py

@@ -26,3 +26,15 @@ class ValidationException(Exception):
     def __init__(self, message):
         # logging.info(f"VALIDATE: {message}")
         self.message = message
+
+
+class InvalidCredentialsException(Exception):
+    def __init__(self, message, detail=""):
+        logging.error(message)
+        prefix = f"InvalidCredentialsException: "
+        if detail:
+            prefix = f"{prefix}{detail} - "
+        self.message = f"{prefix}{message}"
+
+    def __str__(self):
+        return self.message

tests/test_api.py

@@ -209,8 +209,8 @@ def testFailedApiException(self):
         )
         with self.assertRaises(ApiException):
             self.Api.get("/test/no/error")
-        
-		# Test an exception which has no json in response
+
+        # Test an exception which has no json in response
         responses.add(
             responses.GET,
             f"{Config.api_base_url()}/test/invalid/json",
@@ -219,16 +219,14 @@ def testFailedApiException(self):
         )
         with self.assertRaises(ApiException):
             self.Api.get("/test/invalid/json")
-		
+
         # Test an exception with no content
         responses.add(
-            responses.GET,
-            f"{Config.api_base_url()}/test/no/content",
-            status=500,
+            responses.GET, f"{Config.api_base_url()}/test/no/content", status=500,
         )
         with self.assertRaises(ApiException):
             self.Api.get("/test/no/content")
-        
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_auth.py

@@ -12,9 +12,10 @@
 
 from botocore import UNSIGNED
 from botocore.client import Config as BotoConfig
-from botocore.stub import Stubber
+from botocore.stub import Stubber, ANY
 
 from staxapp.auth import StaxAuth
+from staxapp.exceptions import InvalidCredentialsException
 
 
 class StaxAuthTests(unittest.TestCase):
@@ -30,8 +31,16 @@ def setUp(self):
         )
         self.cognito_stub = Stubber(self.cognito_client)
 
+        self.aws_srp_client = botocore.session.get_session().create_client(
+            "cognito-idp",
+            region_name="ap-southeast-2",
+            config=BotoConfig(signature_version=UNSIGNED),
+        )
+        self.aws_srp_stubber = Stubber(self.aws_srp_client)
+
     def tearDown(self):
         self.cognito_stub.deactivate()
+        self.aws_srp_stubber.deactivate()
 
     def testStaxAuthInit(self):
         """
@@ -45,16 +54,58 @@ def testToken(self):
         Test valid JWT is returned
         """
         sa = StaxAuth("ApiAuth")
-        token = sa.id_token_from_cognito()
-        jwt_token = jwt.decode(token, verify=False)
-        self.assertIn("sub", jwt_token)
+        self.stub_aws_srp(sa, "valid_username")
+        token = sa.id_token_from_cognito(
+            username="valid_username", password="correct", client=self.aws_srp_client
+        )
+        self.assertEqual(token, "valid_token")
+
+    def testCredentialErrors(self):
+        """
+        Test that boto errors are caught and converted to InvalidCredentialExceptions
+      """
+
+        sa = StaxAuth("ApiAuth")
+        # Test with invalid username password
+        self.stub_aws_srp(sa, "bad_password", "UserNotFoundException")
+        user_not_found_success = False
+        try:
+            sa.id_token_from_cognito(
+                username="bad_password", password="wrong", client=self.aws_srp_client
+            )
+        except InvalidCredentialsException as e:
+            self.assertIn("Please check your Secret Key is correct", e.message)
+            user_not_found_success = True
+        self.assertTrue(user_not_found_success)
+
+        # Test with no access
+        self.stub_aws_srp(sa, "no_access", "NotAuthorizedException")
+        no_access_success = False
+        try:
+            sa.id_token_from_cognito(
+                username="no_access", password="wrong", client=self.aws_srp_client
+            )
+        except InvalidCredentialsException as e:
+            self.assertIn(
+                "Please check your Access Key, that you have created your Api Token and that you are using the right STAX REGION",
+                e.message,
+            )
+            no_access_success = True
+        self.assertTrue(no_access_success)
+
+        # Test Unknown Error
+        self.stub_aws_srp(sa, "Unknown", "UnitTesting")
+        with self.assertRaises(InvalidCredentialsException):
+            sa.id_token_from_cognito(
+                username="Unknown", password="wrong", client=self.aws_srp_client
+            )
 
     def testCreds(self):
         """
         Test valid credentials are returned
         """
         sa = StaxAuth("ApiAuth")
-        token = sa.id_token_from_cognito()
+        token = jwt.encode({"sub": "unittest"}, "secret", algorithm="HS256")
         jwt_token = jwt.decode(token, verify=False)
         self.stub_cognito_creds(jwt_token.get("sub"))
         creds = sa.sts_from_cognito_identity_pool(
@@ -63,6 +114,56 @@ def testCreds(self):
         self.assertIn("Credentials", creds)
         self.assertTrue(creds.get("IdentityId").startswith("ap-southeast-2"))
 
+    def testAuthErrors(self):
+        """
+      Test that errors are thrown when keys are invalid
+      """
+        sa = StaxAuth("ApiAuth")
+        # Test with no username
+        with self.assertRaises(InvalidCredentialsException):
+            sa.requests_auth(username=None, password="valid")
+
+        # Test with no username
+        with self.assertRaises(InvalidCredentialsException):
+            sa.requests_auth(username="valid", password=None)
+
+    def stub_aws_srp(self, stax_auth, username, error_code=None):
+        expected_parameters = {
+            "AuthFlow": "USER_SRP_AUTH",
+            "AuthParameters": {"SRP_A": ANY, "USERNAME": username},
+            "ClientId": stax_auth.client_id,
+        }
+        if error_code:
+            self.aws_srp_stubber.add_client_error(
+                "initiate_auth",
+                service_error_code=error_code,
+                expected_params=expected_parameters,
+            )
+        else:
+            self.aws_srp_stubber.add_response(
+                "initiate_auth",
+                {
+                    "ChallengeParameters": {
+                        "USER_ID_FOR_SRP": "user",
+                        "SALT": "4",
+                        "SRP_B": "5",
+                        "SECRET_BLOCK": "secblock",
+                    },
+                    "ChallengeName": "PASSWORD_VERIFIER",
+                },
+                expected_parameters,
+            )
+            self.aws_srp_stubber.add_response(
+                "respond_to_auth_challenge",
+                {"AuthenticationResult": {"IdToken": "valid_token"},},
+                {
+                    "ClientId": stax_auth.client_id,
+                    "ChallengeName": ANY,
+                    "ChallengeResponses": ANY,
+                },
+            )
+        self.aws_srp_stubber.activate()
+
     def stub_cognito_creds(self, token: str):
         sa = StaxAuth("ApiAuth")