mantis_api_client 5.5.0__py3-none-any.whl

mantis_api_client/cli_parser/basebox_parser.py

@@ -0,0 +1,258 @@
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import sys
+from typing import Any
+
+from mantis_scenario_model.lab_config_model import LabConfig
+from pydantic.json import pydantic_encoder
+from ruamel.yaml import YAML
+
+from mantis_api_client import scenario_api
+from mantis_api_client.oidc import get_oidc_client
+from mantis_api_client.utils import colored
+from mantis_api_client.utils import wait_lab
+
+
+#
+# 'basebox_list_handler' handler
+def basebox_list_handler(args: Any) -> None:
+    try:
+        baseboxes = sorted(scenario_api.fetch_baseboxes(), key=lambda x: x.id)
+    except Exception as e:
+        print(colored(f"Error when fetching baseboxes: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(baseboxes, default=pydantic_encoder))
+        return
+
+    print("[+] Available baseboxes:")
+    for basebox in baseboxes:
+        print(f"  [+] {basebox.id}")
+
+
+#
+# 'basebox_info_handler' handler
+#
+def basebox_info_handler(args: Any) -> None:
+    # Parameters
+    basebox_id = args.basebox_id
+
+    try:
+        basebox = scenario_api.fetch_basebox_by_id(basebox_id)
+    except Exception as e:
+        print(colored(f"Error when fetching basebox {basebox_id}: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json:
+        print(basebox.json())
+        return
+
+    print("[+] Basebox information:")
+    print(f"  [+] \033[1mId\033[0m: {basebox.id}")
+    print(f"  [+] \033[1mDescription\033[0m: {basebox.description}")
+    print(f"  [+] \033[1mOS\033[0m: {basebox.operating_system} ({basebox.system_type})")
+    print(f"  [+] \033[1mLanguage\033[0m: {basebox.language}")
+    print(f"  [+] \033[1mInstallation date\033[0m: {basebox.installation_date}")
+    print("  [+] \033[1mCredentials\033[0m:")
+    if basebox.username != "" and basebox.username:
+        print(f"       [+] User: {basebox.username}:{basebox.password}")
+    if basebox.admin_username != "" and basebox.admin_username:
+        print(
+            f"       [+] Administrator: {basebox.admin_username}:{basebox.admin_password}"
+        )
+    print("  [+] \033[1mCPE\033[0m:")
+    for cpe in basebox.cpes:
+        print(f"        [+] {cpe}")
+    if basebox.changelog:
+        print("  [+] \033[1mChangelog\033[0m:")
+        for changelog in basebox.changelog:
+            if changelog.get("date", "") != "":
+                print(f"""        [+] {changelog.get("date")}:""")
+                for info in changelog.get("info", "").splitlines():
+                    print(f"""              {info}""")
+
+
+def _basebox_create_or_run_lab(args: Any, do_run: bool = True):
+    # Parameters
+    basebox_id = args.basebox_id
+    lab_config_path = args.lab_config_path
+
+    if not args.workspace_id:
+        try:
+            workspace_id = get_oidc_client().get_default_workspace(raise_exc=True)
+        except Exception as e:
+            print(colored(f"Error when fetching attacks: '{e}'", "red"))
+            sys.exit(1)
+    else:
+        workspace_id = args.workspace_id
+
+    # Retrieve associated group id
+    if workspace_id is None:
+        print(colored("You have to specify a workspace with --workspace", "yellow"))
+        sys.exit(1)
+
+    # Manage lab configuration
+    if lab_config_path is None:
+        lab_config_dict = {}
+    else:
+        with open(lab_config_path, "r") as fd:
+            yaml_content = fd.read()
+        loader = YAML(typ="rt")
+        lab_config_dict = loader.load(yaml_content)
+    lab_config = LabConfig(**lab_config_dict)
+
+    # Launch basebox
+    try:
+        print(f"[+] Going to execute basebox: {basebox_id}")
+
+        if do_run:
+            lab_id = scenario_api.run_lab_basebox(
+                basebox_id,
+                lab_config,
+                workspace_id,
+                args.public,
+            )
+        else:
+            lab_id = scenario_api.create_lab_basebox(
+                basebox_id,
+                lab_config,
+                workspace_id,
+                args.public,
+            )
+
+        print(f"[+] Lab ID: {lab_id}")
+
+        if do_run:
+            wait_lab(lab_id)
+
+    except Exception as e:
+        print(colored(f"Error when running basebox {basebox_id}: '{e}'", "red"))
+        sys.exit(1)
+    finally:
+        if do_run and args.destroy_after_scenario:
+            print("[+] Stopping lab...")
+            scenario_api.stop_lab(lab_id)
+
+
+#
+# 'basebox_create_lab_handler' handler
+#
+def basebox_create_lab_handler(args: Any) -> None:
+    _basebox_create_or_run_lab(args, do_run=False)
+
+
+#
+# 'basebox_run_lab_handler' handler
+#
+def basebox_run_lab_handler(args: Any) -> None:
+    _basebox_create_or_run_lab(args, do_run=True)
+
+
+def add_basebox_parser(
+    root_parser: argparse.ArgumentParser,
+    subparsers: Any,
+) -> None:
+    # -------------------
+    # --- Scenario API options (baseboxes)
+    # -------------------
+
+    parser_basebox = subparsers.add_parser(
+        "basebox",
+        help="Scenario API related commands (basebox)",
+        formatter_class=root_parser.formatter_class,
+    )
+    subparsers_basebox = parser_basebox.add_subparsers()
+
+    # 'basebox_list' command
+    parser_basebox_list = subparsers_basebox.add_parser(
+        "list",
+        help="List all available baseboxes",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_basebox_list.set_defaults(func=basebox_list_handler)
+    parser_basebox_list.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'basebox_info' command
+    parser_basebox_info = subparsers_basebox.add_parser(
+        "info",
+        help="Get information about a basebox",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_basebox_info.set_defaults(func=basebox_info_handler)
+    parser_basebox_info.add_argument("basebox_id", type=str, help="The basebox ID")
+    parser_basebox_info.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'basebox_create_lab' command
+    parser_basebox_create_lab = subparsers_basebox.add_parser(
+        "create",
+        help="Create a lab with a specific basebox",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_basebox_create_lab.set_defaults(func=basebox_create_lab_handler)
+    parser_basebox_create_lab.add_argument(
+        "basebox_id", type=str, help="The basebox ID"
+    )
+    parser_basebox_create_lab.add_argument(
+        "--destroy",
+        action="store_true",
+        dest="destroy_after_scenario",
+        help="Do not keep the lab up after basebox execution (False by default)",
+    )
+    parser_basebox_create_lab.add_argument(
+        "--workspace",
+        dest="workspace_id",
+        help="The group ID that have ownership on lab",
+    )
+    parser_basebox_create_lab.add_argument(
+        "--lab_config",
+        action="store",
+        required=False,
+        dest="lab_config_path",
+        help="Input path of a YAML configuration for the scenario to run",
+    )
+    parser_basebox_create_lab.add_argument(
+        "--public",
+        action="store_true",
+        help="Create a public link to access this lab",
+    )
+
+    # 'basebox_run_lab' command
+    parser_basebox_run_lab = subparsers_basebox.add_parser(
+        "run",
+        help="Create and run a lab with a specific basebox",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_basebox_run_lab.set_defaults(func=basebox_run_lab_handler)
+    parser_basebox_run_lab.add_argument("basebox_id", type=str, help="The basebox ID")
+    parser_basebox_run_lab.add_argument(
+        "--destroy",
+        action="store_true",
+        dest="destroy_after_scenario",
+        help="Do not keep the lab up after basebox execution (False by default)",
+    )
+    parser_basebox_run_lab.add_argument(
+        "--workspace",
+        dest="workspace_id",
+        help="The group ID that have ownership on lab",
+    )
+    parser_basebox_run_lab.add_argument(
+        "--lab_config",
+        "-lc",
+        action="store",
+        required=False,
+        dest="lab_config_path",
+        help="Input path of a YAML configuration for the scenario run",
+    )
+    parser_basebox_run_lab.add_argument(
+        "--public",
+        action="store_true",
+        help="Create a public link to access this lab",
+    )
+
+    parser_basebox.set_defaults(func=lambda _: parser_basebox.print_help())

mantis_api_client/cli_parser/dataset_parser.py

@@ -0,0 +1,200 @@
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import time
+from pathlib import Path
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import NoReturn
+from uuid import UUID
+
+from mantis_dataset_model.dataset_model import PartialManifest
+from pydantic.json import pydantic_encoder
+from rich.console import Console
+
+from mantis_api_client import dataset_api
+from mantis_api_client.exceptions import ConfigException
+from mantis_api_client.oidc import get_oidc_client
+from mantis_api_client.utils import colored
+
+
+print = Console().print
+
+
+#
+# 'datasets_list_handler' handler
+#
+def datasets_list_handler(args: Any) -> None:
+    try:
+        datasets: Dict[str, List[Dict[UUID, PartialManifest]]] = (
+            dataset_api.fetch_datasets()
+        )
+    except ConfigException as e:
+        print("{} You should run `mantis init` again to login to M&ntis SSO.", e)
+        exit(1)
+    except Exception as e:
+        print(colored(f"Error when listing datasets: '{e}'", "red"))
+        exit(1)
+
+    if args.json:
+        print(json.dumps(datasets, default=pydantic_encoder))
+        return
+
+    print("[+] Available datasets:")
+    for dataset in datasets:
+        if len(datasets[dataset]) > 0:
+            print(f"  [+] Datasets of '{dataset}': ")
+            for d in datasets[dataset]:
+                for id in d:
+                    print(
+                        f"""    [+] {id} - {d[id]['name']} ({d[id]["date_dataset_created"]})"""
+                    )
+
+
+#
+# 'datasets_info_handler' handler
+#
+def datasets_info_handler(args: Any) -> None:
+    # Parameters
+    dataset_id = args.dataset_id
+
+    try:
+        dataset = dataset_api.fetch_dataset_by_uuid(dataset_id)
+    except Exception as e:
+        print(colored(f"Error when fetching dataset: '{e}'", "red"))
+        exit(1)
+
+    if args.json:
+        print(dataset.json())
+        return
+
+    print("[+] Dataset information: ")
+    print(f"  [+] \033[1mName\033[0m: {dataset.name}")
+    print(f"  [+] \033[1mDescription\033[0m: {dataset.description}")
+    print(f"  [+] \033[1mCreation date\033[0m: {dataset.date_dataset_created}")
+    if dataset.scenario:  # scenario is optional
+        if (
+            "timestamps" in dataset.scenario
+            and "duration" in dataset.scenario["timestamps"]
+        ):
+            print(f"  [+] Duration: {dataset.scenario['timestamps']['duration']}")
+
+
+#
+# 'datasets_push_handler' handler
+#
+def datasets_push_handler(args: Any) -> None | NoReturn:
+    """
+    Handler to upload dataset to remote
+
+    An attempt is performed to send the dataset and make it available to workspace WORKSPACE (--workspace).
+    """
+    # Parameters
+    dataset_path: Path = args.dataset_path
+    workspace_id: UUID | None = args.workspace
+    uploaded: bool = False
+    error_msg: str = ""
+
+    if not args.public:
+        _workspace_id: str | None = None
+        if not workspace_id:
+            _workspace_id = get_oidc_client().get_default_workspace()
+        if not _workspace_id:
+            print("[b red]No default workspace found. [/b red]Run `mantis login`.")
+            exit(1)
+        workspace_id = UUID(_workspace_id)
+
+    if dataset_path.exists() is False:
+        print(f"[b red]The dataset path '{dataset_path}' does not exist.")
+        exit(1)
+
+    if args.public:
+        print("[+] Uploading dataset accessible to anyone")
+    else:
+        print(f"[+] Uploading dataset accessible to workspace: {workspace_id}")
+
+    # Uploading the dataset
+    dataset_id = dataset_api.upload_dataset(dataset_path, workspace_id)
+    print(f"  [+] Dataset id: {dataset_id}")
+
+    # Checking the status of the task
+    timeout = 300  # seconds
+    start = time.time()
+    finished = False
+    while time.time() - start < timeout and finished is False:
+        res = dataset_api.upload_dataset_status(dataset_id)
+        try:
+            state = res["status"]
+            error_msg = res["message"]
+        except KeyError:
+            error_msg = res.get("detail", "Unknown error")
+            break
+        if state in ["FINISHED", "FINISHED_ERROR"]:
+            finished = True
+            uploaded = state == "FINISHED"
+        else:
+            print("  [+] Dataset verification in progress")
+            time.sleep(1)
+
+    if uploaded is False:
+        print(colored(f"  [-] The dataset could not be uploaded: {error_msg}", "red"))
+    else:
+        print("  [+] Successfully uploaded the dataset.")
+    return None
+
+
+def add_dataset_parser(root_parser: argparse.ArgumentParser, subparsers: Any) -> None:
+
+    # -------------------
+    # --- Dataset API options
+    # -------------------
+
+    parser_dataset = subparsers.add_parser(
+        "dataset",
+        help="Dataset API related commands",
+        formatter_class=root_parser.formatter_class,
+    )
+    subparsers_dataset = parser_dataset.add_subparsers()
+
+    # 'datasets_list' command
+    parser_dataset_list = subparsers_dataset.add_parser(
+        "list",
+        help="List all available datasets",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_dataset_list.set_defaults(func=datasets_list_handler)
+    parser_dataset_list.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'datasets_info' command
+    parser_datasets_info = subparsers_dataset.add_parser(
+        "info",
+        help="Get information about a dataset",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_datasets_info.set_defaults(func=datasets_info_handler)
+    parser_datasets_info.add_argument("dataset_id", type=str, help="The dataset id")
+    parser_datasets_info.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    parser_datasets_push = subparsers_dataset.add_parser(
+        "push",
+        help="Upload a local dataset .zip file to the remote",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_datasets_push.set_defaults(func=datasets_push_handler)
+    pubpriv_group = parser_datasets_push.add_mutually_exclusive_group()
+    pubpriv_group.add_argument(
+        "--workspace",
+        type=UUID,
+        metavar="WORKSPACE_UUID",
+        help="Workspace to which this dataset belongs",
+    )
+    pubpriv_group.add_argument("--public", action="store_true")
+    parser_datasets_push.add_argument(
+        "dataset_path", type=Path, help="Path to the dataset .zip file to upload"
+    )
+    parser_dataset.set_defaults(func=lambda _: parser_dataset.print_help())