Added token validation

Author: VaghinakDev

src/superannotate/lib/app/interface/base_interface.py

@@ -15,6 +15,7 @@
 from lib.core import CONFIG
 from lib.core import setup_logging
 from lib.core.entities.base import ConfigEntity
+from lib.core.entities.base import TokenStr
 from lib.core.exceptions import AppException
 from lib.infrastructure.controller import Controller
 from lib.infrastructure.utils import extract_project_folder
@@ -26,39 +27,41 @@
 class BaseInterfaceFacade:
     REGISTRY = []
 
-    def __init__(self, token: str = None, config_path: str = None):
-        if token:
-            config = ConfigEntity(SA_TOKEN=token)
-        elif config_path:
-            config_path = Path(config_path)
-            if not Path(config_path).is_file() or not os.access(config_path, os.R_OK):
-                raise AppException(
-                    f"SuperAnnotate config file {str(config_path)} not found."
-                )
-            try:
+    @validate_arguments
+    def __init__(self, token: TokenStr = None, config_path: str = None):
+        try:
+            if token:
+                config = ConfigEntity(SA_TOKEN=token)
+            elif config_path:
+                config_path = Path(config_path)
+                if not Path(config_path).is_file() or not os.access(config_path, os.R_OK):
+                    raise AppException(
+                        f"SuperAnnotate config file {str(config_path)} not found."
+                    )
                 if config_path.suffix == ".json":
                     config = self._retrieve_configs_from_json(config_path)
                 else:
                     config = self._retrieve_configs_from_ini(config_path)
-            except pydantic.ValidationError as e:
-                raise AppException(wrap_error(e))
-            except KeyError:
-                raise
-        else:
-            config = self._retrieve_configs_from_env()
-            if not config:
-                if Path(constants.CONFIG_INI_FILE_LOCATION).exists():
-                    config = self._retrieve_configs_from_ini(
-                        constants.CONFIG_INI_FILE_LOCATION
-                    )
-                elif Path(constants.CONFIG_JSON_FILE_LOCATION).exists():
-                    config = self._retrieve_configs_from_json(
-                        constants.CONFIG_JSON_FILE_LOCATION
-                    )
-                else:
-                    raise AppException(
-                        f"SuperAnnotate config file {constants.CONFIG_INI_FILE_LOCATION} not found."
-                    )
+
+            else:
+                config = self._retrieve_configs_from_env()
+                if not config:
+                    if Path(constants.CONFIG_INI_FILE_LOCATION).exists():
+                        config = self._retrieve_configs_from_ini(
+                            constants.CONFIG_INI_FILE_LOCATION
+                        )
+                    elif Path(constants.CONFIG_JSON_FILE_LOCATION).exists():
+                        config = self._retrieve_configs_from_json(
+                            constants.CONFIG_JSON_FILE_LOCATION
+                        )
+                    else:
+                        raise AppException(
+                            f"SuperAnnotate config file {constants.CONFIG_INI_FILE_LOCATION} not found."
+                        )
+        except pydantic.ValidationError as e:
+            raise AppException(wrap_error(e))
+        except KeyError:
+            raise
         if not config:
             raise AppException("Credentials not provided.")
         setup_logging(config.LOGGING_LEVEL, config.LOGGING_PATH)
@@ -70,7 +73,13 @@ def _retrieve_configs_from_json(path: Path) -> typing.Union[ConfigEntity]:
         with open(path) as json_file:
             json_data = json.load(json_file)
         token = json_data["token"]
-        config = ConfigEntity(SA_TOKEN=token)
+        try:
+            config = ConfigEntity(SA_TOKEN=token)
+        except pydantic.ValidationError:
+            raise pydantic.ValidationError(
+                [pydantic.error_wrappers.ErrorWrapper(ValueError("Invalid token."), loc='token')],
+                model=ConfigEntity
+            )
         host = json_data.get("main_endpoint")
         verify_ssl = json_data.get("ssl_verify")
         if host:
@@ -89,10 +98,7 @@ def _retrieve_configs_from_ini(path: Path) -> typing.Union[ConfigEntity]:
         config_data = {}
         for key in config_parser["DEFAULT"]:
             config_data[key.upper()] = config_parser["DEFAULT"][key]
-        try:
-            return ConfigEntity(**config_data)
-        except pydantic.ValidationError as e:
-            raise AppException(wrap_error(e))
+        return ConfigEntity(**config_data)
 
     @staticmethod
     def _retrieve_configs_from_env() -> typing.Union[ConfigEntity, None]:

src/superannotate/lib/app/interface/types.py

@@ -51,5 +51,4 @@ def wrapped(self, *args, **kwargs):
             return pydantic_validate_arguments(func)(self, *args, **kwargs)
         except ValidationError as e:
             raise AppException(wrap_error(e)) from e
-
     return wrapped

src/superannotate/lib/core/entities/base.py

@@ -1,3 +1,4 @@
+import re
 import warnings
 from datetime import datetime
 from enum import Enum
@@ -14,6 +15,7 @@
 from lib.core.enums import BaseTitledEnum
 from pydantic import BaseModel as PydanticBaseModel
 from pydantic import Extra
+from pydantic import StrictStr
 from pydantic import Field
 from pydantic.datetime_parse import parse_datetime
 from pydantic.typing import is_namedtuple
@@ -292,8 +294,21 @@ def map_fields(entity: dict) -> dict:
         return entity
 
 
+class TokenStr(StrictStr):
+    regex = r'^[-.@_A-Za-z0-9]+=\d+$'
+
+    @classmethod
+    def validate(cls, value: Union[str]) -> Union[str]:
+        if cls.curtail_length and len(value) > cls.curtail_length:
+            value = value[: cls.curtail_length]
+        if cls.regex:
+            if not re.match(cls.regex, value):
+                raise ValueError("Invalid token.")
+        return value
+
+
 class ConfigEntity(BaseModel):
-    API_TOKEN: str = Field(alias="SA_TOKEN")
+    API_TOKEN: TokenStr = Field(alias="SA_TOKEN")
     API_URL: str = Field(alias="SA_URL", default=BACKEND_URL)
     LOGGING_LEVEL: Literal[
         "NOTSET", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"

src/superannotate/lib/infrastructure/services/http_client.py

@@ -183,27 +183,30 @@ def serialize_response(
             "status": response.status_code,
         }
         try:
-            data_json = response.json()
             if not response.ok:
                 if response.status_code in (502, 504):
                     data[
                         "_error"
                     ] = "Our service is currently unavailable, please try again later."
+                    return content_type(**data)
                 else:
+                    data_json = response.json()
                     data["_error"] = data_json.get(
                         "error", data_json.get("errors", "Unknown Error")
                     )
-            else:
-                if dispatcher:
-                    if dispatcher in data_json:
-                        data["data"] = data_json.pop(dispatcher)
-                    else:
-                        data["data"] = data_json
-                        data_json = {}
-                    data.update(data_json)
+                    return content_type(**data)
+            data_json = response.json()
+            if dispatcher:
+                if dispatcher in data_json:
+                    data["data"] = data_json.pop(dispatcher)
                 else:
                     data["data"] = data_json
+                    data_json = {}
+                data.update(data_json)
+            else:
+                data["data"] = data_json
             return content_type(**data)
         except json.decoder.JSONDecodeError:
+            data['_error'] = response.content
             data["reason"] = response.reason
             return content_type(**data)

tests/integration/annotations/test_upload_annotations_from_folder_to_project.py

@@ -8,7 +8,8 @@
 sa = SAClient()
 
 
-class TestAnnotationUploadVector(BaseTestCase):
+class \
+        TestAnnotationUploadVector(BaseTestCase):
     PROJECT_NAME = "Test-Upload_annotations_from_folder_to_project"
     PROJECT_DESCRIPTION = "Desc"
     PROJECT_TYPE = "Vector"

tests/unit/test_init.py

@@ -6,6 +6,7 @@
 from unittest.mock import patch
 
 import superannotate.lib.core as constants
+from superannotate.lib.app.interface.types import validate_arguments
 from superannotate import AppException
 from superannotate import SAClient
 
@@ -15,10 +16,7 @@ class ClientInitTestCase(TestCase):
 
     def test_init_via_invalid_token(self):
         _token = "123"
-        with self.assertRaisesRegexp(
-            AppException,
-            "Unable to retrieve team data. Please verify your credentials.",
-        ):
+        with self.assertRaisesRegexp(AppException, r"(\s+)token(\s+)Invalid token."):
             SAClient(token=_token)
 
     @patch("lib.core.usecases.GetTeamUseCase")
@@ -49,6 +47,20 @@ def test_init_via_config_json(self, get_team_use_case):
                         self._token.split("=")[-1]
                     )
 
+    def test_init_via_config_json_invalid_json(self):
+        with tempfile.TemporaryDirectory() as config_dir:
+            constants.HOME_PATH = config_dir
+            config_ini_path = f"{config_dir}/config.ini"
+            config_json_path = f"{config_dir}/config.json"
+            with patch("lib.core.CONFIG_INI_FILE_LOCATION", config_ini_path), patch(
+                    "lib.core.CONFIG_JSON_FILE_LOCATION", config_json_path
+            ):
+                with open(f"{config_dir}/config.json", "w") as config_json:
+                    json.dump({"token": "INVALID_TOKEN"}, config_json)
+                for kwargs in ({}, {"config_path": f"{config_dir}/config.json"}):
+                    with self.assertRaisesRegexp(AppException, r"(\s+)token(\s+)Invalid token."):
+                        SAClient(**kwargs)
+
     @patch("lib.core.usecases.GetTeamUseCase")
     def test_init_via_config_ini(self, get_team_use_case):
         with tempfile.TemporaryDirectory() as config_dir:
@@ -83,6 +95,32 @@ def test_init_env(self, get_team_use_case):
         assert sa.controller._config.API_URL == "SOME_URL"
         assert get_team_use_case.call_args_list[0].kwargs["team_id"] == 123
 
+    @patch.dict(os.environ, {"SA_URL": "SOME_URL", "SA_TOKEN": "SOME_TOKEN"})
+    def test_init_env_invalid_token(self):
+        with self.assertRaisesRegexp(AppException, r"(\s+)SA_TOKEN(\s+)Invalid token."):
+            SAClient()
+
+    def test_init_via_config_ini_invalid_token(self):
+        with tempfile.TemporaryDirectory() as config_dir:
+            constants.HOME_PATH = config_dir
+            config_ini_path = f"{config_dir}/config.ini"
+            config_json_path = f"{config_dir}/config.json"
+            with patch("lib.core.CONFIG_INI_FILE_LOCATION", config_ini_path), patch(
+                    "lib.core.CONFIG_JSON_FILE_LOCATION", config_json_path
+            ):
+                with open(f"{config_dir}/config.ini", "w") as config_ini:
+                    config_parser = ConfigParser()
+                    config_parser.optionxform = str
+                    config_parser["DEFAULT"] = {
+                        "SA_TOKEN": "INVALID_TOKEN",
+                        "LOGGING_LEVEL": "DEBUG",
+                    }
+                    config_parser.write(config_ini)
+
+                for kwargs in ({}, {"config_path": f"{config_dir}/config.ini"}):
+                    with self.assertRaisesRegexp(AppException, r"(\s+)SA_TOKEN(\s+)Invalid token."):
+                        SAClient(**kwargs)
+
     def test_invalid_config_path(self):
         _path = "something"
         with self.assertRaisesRegexp(