Merge pull request #225 from Axonius/develop

Bugfixes for QA gold release SA-3419

axonius_api_client/api/api_endpoint.py

@@ -110,15 +110,11 @@ def __post_init__(self):
     def str_properties(self) -> t.List[str]:
         """Get the properties for this endpoint as a list of strs."""
         return [
-            f"method={self.method!r}",
-            f"path={self.path!r}",
+            f"method={self.method!r}, path={self.path!r}",
             f"request_schema={get_cls_path(self.request_schema_cls)}",
             f"request_model={get_cls_path(self.request_model_cls)}",
             f"response_schema={get_cls_path(self.response_schema_cls)}",
             f"response_model={get_cls_path(self.response_model_cls)}",
-            f"request_as_none={self.request_as_none}",
-            f"response_as_text={self.response_as_text}",
-            f"http_args_required={self.http_args_required}",
         ]
 
     def perform_request(
@@ -253,6 +249,10 @@ def handle_response(
             data = self.load_response(
                 http=http, response=response, **combo_dicts(kwargs, data=data)
             )
+            try:
+                data.RESPONSE = response
+            except Exception:
+                pass
         return data
 
     def get_response_json(self, response: requests.Response) -> JSON_TYPES:

axonius_api_client/api/api_endpoints.py

@@ -1290,6 +1290,31 @@ class DataScopes(ApiEndpointGroup):
     )
 
 
+@dataclasses.dataclass(eq=True, frozen=True, repr=False)
+class Account(ApiEndpointGroup):
+    """Pass."""
+
+    login: ApiEndpoint = ApiEndpoint(
+        method="post",
+        path="api/login",
+        request_schema_cls=json_api.account.LoginRequestSchema,
+        request_model_cls=json_api.account.LoginRequest,
+        response_schema_cls=json_api.account.LoginResponseSchema,
+        response_model_cls=json_api.account.LoginResponse,
+    )
+
+    get_api_keys: ApiEndpoint = ApiEndpoint(
+        method="get",
+        path="api/settings/api_key",
+        request_schema_cls=None,
+        request_model_cls=None,
+        response_schema_cls=None,
+        response_model_cls=None,
+    )
+
+    validate: ApiEndpoint = SystemSettings.get_constants
+
+
 @dataclasses.dataclass(eq=True, frozen=True, repr=False)
 class ApiEndpoints(BaseData):
     """Pass."""
@@ -1313,6 +1338,7 @@ class ApiEndpoints(BaseData):
     dashboard_spaces: ApiEndpointGroup = DashboardSpaces()
     folders_queries: ApiEndpointGroup = FoldersQueries()
     folders_enforcements: ApiEndpointGroup = FoldersEnforcements()
+    account: ApiEndpointGroup = Account()
 
     @classmethod
     def get_groups(cls) -> Dict[str, ApiEndpointGroup]:

axonius_api_client/api/json_api/__init__.py

@@ -2,6 +2,7 @@
 """Models for API requests & responses."""
 
 from . import (
+    account,
     adapters,
     assets,
     audit_logs,
@@ -57,4 +58,5 @@
     "dashboard_spaces",
     "spaces_export",
     "folders",
+    "account",
 )

axonius_api_client/api/json_api/account.py

@@ -0,0 +1,132 @@
+# -*- coding: utf-8 -*-
+"""Models for API requests & responses."""
+import dataclasses
+import typing as t
+
+import marshmallow
+
+from ...exceptions import AuthError
+from .base import BaseModel, BaseSchemaJson
+from .custom_fields import SchemaBool
+from .generic import Metadata, MetadataSchema
+
+
+class LoginRequestSchema(BaseSchemaJson):
+    """Schema for issuing a login request."""
+
+    user_name = marshmallow.fields.Str(
+        allow_none=True,
+        dump_default=None,
+        load_default=None,
+        description="Axonius User Name",
+    )
+    password = marshmallow.fields.Str(
+        allow_none=True,
+        dump_default=None,
+        load_default=None,
+        description="Axonius Password",
+    )
+    saml_token = marshmallow.fields.Str(
+        allow_none=True,
+        dump_default=None,
+        load_default=None,
+        description="SAML token from 2FA negotiation",
+    )
+    remember_me = SchemaBool(
+        load_default=False, dump_default=False, description="Used for browser controls"
+    )
+    eula_agreed = SchemaBool(
+        load_default=False, dump_default=False, description="EULA has been agreed to by user"
+    )
+
+    class Meta:
+        """Pass."""
+
+        type_ = "login_schema"
+
+    @staticmethod
+    def get_model_cls() -> t.Optional[type]:
+        """Pass."""
+        return LoginRequest
+
+
+LOGIN_REQUEST_SCHEMA = LoginRequestSchema()
+
+
+@dataclasses.dataclass
+class LoginRequest(BaseModel):
+    """Model for issuing a login request."""
+
+    user_name: t.Optional[str] = None
+    password: t.Optional[str] = None
+    saml_token: t.Optional[str] = None
+    remember_me: bool = False
+    eula_agreed: bool = False
+
+    @staticmethod
+    def get_schema_cls() -> t.Optional[type]:
+        """Pass."""
+        return LoginRequestSchema
+
+    def _check_credential(self, attr: str) -> str:
+        """Check that a credential is a non-empty string."""
+        value: t.Any = getattr(self, attr)
+
+        if isinstance(value, str) and value.strip():
+            value = value.strip()
+            setattr(self, attr, value)
+            return value
+
+        field: marshmallow.Field = LOGIN_REQUEST_SCHEMA.declared_fields[attr]
+        description: str = field.metadata.get("description", f"{attr}")
+        msgs: t.List[str] = [
+            f"Value provided for {description} is not a non-empty string"
+            f"Provided type {type(value)}, value: {value!r}"
+        ]
+        raise AuthError(msgs)
+
+    def check_credentials(self):
+        """Check that username and password are not empty."""
+        self._check_credential(attr="user_name")
+        self._check_credential(attr="password")
+
+
+class LoginResponseSchema(MetadataSchema):
+    """Schema for receiving a login response."""
+
+    class Meta:
+        """Pass."""
+
+        type_ = "metadata_schema"
+
+    @staticmethod
+    def get_model_cls() -> t.Optional[type]:
+        """Pass."""
+        return LoginResponse
+
+
+@dataclasses.dataclass
+class LoginResponse(Metadata):
+    """Model for receiving a login response."""
+
+    document_meta: t.Optional[dict] = dataclasses.field(default_factory=dict)
+
+    @staticmethod
+    def get_schema_cls() -> t.Optional[type]:
+        """Pass."""
+        return LoginResponseSchema
+
+    @property
+    def access_token(self) -> str:
+        """Get the Access token for use in auth header."""
+        return self.document_meta["access_token"]
+
+    @property
+    def refresh_token(self) -> str:
+        """Get the Refresh token for use in auth header."""
+        return self.document_meta["refresh_token"]
+
+    @property
+    def authorization(self) -> str:
+        """Get the value to use in the authorization header."""
+        return f"Bearer {self.access_token}"

axonius_api_client/api/system/signup.py

@@ -143,4 +143,4 @@ def __init__(self, url, **kwargs):
         log_level = kwargs.get("log_level", LOG_LEVEL_API)
         self.LOG = get_obj_log(obj=self, level=log_level)
         kwargs.setdefault("certwarn", False)
-        self.http = Http(url=url, **kwargs)
+        self.HTTP = self.http = Http(url=url, **kwargs)

axonius_api_client/auth/__init__.py

@@ -1,7 +1,8 @@
 # -*- coding: utf-8 -*-
 """Authenticating with Axonius."""
-from . import api_key, models
+from . import api_key, credentials, models
 from .api_key import ApiKey
+from .credentials import Credentials
 from .models import Mixins, Model
 
 __all__ = (
@@ -10,4 +11,6 @@
     "Model",
     "Mixins",
     "ApiKey",
+    "credentials",
+    "Credentials",
 )

axonius_api_client/auth/credentials.py

@@ -0,0 +1,70 @@
+# -*- coding: utf-8 -*-
+"""Authentication via API key and API secret."""
+import typing as t
+
+from ..api.api_endpoints import ApiEndpoint
+from ..api.json_api.account import LoginRequest, LoginResponse
+from ..http import Http
+from .models import Mixins
+
+
+class Credentials(Mixins):
+    """Authentication method using username and password credentials."""
+
+    def __init__(
+        self,
+        http: Http,
+        username: t.Optional[str] = None,
+        password: t.Optional[str] = None,
+        **kwargs,
+    ):
+        """Authenticate using username and password.
+
+        Args:
+            http (Http): HTTP client to use to send requests
+            username (t.Optional[str], optional): Axonius User Name
+            password (t.Optional[str], optional): Axonius Password
+            prompt (bool, optional): Prompt for credentials that are not non-empty strings
+        """
+        creds: LoginRequest = LoginRequest(user_name=username, password=password, eula_agreed=True)
+        super().__init__(http=http, creds=creds, **kwargs)
+
+    def login(self):
+        """Login to API."""
+        if not self.is_logged_in:
+            self._creds.check_credentials()
+            self._login_response: LoginResponse = self._login(request_obj=self._creds)
+            headers: dict = {"authorization": self._login_response.authorization}
+            self._api_keys: dict = self._get_api_keys(headers=headers)
+            self.http.session.headers["api-key"] = self._api_keys["api_key"]
+            self.http.session.headers["api-secret"] = self._api_keys["api_secret"]
+            self._validate()
+            self._logged_in = True
+            self.LOG.debug(f"Successfully logged in using {self._cred_fields}")
+
+    def logout(self):
+        """Logout from API."""
+        super().logout()
+
+    def _login(self, request_obj: LoginRequest) -> LoginResponse:
+        """Direct API method to issue a login request.
+
+        Args:
+            request_obj (LoginRequest): Request object to send
+
+        Returns:
+            LoginResponse: Response object received
+        """
+        endpoint: ApiEndpoint = self.endpoints.login
+        response: LoginResponse = endpoint.perform_request(http=self.http, request_obj=request_obj)
+        return response
+
+    @property
+    def _cred_fields(self) -> t.List[str]:
+        """Credential fields used by this auth model."""
+        return [f"username={self._creds.user_name!r}", "password"]
+
+    def _logout(self):
+        """Logout from API."""
+        self._logged_in = False
+        self.http.session.headers = {}

axonius_api_client/auth/models.py

@@ -2,6 +2,7 @@
 """Authentication models."""
 import abc
 import logging
+import typing as t
 
 from ..api.api_endpoint import ApiEndpoint
 from ..api.api_endpoints import ApiEndpoints
@@ -46,10 +47,7 @@ class Mixins(Model):
     _logged_in: bool = False
     """Attribute checked by :meth:`is_logged_in`."""
 
-    _validate_endpoint: ApiEndpoint = ApiEndpoints.system_settings.get_constants
-    """Endpoint to use to validate logins."""
-
-    def __init__(self, http: Http, creds: dict, **kwargs):
+    def __init__(self, http: Http, creds: t.Any, **kwargs):
         """Mixins for Auth Models.
 
         Args:
@@ -64,7 +62,7 @@ def __init__(self, http: Http, creds: dict, **kwargs):
         self._http: Http = http
         """HTTP Client."""
 
-        self._creds: dict = creds
+        self._creds: t.Any = creds
         """Credential store."""
 
         self._check_http_lock()
@@ -97,6 +95,25 @@ def is_logged_in(self) -> bool:
         """Check if login has been called."""
         return self._logged_in
 
+    def get_api_keys(self) -> dict:
+        """Get the API key and secret for the current user."""
+        return self._get_api_keys()
+
+    @property
+    def endpoints(self) -> ApiEndpoints:
+        """Get the endpoint group for this module."""
+        return ApiEndpoints.account
+
+    def _validate(self):
+        """Validate credentials."""
+        try:
+            self.endpoints.validate.perform_request(http=self.http)
+        except Exception:
+            self._logged_in = False
+            raise
+        else:
+            self._logged_in = True
+
     def __str__(self) -> str:
         """Show object info."""
         bits = [f"url={self.http.url!r}", f"is_logged_in={self.is_logged_in}"]
@@ -119,16 +136,12 @@ def _check_http_lock(self):
         if auth_lock:
             raise AuthError(f"{self.http} already being used by {auth_lock}")
 
+    def _get_api_keys(self, **http_args) -> dict:
+        """Direct API method to get the API keys for the current user."""
+        endpoint: ApiEndpoint = self.endpoints.get_api_keys
+        response: dict = endpoint.perform_request(http=self.http, http_args=http_args)
+        return response
+
     def _set_http_lock(self):
         """Set HTTP Client auth lock."""
         self._http._auth_lock = self
-
-    def _validate(self):
-        """Validate credentials."""
-        try:
-            self._validate_endpoint.perform_request(http=self.http)
-        except Exception:
-            self._logged_in = False
-            raise
-        else:
-            self._logged_in = True