mantis_api_client 5.5.0__py3-none-any.whl

mantis_api_client/cli_parser/user_parser.py

@@ -0,0 +1,225 @@
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import sys
+from typing import Any
+
+from pydantic.json import pydantic_encoder
+
+from mantis_api_client import user_api
+from mantis_api_client.oidc import get_oidc_client
+from mantis_api_client.utils import colored
+
+
+# #
+# # 'user_list_handler' handler
+# #
+# def user_list_handler(args: Any) -> None:
+#     try:
+#         users = user_api.fetch_users()
+#     except Exception as e:
+#         print(colored(f"Error when fetching users: '{e}'", "red"))
+#         sys.exit(1)
+
+#     if args.json:
+#         print(json.dumps(users, default=pydantic_encoder))
+#         return
+
+#     print("[+] Organization users:")
+#     for user in users["data"]:
+#         print(f"  [+] {user['id']}: {user['given_name']} {user['last_name']}")
+
+
+#
+# 'user_self_handler' handler
+#
+def user_self_handler(args: Any) -> None:
+    try:
+        user = user_api.fetch_current_user()
+    except Exception as e:
+        print(colored(f"Error when fetching current user information: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(user, default=pydantic_encoder))
+        return
+
+    print("[+] Current user information:")
+    print(f"  [+] ID: {user['id']}")
+    print(f"  [+] Name: {user['given_name']} {user['last_name']}")
+    print(f"  [+] Email: {user['email']}")
+    print(f"  [+] Creation date: {user['creation_date']}")
+
+
+#
+# 'user_info_handler' handler
+#
+def user_info_handler(args: Any) -> None:
+    # Parameters
+    user_id = args.user_id
+
+    try:
+        user = user_api.fetch_user(user_id)
+    except Exception as e:
+        print(
+            colored(
+                f"Error when fetching information for user '{user_id}': '{e}'",
+                "red",
+                "on_white",
+            )
+        )
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(user, default=pydantic_encoder))
+        return
+
+    print("[+] User information:")
+    print(f"  [+] ID: {user['id']}")
+    print(f"  [+] Name: {user['given_name']} {user['last_name']}")
+    print(f"  [+] Email: {user['email']}")
+    print(f"  [+] Creation date: {user['creation_date']}")
+
+
+#
+# 'organization_info_handler' handler
+#
+def organization_info_handler(args: Any) -> None:
+    try:
+        # Retrieve current organization name
+        workspace_id = get_oidc_client().get_default_workspace(raise_exc=True)
+
+        # Retrieve associated organization id
+        if workspace_id is None:
+            print("Current user not linked to a workspace")
+            sys.exit(0)
+
+        # Retrieve organization info
+        workspaces_info = iter(user_api.fetch_current_workspaces())
+        organization_id = next(
+            workspace_info["organization_id"]
+            for workspace_info in workspaces_info
+            if workspace_info["id"] == workspace_id
+        )
+        organization_info = user_api.fetch_organization(organization_id)
+
+        # Retrieve plan info
+        plan_info = user_api.get_plan(organization_info["plan_id"])
+
+        # Retrieve plan limits
+        plan_limits = user_api.get_plan_limits(organization_info["plan_id"])
+
+        # Retrieve credits list
+        credits_list = user_api.get_credits_list(organization_id)
+
+        # Retrieve consumption history
+        consumption_history = user_api.get_consumption_history(organization_id)
+    except Exception as e:
+        print(
+            colored(
+                f"Error when fetching information for organization '{workspace_id}': '{e}'",
+                "red",
+                "on_white",
+            )
+        )
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(organization_info, default=pydantic_encoder))
+        print(json.dumps(plan_info, default=pydantic_encoder))
+        print(json.dumps(plan_limits, default=pydantic_encoder))
+        print(json.dumps(credits_list, default=pydantic_encoder))
+        print(json.dumps(consumption_history, default=pydantic_encoder))
+        return
+
+    print("[+] Organization information:")
+    print(f"  [+] ID: {organization_info['id']}")
+    print(f"  [+] Name: {organization_info['name']}")
+    print(f"  [+] Subscription start: {organization_info['subscription_start']}")
+    print(f"  [+] Subscription end: {organization_info['subscription_end']}")
+
+    print("[+] Plan information:")
+    print(f"  [+] ID: {plan_info['id']}")
+    print(f"  [+] Name: {plan_info['name']}")
+    print(f"  [+] Frequency: {plan_info['frequency']}")
+
+    print("[+] Plan limits:")
+    print(f"  [+] Max active labs: {plan_limits['nb_active_labs']}")
+    print(f"  [+] Max accounts: {plan_limits['nb_accounts']}")
+    print(f"  [+] Max workspaces: {plan_limits['nb_workspaces']}")
+    print(f"  [+] Max lab duration: {plan_limits['lab_max_duration']}")
+
+    print("[+] Credits:")
+    if "data" in credits_list:
+        for credit in credits_list["data"]:
+            print(f"  [+] ID: {credit['id']}")
+            print(f"    [+] Purpose: {credit['purpose']}")
+            print(
+                f"    [+] Period: {credit['creation_date']} to {credit['expiration_date']}"
+            )
+            print(f"    [+] Credit: {credit['remaining']}/{credit['total']}")
+    else:
+        print("  [+] No credits")
+
+
+def add_user_parser(
+    root_parser: argparse.ArgumentParser,
+    subparsers: Any,
+) -> None:
+    # --------------------
+    # --- User API options (users)
+    # --------------------
+
+    parser_user = subparsers.add_parser(
+        "user",
+        help="User API related commands",
+        formatter_class=root_parser.formatter_class,
+    )
+    subparsers_user = parser_user.add_subparsers()
+
+    # # 'user_list' command
+    # parser_user_list = subparsers_user.add_parser(
+    #     "list",
+    #     help="List all organization users",
+    #     formatter_class=root_parser.formatter_class,
+    # )
+    # parser_user_list.set_defaults(func=user_list_handler)
+    # parser_user_list.add_argument(
+    #     "--json", help="Return JSON result.", action="store_true"
+    # )
+
+    # 'user_self' command
+    parser_user_self = subparsers_user.add_parser(
+        "self",
+        help="Retrieve information for current user",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_user_self.set_defaults(func=user_self_handler)
+    parser_user_self.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'user_info' command
+    parser_user_info = subparsers_user.add_parser(
+        "info",
+        help="Retrieve information about a specific user",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_user_info.set_defaults(func=user_info_handler)
+    parser_user_info.add_argument("user_id", type=str, help="The user ID")
+    parser_user_info.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'orga_info' command
+    parser_user_info = subparsers_user.add_parser(
+        "organization",
+        help="Retrieve information about current organization",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_user_info.set_defaults(func=organization_info_handler)
+    parser_user_info.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    parser_user.set_defaults(func=lambda _: parser_user.print_help())

mantis_api_client/config.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2016-2021 AMOSSYS
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+from dataclasses import dataclass
+from dataclasses import field
+from pathlib import Path
+from typing import Optional
+
+from mantis_authz.config import AuthzConfig
+from omegaconf import OmegaConf
+from omegaconf import SI
+
+
+@dataclass
+class MantisOIDCConfig(AuthzConfig):
+    realm: str = "mantis"
+    client_id: str = "mantis-cli"
+    domain: Optional[str] = None
+    redirect_url: dict = field(
+        default_factory=lambda: {
+            "scheme": "http",
+            "host": "localhost",
+            "port": 9876,
+            "path": "/callback",
+        }
+    )
+    redirect_url_timeout: int = SI(
+        "${oc.decode:${oc.env:OIDC_REDIRECT_URL_TIMEOUT,60}}"
+    )
+    redirect_uri: str = (
+        "${oc.env:OIDC_REDIRECT_URI,${.redirect_url.scheme}://${.redirect_url.host}:${.redirect_url.port}${.redirect_url.path}}"
+    )
+
+
+@dataclass
+class MantisApiClientConfig:
+    dataset_api_url: str
+    dataset_web_url: str
+    scenario_api_url: str
+    user_api_url: str
+    cacert: Optional[Path] = SI("${oc.env:CR_CA_CERT,null}")
+    cert: Optional[Path] = SI("${oc.env:CR_CLIENT_CERT,null}")
+    key: Optional[Path] = SI("${oc.env:CR_CLIENT_KEY,null}")
+
+    config_file: Path = SI(
+        "${oc.env:MANTIS_CONFIG,${oc.env:XDG_CONFIG_HOME,${oc.env:HOME}/.config}/mantis/config.yml}"
+    )
+
+    # OIDC
+    oidc: MantisOIDCConfig = field(default_factory=MantisOIDCConfig)
+
+
+mantis_api_client_config = OmegaConf.structured(MantisApiClientConfig)

mantis_api_client/dataset_api.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2016-2021 AMOSSYS
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+from pathlib import Path
+from typing import Any
+from uuid import UUID
+
+import requests
+from mantis_dataset_model.dataset_model import Manifest
+
+from mantis_api_client.config import mantis_api_client_config
+from mantis_api_client.oidc import authorize
+
+# -------------------------------------------------------------------------- #
+# Internal helpers
+# -------------------------------------------------------------------------- #
+
+
+@authorize
+def _get(route: str, **kwargs: Any) -> requests.Response:
+    return requests.get(
+        f"{mantis_api_client_config.dataset_api_url}{route}",
+        verify=mantis_api_client_config.cacert,
+        cert=(mantis_api_client_config.cert, mantis_api_client_config.key),
+        **kwargs,
+    )
+
+
+@authorize
+def _post(route: str, **kwargs: Any) -> requests.Response:
+    timeout = kwargs.pop("timeout", 30)
+    return requests.post(
+        f"{mantis_api_client_config.dataset_api_url}{route}",
+        verify=mantis_api_client_config.cacert,
+        cert=(mantis_api_client_config.cert, mantis_api_client_config.key),
+        timeout=timeout,
+        **kwargs,
+    )
+
+
+@authorize
+def _put(route: str, **kwargs: Any) -> requests.Response:
+    return requests.put(
+        f"{mantis_api_client_config.dataset_api_url}{route}",
+        verify=mantis_api_client_config.cacert,
+        cert=(mantis_api_client_config.cert, mantis_api_client_config.key),
+        **kwargs,
+    )
+
+
+@authorize
+def _delete(route: str, **kwargs: Any) -> requests.Response:
+    return requests.delete(
+        f"{mantis_api_client_config.dataset_api_url}{route}",
+        verify=mantis_api_client_config.cacert,
+        cert=(mantis_api_client_config.cert, mantis_api_client_config.key),
+        **kwargs,
+    )
+
+
+def _handle_error(
+    result: requests.Response, context_error_msg: str
+) -> requests.Response:
+    if (
+        result.headers.get("content-type") == "application/json"
+        and "message" in result.json()
+    ):
+        error_msg = result.json()["message"]
+    else:
+        error_msg = result.text
+
+    raise Exception(
+        f"{context_error_msg}. "
+        f"Status code: '{result.status_code}'.\n"
+        f"Error message: '{error_msg}'."
+    )
+
+
+# -------------------------------------------------------------------------- #
+# Internal helpers
+# -------------------------------------------------------------------------- #
+
+
+def get_version() -> str:
+    """
+    Return dataset API version.
+
+    :return: The version inumber n a string
+    """
+    result = _get("/version")
+
+    if result.status_code != 200:
+        _handle_error(result, "Cannot retrieve dataset API version")
+
+    return result.json()
+
+
+# -------------------------------------------------------------------------- #
+# Dataset API
+# -------------------------------------------------------------------------- #
+
+
+def fetch_datasets() -> Any:
+    """
+    List all available datasets
+
+    :return: the list of manifests
+    """
+    result = _get("/")
+
+    if result.status_code != 200:
+        _handle_error(result, "Cannot retrieve datasets from dataset API")
+
+    return result.json()
+
+
+def fetch_dataset_by_uuid(partial_dataset_id: str) -> Any:
+    """
+    Get the full JSON manifest of a specific dataset
+
+    :param partial_dataset_id: UUID of the dataset to fetch, or prefix of it that uniquely identifies the dataset
+
+    If a full UUID is given, returns the identified dataset metadata if it
+    exists. If a partial UUID (prefix of an existing UUID) is given, and this
+    prefix uniquely identifies one and only one dataset, the latter's metadata
+    is returned.
+    """
+    result = _get(f"/{partial_dataset_id}")
+
+    if result.status_code != 200:
+        _handle_error(result, "Cannot retrieve dataset from dataset API")
+    else:
+        dataset_data = result.json()
+        dataset = Manifest(**dataset_data)
+
+    return dataset
+
+
+def fetch_dataset_resource(
+    partial_dataset_id: str, resource_type_str: str, resource_id: UUID
+) -> Any:
+    """
+    Get the description of a specific resource.
+
+    This function *does not* fetch the file(s) associated with one resource,
+    it merely returns the portion of the JSON manifest that describes that
+    particular resource. This JSON includes, in particular, the list
+    of files included in the resource, and the URLs to download them. It is
+    up to the user to manually fetch these URLs.
+
+    :param partial_dataset_id: UUID of the dataset to which the resource belongs, or a prefix of it that uniquely identifies it
+    :param resource_type_str: type of the resource to fetch. Must be a string among "log", "pcap", "memory_dump", "redteam_report", "life_report"
+    :param resource_id: UUID of the resource
+    :return: JSON structure describing the resource
+    """
+    result = _get(f"/{partial_dataset_id}/{resource_type_str}/{resource_id}")
+
+    if result.status_code != 200:
+        _handle_error(result, "Cannot retrieve dataset from dataset API")
+
+    return result.json()
+
+
+# def delete_dataset(
+#     dataset_id: str, force: bool = False
+# ) -> None:
+#     """
+#     Delete a specific dataset
+
+#     :param dataset_id: the full UUID of the dataset to delete
+#     :param force: (optional, default False) if True, forces the deletion of the dataset, even if there are validation errors in the manifest, or in the dataset contents. If the manifest is corrupted, remotely stored resources may not be deleted even though delete_remote_data was set to True.
+
+#     """
+
+#     raise NotImplementedError("This API method is not currently available")
+
+#     # result = _delete(
+#     #     f"/{dataset_id}",
+#     #     params={"force": force},
+#     # )
+
+#     # if result.status_code != 200:
+#     #     _handle_error(result, "Cannot delete dataset from dataset API")
+
+#     # return None
+
+
+# def delete_all_datasets(force: bool = False) -> Any:
+#     """
+#     Delete all datasets
+
+#     :param force: (optional, default False) if True, forces the deletion of the datasets, even if there are validation errors in the manifests, or in the datasets contents. If a manifest is corrupted, remotely stored resources for that dataset may not be deleted even though delete_remote_data was set to True.
+#     :return: a JSON structure specifying the datasets that were successfully deleted, and the errors encountered
+
+#     """
+
+#     raise NotImplementedError("This API method is not currently available")
+
+#     # result = _delete(
+#     #     "/all", params={"force": force}
+#     # )
+
+#     # if result.status_code != 200:
+#     #     _handle_error(result, "Cannot delete datasets from dataset API")
+
+#     # return result.json()
+
+
+def get_archive_infos_for_dataset(partial_dataset_id: str) -> Any:
+    """
+    Retrieve information about a zip archive for a partial_dataset_id.
+
+    :param partial_dataset_id: the partial or complete UUID of the dataset
+    :return: a JSON structure containing size, url and if the archive exists.
+    """
+    result = _get(f"/{partial_dataset_id}/zip")
+
+    if result.status_code != 200:
+        _handle_error(
+            result,
+            "Cannot retrieve dataset archive information from dataset API",
+        )
+
+    return result.json()
+
+
+def upload_dataset(dataset_path: Path, workspace_id: UUID | None) -> UUID:
+    # Uploading the dataset
+    params = {"filename": dataset_path.name, "owner": workspace_id}
+    with dataset_path.open("rb") as f:
+        result = _post("/upload", params=params, data=f, timeout=None)
+
+    if result.status_code != 200:
+        _handle_error(
+            result,
+            "Cannot retrieve dataset archive information from dataset API",
+        )
+    dataset_id = UUID(result.json()["dataset_id"])
+
+    return dataset_id
+
+
+def upload_dataset_status(dataset_id: UUID) -> Any:
+    # Getting information about the Dataset
+    state = _get(f"/upload/{dataset_id}/status")
+    state = state.json()
+    return state

mantis_api_client/exceptions.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2016-2021 AMOSSYS
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+
+
+class ConfigException(Exception):
+    pass