Merge pull request #31 from stax-labs/feat/cleanup

chore(staxapp): code cleanup

Author: Mat Lord

Makefile

@@ -27,9 +27,9 @@ lint: install
 	${ISORT} --diff staxapp/*.py
 	${BLACK} -t py37 --check --diff staxapp/
 
-format: lint
+format:
 	${ISORT} --apply staxapp/*.py
-	${BLACK} -t py37 staxapp/
+	${BLACK} -t py37 staxapp/*.py
 
 download-schema:
 	curl --fail --compressed -s -o staxapp/data/schema.json https://api.au1.staxapp.cloud/20190206/public/api-document

staxapp/api.py

@@ -1,6 +1,3 @@
-import json
-import logging
-
 import requests
 
 from staxapp.config import Config
@@ -13,7 +10,7 @@ class Api:
     @classmethod
     def _auth(cls):
         if not cls._requests_auth:
-            cls._requests_auth = Config.auth().requests_auth(
+            cls._requests_auth = Config.get_auth_class().requests_auth(
                 Config.access_key, Config.secret_key
             )
         return cls._requests_auth
@@ -31,16 +28,15 @@ def get(cls, url_frag, params={}, **kwargs):
         url = f"{Config.api_base_url()}/{url_frag.lstrip('/')}"
 
         response = requests.get(url, auth=cls._auth(), params=params, **kwargs)
-        # logging.debug(f"GET: {response.text}")
         cls.handle_api_response(response)
         return response.json()
 
     @classmethod
     def post(cls, url_frag, payload={}, **kwargs):
         url_frag = url_frag.replace(f"/{Config.API_VERSION}", "")
         url = f"{Config.api_base_url()}/{url_frag.lstrip('/')}"
+
         response = requests.post(url, json=payload, auth=cls._auth(), **kwargs)
-        # logging.debug(f"POST: {response.text}")
         cls.handle_api_response(response)
         return response.json()
 

staxapp/auth.py

@@ -1,30 +1,27 @@
 #!/usr/local/bin/python3
-import logging
-import sys
-from datetime import date, datetime, timedelta, timezone
+from datetime import datetime, timezone
 
 import boto3
-import jwt
 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 warrant import AWSSRP
 
-from staxapp.config import Config as JumaConfig
+from staxapp.config import Config as StaxConfig
 from staxapp.exceptions import InvalidCredentialsException
 
 
 class StaxAuth:
     def __init__(self, config_branch):
-        config = JumaConfig.api_config
+        config = StaxConfig.api_config
 
         self.identity_pool = config.get(config_branch).get("identityPoolId")
         self.user_pool = config.get(config_branch).get("userPoolId")
         self.client_id = config.get(config_branch).get("userPoolWebClientId")
         self.aws_region = config.get(config_branch).get("region")
 
-    def requests_auth(self, username, password):
+    def requests_auth(self, username, password, **kwargs):
         if username is None:
             raise InvalidCredentialsException(
                 "Please provide an Access Key to your config"
@@ -34,51 +31,52 @@ def requests_auth(self, username, password):
                 "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)
+        id_token = self.id_token_from_cognito(username, password, **kwargs)
+        id_creds = self.sts_from_cognito_identity_pool(id_token, **kwargs)
         auth = self.sigv4_signed_auth_headers(id_creds)
 
-        JumaConfig.expiration = id_creds.get("Credentials").get("Expiration")
-        JumaConfig.auth = auth
+        StaxConfig.expiration = id_creds.get("Credentials").get("Expiration")
+        StaxConfig.auth = auth
 
-        return JumaConfig.auth
+        return StaxConfig.auth
 
-    def id_token_from_cognito(self, username=None, password=None, client=None):
+    def id_token_from_cognito(
+        self, username=None, password=None, srp_client=None, **kwargs
+    ):
         token = None
-        if username and password:
-            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,
+        if not srp_client:
+            srp_client = boto3.client(
+                "cognito-idp",
+                region_name=self.aws_region,
+                config=BotoConfig(signature_version=UNSIGNED),
             )
-            try:
-                tokens = aws.authenticate_user()
-            except ClientError as e:
-                if e.response["Error"]["Code"] == "NotAuthorizedException":
-                    raise InvalidCredentialsException(
-                        message=str(e), detail="Please check your Secret Key is correct"
-                    )
-                elif e.response["Error"]["Code"] == "UserNotFoundException":
-                    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"]
+        aws = AWSSRP(
+            username=username,
+            password=password,
+            pool_id=self.user_pool,
+            client_id=self.client_id,
+            client=srp_client,
+        )
+        try:
+            tokens = aws.authenticate_user()
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NotAuthorizedException":
+                raise InvalidCredentialsException(
+                    message=str(e), detail="Please check your Secret Key is correct"
+                )
+            elif e.response["Error"]["Code"] == "UserNotFoundException":
+                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"]
         return token
 
-    def sts_from_cognito_identity_pool(self, token, cognito_client=None):
+    def sts_from_cognito_identity_pool(self, token, cognito_client=None, **kwargs):
         if not cognito_client:
             cognito_client = boto3.client(
                 "cognito-identity",
@@ -91,43 +89,39 @@ def sts_from_cognito_identity_pool(self, token, cognito_client=None):
                 f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
             },
         )
-        # logging.debug(f"ID: {id}")
-
         id_creds = cognito_client.get_credentials_for_identity(
             IdentityId=id["IdentityId"],
             Logins={
                 f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
             },
         )
-        # logging.debug(f"CREDS: {id_creds}")
         return id_creds
 
     def sigv4_signed_auth_headers(self, id_creds):
         auth = AWSRequestsAuth(
             aws_access_key=id_creds.get("Credentials").get("AccessKeyId"),
             aws_secret_access_key=id_creds.get("Credentials").get("SecretKey"),
             aws_token=id_creds.get("Credentials").get("SessionToken"),
-            aws_host=f"{JumaConfig.hostname}",
+            aws_host=f"{StaxConfig.hostname}",
             aws_region=self.aws_region,
             aws_service="execute-api",
         )
-        # logging.debug(f"AUTH: {auth}")
         return auth
 
 
 class RootAuth:
     @staticmethod
-    def requests_auth(username, password):
-        if JumaConfig.expiration and JumaConfig.expiration > datetime.now(timezone.utc):
-            return JumaConfig.auth
+    def requests_auth(username, password, **kwargs):
+        if StaxConfig.expiration and StaxConfig.expiration > datetime.now(timezone.utc):
+            return StaxConfig.auth
 
-        return StaxAuth("JumaAuth").requests_auth(username, password)
+        return StaxAuth("JumaAuth").requests_auth(username, password, **kwargs)
 
 
 class ApiTokenAuth:
     @staticmethod
-    def requests_auth(username, password):
-        if JumaConfig.expiration and JumaConfig.expiration > datetime.now(timezone.utc):
-            return JumaConfig.auth
+    def requests_auth(username, password, **kwargs):
+        if StaxConfig.expiration and StaxConfig.expiration > datetime.now(timezone.utc):
+            return StaxConfig.auth
 
-        return StaxAuth("ApiAuth").requests_auth(username, password)
+        return StaxAuth("ApiAuth").requests_auth(username, password, **kwargs)

staxapp/config.py

@@ -37,7 +37,6 @@ def set_config(cls):
         cls.base_url = f"https://api.{cls.STAX_REGION}/{cls.API_VERSION}"
         config_url = f"{cls.api_base_url()}/public/config"
         config_response = requests.get(config_url)
-        # logging.debug(f"IDAM: get config from {config_url}")
         try:
             config_response.raise_for_status()
         except requests.exceptions.HTTPError as e:
@@ -71,8 +70,7 @@ def schema_url(cls):
         return f"{cls.base_url}/public/api-document"
 
     @classmethod
-    def auth(cls):
-        # logging.debug(f"AUTHCLASS: {cls.auth_class}")
+    def get_auth_class(cls):
         if cls.auth_class is None:
             from staxapp.auth import ApiTokenAuth
 

staxapp/contract.py

@@ -38,7 +38,6 @@ def validate(cls, data, component):
         Validates a request body against an component in a openapi3.0 template
         """
         if not cls._swagger_doc:
-            # logging.info(f"SCHEMA: no swagger defined, loading default template")
             cls.set_schema(cls.default_swagger_template())
 
         components = cls._resolved_schema.get("components")

staxapp/exceptions.py

@@ -24,7 +24,6 @@ def __str__(self):
 
 class ValidationException(Exception):
     def __init__(self, message):
-        # logging.info(f"VALIDATE: {message}")
         self.message = message
 
 

staxapp/openapi.py

@@ -1,11 +1,7 @@
 import json
-import logging
 import os
-import uuid
 
 import requests
-from jsonschema import validate
-from prance import ResolvingParser
 
 from staxapp.api import Api
 from staxapp.auth import ApiTokenAuth
@@ -26,7 +22,6 @@ def __init__(self, classname, lambda_client=None):
         if not self._initialized:
             self._map_paths_to_operations()
             StaxContract.set_schema(self._schema)
-            # logging.info(f"{self._operation_map}")
             if not self._operation_map.get(self.classname):
                 raise ValidationException(
                     f"No such class: {self.classname}. Please use one of {list(self._operation_map)}"
@@ -44,14 +39,12 @@ def __init__(self, classname, lambda_client=None):
     @classmethod
     def _load_schema(cls):
         if Config.load_live_schema:
-            # logging.info(f"SCHEMA: loading from {Config.schema_url()}")
             schema_response = requests.get(Config.schema_url())
             schema_response.raise_for_status()
             cls._schema = schema_response.json()
         else:
             current_dir = os.path.abspath(os.path.dirname(__file__))
             schema_file = f"{current_dir}/data/schema.json"
-            # logging.info(f"SCHEMA: loading from {schema_file}")
             with open(schema_file, "r") as f:
                 cls._schema = json.load(f)