Pebryan/7 16 21 sent api update (#187)

* added incident collection

* updated pandas version

* pandas version check

* package versions

* Updated docstrings

* linting fixes

Co-authored-by: Pete Bryan <pebryan@microsoft.com>

conda/conda-reqs.txt

@@ -17,7 +17,7 @@ msrest>=0.6.0
 msrestazure>=0.6.0
 networkx>=2.2
 numpy>=1.15.4
-pandas>=0.25.0
+pandas>=1.1.5
 python-dateutil>=2.8.1
 pytz>=2019.2
 pyyaml>=3.13

msticpy/data/azure_data.py

@@ -121,12 +121,23 @@ def __init__(self, connect: bool = False):
         if connect:
             self.connect()
 
-    def connect(
-        self,
-        auth_methods: List = None,
-        silent: bool = False,
-    ):
-        """Authenticate with the SDK."""
+    def connect(self, auth_methods: List = None, silent: bool = False):
+        """
+        Authenticate to the Azure SDK.
+
+        Parameters
+        ----------
+        auth_methods : List, optional
+            list of prefered authentication methods to use, by default None
+        silent : bool, optional
+            Set true to prevent output during auth process, by default False
+
+        Raises
+        ------
+        CloudError
+            If no valid credentials are found or if subscription client can't be created
+
+        """
         self.credentials = az_connect(auth_methods=auth_methods, silent=silent)
         if not self.credentials:
             raise CloudError("Could not obtain credentials.")
@@ -137,7 +148,20 @@ def connect(
         self.connected = True
 
     def get_subscriptions(self) -> pd.DataFrame:
-        """Get details of all subscriptions within the tenant."""
+        """
+        Get details of all subscriptions within the tenant.
+
+        Returns
+        -------
+        pd.DataFrame
+            Details of the subscriptions present in the users tentant.
+
+        Raises
+        ------
+        MsticpyNotConnectedError
+            If .connect() has not been called
+
+        """
         if self.connected is False:
             raise MsticpyNotConnectedError(
                 "You need to connect to the service before using this function.",
@@ -181,6 +205,16 @@ def get_subscription_info(self, sub_id: str) -> dict:
         sub_id : str
             The ID of the subscription to return details on.
 
+        Returns
+        -------
+        dict
+            Details on the selected subscription.
+
+        Raises
+        ------
+        MsticpyNotConnectedError
+            If .connect() has not been called.
+
         """
         if self.connected is False:
             raise MsticpyNotConnectedError(

msticpy/data/azure_sentinel.py

@@ -12,6 +12,7 @@
 from azure.common.exceptions import CloudError
 
 from .azure_data import AzureData
+from ..common.azure_auth_core import AzCredentials
 from ..common.exceptions import MsticpyAzureConfigError
 from ..common.wsconfig import WorkspaceConfig
 
@@ -28,18 +29,36 @@ class AzureSentinel(AzureData):
     """Class for returning key Azure Sentinel elements."""
 
     def __init__(self, connect: bool = False):
-        """Initialize connector for Azure APIs."""
+        """
+        Initialize connector for Azure APIs.
+
+        Parameters
+        ----------
+        connect : bool, optional
+            Set true if you want to connect to API on initalization, by default False
+
+        """
         super().__init__()
         self.config = None
 
-    def connect(
-        self,
-        auth_methods: List = None,
-        silent: bool = False,
-    ):
-        """Authenticate with the SDK & API."""
+    def connect(self, auth_methods: List = None, silent: bool = False, **kwargs):
+        """
+        Authenticate with the SDK & API.
+
+        Parameters
+        ----------
+        auth_methods : List, optional
+            list of prefered authentication methods to use, by default None
+        silent : bool, optional
+            Set true to prevent output during auth process, by default False
+
+        """
         super().connect(auth_methods=auth_methods, silent=silent)
-        self.token = _get_token(self.credentials)
+        if "token" in kwargs:
+            self.token = kwargs["token"]
+        else:
+            self.token = _get_token(self.credentials)  # type: ignore
+
         self.res_group_url = None
         self.prov_path = None
 
@@ -318,13 +337,14 @@ def get_incidents(
 
         return incidents_df
 
-    def get_incident(
+    def get_incident(  # pylint: disable=too-many-locals
         self,
         incident_id: str,
         res_id: str = None,
         sub_id: str = None,
         res_grp: str = None,
         ws_name: str = None,
+        entities: bool = False,
     ) -> pd.DataFrame:
         """
         Get details on a specific incident.
@@ -341,7 +361,8 @@ def get_incident(
             Resource Group name of the workspace, to be used if not providing Resource ID.
         ws_name : str, optional
             Workspace name of the workspace, to be used if not providing Resource ID.
-
+        entities : bool, optional
+            If True, include all entities in the response. Default is False.
 
         Returns
         -------
@@ -379,6 +400,20 @@ def get_incident(
         else:
             raise CloudError(response=response)
 
+        if entities:
+            entities_url = incident_url + "/entities"
+            ent_parameters = {"api-version": "2019-01-01-preview"}
+            ents = requests.post(
+                entities_url,
+                headers=_get_api_headers(self.token),
+                params=ent_parameters,
+            )
+            if ents.status_code == 200:
+                unique_entities = [
+                    (ent["kind"], ent["properties"]) for ent in ents.json()["entities"]
+                ]
+                incident_df["Entities"] = [unique_entities]
+
         return incident_df
 
     def update_incident(
@@ -539,8 +574,21 @@ def _check_config(self, items: List) -> Dict:
         return config_items
 
 
-def _build_paths(resid) -> str:
-    """Build a API URL from an Azure resource ID."""
+def _build_paths(resid: str) -> str:
+    """
+    Build an API URL from an Azure resource ID.
+
+    Parameters
+    ----------
+    resid : str
+        An Azure resource ID.
+
+    Returns
+    -------
+    str
+        A URI to that resource.
+
+    """
     res_info = {
         "subscription_id": resid.split("/")[2],
         "resource_group": resid.split("/")[4],
@@ -556,22 +604,66 @@ def _build_paths(resid) -> str:
     return url_part1 + url_part2 + url_part3
 
 
-def _get_token(credential) -> str:
-    """Extract token from a azure.identity object."""
+def _get_token(credential: AzCredentials) -> str:
+    """
+    Extract token from a azure.identity object.
+
+    Parameters
+    ----------
+    credential : AzCredentials
+        Azure OAuth credentials.
+
+    Returns
+    -------
+    str
+        A token to be used in API calls.
+
+    """
     token = credential.modern.get_token("https://management.azure.com/.default")
     return token.token
 
 
-def _get_api_headers(token):
-    """Return authorization header with current token."""
+def _get_api_headers(token: str) -> Dict:
+    """
+    Return authorization header with current token.
+
+    Parameters
+    ----------
+    token : str
+        Auzre auth token.
+
+    Returns
+    -------
+    Dict
+        A dictionary of headers to be used in API calls.
+
+    """
     return {
         "Authorization": f"Bearer {token}",
         "Content-Type": "application/json",
     }
 
 
-def _azs_api_result_to_df(response) -> pd.DataFrame:
-    """Convert API response to a Pandas dataframe."""
+def _azs_api_result_to_df(response: requests.Response) -> pd.DataFrame:
+    """
+    Convert API response to a Pandas dataframe.
+
+    Parameters
+    ----------
+    response : requests.Response
+        A response object from an Azure REST API call.
+
+    Returns
+    -------
+    pd.DataFrame
+        The API response as a Pandas dataframe.
+
+    Raises
+    ------
+    ValueError
+        If the response is not valid JSON.
+
+    """
     j_resp = response.json()
     if response.status_code != 200 or not j_resp:
         raise ValueError("No valid JSON result in response")
@@ -581,7 +673,20 @@ def _azs_api_result_to_df(response) -> pd.DataFrame:
 
 
 def _build_data(items: dict, **kwargs) -> dict:
-    """Build request data body from items."""
+    """
+    Build request data body from items.
+
+    Parameters
+    ----------
+    items : dict
+        A set pf items to be formated in the request body.
+
+    Returns
+    -------
+    dict
+        The request body formatted for the API.
+
+    """
     data_body = {"properties": {}}  # type: Dict[str, Dict[str, str]]
     for key in items.keys():
         if key in ["severity", "status", "title", "message"]:

requirements-all.txt

@@ -35,13 +35,14 @@ nest_asyncio>=1.4.0
 networkx>=2.2
 numpy>=1.15.4  # pandas
 openpyxl>=3.0
-pandas>=0.25.0
+pandas>=1.1.5
 python-dateutil>=2.8.1  # pandas
 pytz>=2019.2  # pandas
 pyyaml>=3.13
 requests>=2.21.1
 scikit-learn>=0.20.2
 scipy>=1.1.0
+seaborn>=0.9.0
 setuptools>=40.6.3
 splunk-sdk>=1.6.0
 statsmodels>=0.11.1

requirements.txt

@@ -33,7 +33,7 @@ msrest>=0.6.0
 # msrestazure>=0.6.0  # azure
 networkx>=2.2
 numpy>=1.15.4  # pandas
-pandas>=0.25.0
+pandas>=1.1.5
 python-dateutil>=2.8.1  # pandas
 pytz>=2019.2  # pandas
 pyyaml>=3.13

tools/misc/chk_pkgs.py

@@ -12,7 +12,7 @@
     "plotly>=3.10.0",
     "prettytable>=0.7.2",
     "matplotlib>=3.0.0",
-    "pandas>=0.23.4",
+    "pandas>=1.1.5",
     "adal>=1.2.1",
     "Pygments>=2.2.0",
     "seaborn>=0.9.0",