mantis_api_client 5.5.0__py3-none-any.whl

mantis_api_client/cli_parser/scenario_parser.py

@@ -0,0 +1,311 @@
+# -*- 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
+
+
+#
+# 'scenario_list_handler' handler
+#
+def scenario_list_handler(args: Any) -> None:
+    try:
+        scenarios = sorted(scenario_api.fetch_scenarios(), key=lambda x: x.name)
+    except Exception as e:
+        print(colored(f"Error when fetching scenarios: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(scenarios, default=pydantic_encoder))
+        return
+
+    print(f"[+] Available scenarios ({len(scenarios)}):")
+    for scenario in scenarios:
+        print(f"  [+] \033[1m{scenario.name}\033[0m")
+        print(f"         {scenario.description}")
+
+
+#
+# 'scenario_info_handler' handler
+#
+def scenario_info_handler(args: Any) -> None:
+    # Parameters
+    scenario_name = args.scenario_name
+
+    try:
+        scenario = scenario_api.fetch_scenario_by_name(scenario_name)
+    except Exception as e:
+        print(
+            colored(
+                f"Error when fetching scenario {scenario_name}: '{e}'",
+                "red",
+                "on_white",
+            )
+        )
+        sys.exit(1)
+
+    if args.json:
+        print(scenario.json())
+        return
+
+    print("[+] Scenario information:")
+    print(f"  [+] \033[1mName\033[0m: {scenario.name}")
+    print("  [+] \033[1mKeywords\033[0m: ", end="")
+    print(", ".join(scenario.keywords))
+    print(f"  [+] \033[1mDescription\033[0m: {scenario.description}")
+    if len(scenario.long_description) > 0:
+        print("  [+] \033[1mLong description\033[0m:")
+        for ld in scenario.long_description:
+            print(f"        - {ld}")
+    print(f"  [+] \033[1mLearning context\033[0m: {scenario.learning_context}")
+    print("  [+] \033[1mAvailable scenario profiles\033[0m:")
+    for scenario_profile in scenario.scenario_profiles:
+        target = f", target: {scenario_profile.compromission.target_name}"
+        print(
+            f"""       [+] {scenario_profile.name} (Topology: {scenario_profile.topology_name}){target}"""
+        )
+    print("  [+] \033[1mUnit attacks available\033[0m:")
+    for attack in scenario.attacks:
+        print(f"       - {attack}")
+
+
+def _scenario_create_or_run_lab(args: Any, do_run: bool = True):
+    # Parameters
+    scenario_name = args.scenario_name
+    lab_config_path = args.lab_config_path
+    scenario_profile = args.scenario_profile
+
+    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 scenario
+    try:
+        scenario = scenario_api.fetch_scenario_by_name(scenario_name)
+        print(f"[+] Going to execute scenario: {scenario.name}")
+
+        if scenario_profile is None:
+            print(
+                "Needed argument --scenario_profile in order to choose the unit scenario to run"
+            )
+            print("Available unit scenarios:")
+            for available_scenario_profile in scenario.scenario_profiles:
+                print(f"  [+] {available_scenario_profile.name}")
+            sys.exit(-1)
+
+        if not any(scenario_profile == c.name for c in scenario.scenario_profiles):
+            print(
+                colored(
+                    f"Error '{scenario_profile}' not supported for this scenario.",
+                    "red",
+                    "on_white",
+                )
+            )
+            sys.exit(-1)
+
+        if do_run:
+            lab_id = scenario_api.run_lab_scenario(
+                scenario,
+                scenario_profile,
+                lab_config,
+                workspace_id,
+                args.public,
+            )
+        else:
+            lab_id = scenario_api.create_lab_scenario(
+                scenario,
+                scenario_profile,
+                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 scenario {scenario_name}: '{e}'", "red"))
+        sys.exit(1)
+    finally:
+        if do_run and args.destroy_after_scenario:
+            print("[+] Stopping lab...")
+            scenario_api.stop_lab(lab_id)
+
+
+#
+# 'scenario_create_lab_handler' handler
+#
+def scenario_create_lab_handler(args: Any) -> None:
+    _scenario_create_or_run_lab(args, do_run=False)
+
+
+#
+# 'scenario_run_lab_handler' handler
+#
+def scenario_run_lab_handler(args: Any) -> None:
+    _scenario_create_or_run_lab(args, do_run=True)
+
+
+def add_scenario_parser(
+    root_parser: argparse.ArgumentParser,
+    subparsers: Any,
+) -> None:
+    # --------------------
+    # --- Scenario API options (scenario)
+    # --------------------
+
+    parser_scenario = subparsers.add_parser(
+        "scenario",
+        help="Scenario API related commands (scenario)",
+        formatter_class=root_parser.formatter_class,
+    )
+    subparsers_scenario = parser_scenario.add_subparsers()
+
+    # 'scenario_list' command
+    parser_scenario_list = subparsers_scenario.add_parser(
+        "list",
+        help="List all available scenarios",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_scenario_list.set_defaults(func=scenario_list_handler)
+    parser_scenario_list.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'scenario_info' command
+    parser_scenario_info = subparsers_scenario.add_parser(
+        "info",
+        help="Get information about a scenario",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_scenario_info.set_defaults(func=scenario_info_handler)
+    parser_scenario_info.add_argument(
+        "scenario_name", type=str, help="The scenario name"
+    )
+    parser_scenario_info.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'scenario_create_lab' command
+    parser_scenario_create_lab = subparsers_scenario.add_parser(
+        "create",
+        help="Create a lab with a specific scenario",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_scenario_create_lab.set_defaults(func=scenario_create_lab_handler)
+    parser_scenario_create_lab.add_argument(
+        "scenario_name", type=str, help="The scenario name"
+    )
+    parser_scenario_create_lab.add_argument(
+        "--workspace",
+        dest="workspace_id",
+        help="The group ID that have ownership on lab",
+    )
+    parser_scenario_create_lab.add_argument(
+        "--destroy",
+        action="store_true",
+        dest="destroy_after_scenario",
+        help="Do not keep the lab up after scenario execution (False by default)",
+    )
+    parser_scenario_create_lab.add_argument(
+        "--lab_config",
+        "-lc",
+        action="store",
+        required=False,
+        dest="lab_config_path",
+        help="Input path of a YAML configuration for the scenario to run",
+    )
+    parser_scenario_create_lab.add_argument(
+        "--scenario_profile",
+        "-sp",
+        required=False,
+        dest="scenario_profile",
+        help="Allows to define the scenario config to run",
+    )
+    parser_scenario_create_lab.add_argument(
+        "--public",
+        action="store_true",
+        help="Create a public link to access this lab",
+    )
+
+    # 'scenario_run_lab' command
+    parser_scenario_run_lab = subparsers_scenario.add_parser(
+        "run",
+        help="Create and run a lab with a specific scenario",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_scenario_run_lab.set_defaults(func=scenario_run_lab_handler)
+    parser_scenario_run_lab.add_argument(
+        "scenario_name", type=str, help="The scenario name"
+    )
+    parser_scenario_run_lab.add_argument(
+        "--workspace",
+        dest="workspace_id",
+        help="The group ID that have ownership on lab",
+    )
+    parser_scenario_run_lab.add_argument(
+        "--topology",
+        action="store",
+        required=False,
+        dest="topology_file",
+        help="Input path of a YAML topology file to override the default one",
+    )
+    parser_scenario_run_lab.add_argument(
+        "--destroy",
+        action="store_true",
+        dest="destroy_after_scenario",
+        help="Do not keep the lab up after scenario execution (False by default)",
+    )
+    parser_scenario_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_scenario_run_lab.add_argument(
+        "--scenario_profile",
+        "-sp",
+        required=False,
+        dest="scenario_profile",
+        help="Allows to define the scenario config to run",
+    )
+    parser_scenario_run_lab.add_argument(
+        "--public",
+        action="store_true",
+        help="Create a public link to access this lab",
+    )
+
+    parser_scenario.set_defaults(func=lambda _: parser_scenario.print_help())

mantis_api_client/cli_parser/signature_parser.py

@@ -0,0 +1,147 @@
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import sys
+from typing import Any
+from typing import List
+
+from mantis_scenario_model.signature_model import Signature
+from pydantic.json import pydantic_encoder
+
+from mantis_api_client import scenario_api
+from mantis_api_client.utils import colored
+
+
+#
+# 'signature_list_handler' handler
+#
+def signature_list_handler(args: Any) -> None:
+    try:
+        signatures = scenario_api.fetch_signatures()
+    except Exception as e:
+        print(colored(f"Error when fetching signatures: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(signatures, default=pydantic_encoder))
+        return
+    for key, signature in sorted(
+        signatures.items(), key=lambda kv: kv[1].attack_reference_id  # (clé, Signature)
+    ):
+        print(f"[+] \033[1m{signature.attack_reference_id}\033[0m")
+        print(f"  [+] id: {key}")
+        implems = ", ".join(
+            imp.implementation_type.value for imp in signature.implementations
+        )
+        print(f"  [+] available implementations: {implems}")
+
+
+def add_signature_parser(
+    root_parser: argparse.ArgumentParser,
+    subparsers: Any,
+) -> None:
+    # -------------------
+    # --- Scenario API options (signatures)
+    # -------------------
+
+    parser_signature = subparsers.add_parser(
+        "signature",
+        help="Scenario API related commands (signature)",
+        formatter_class=root_parser.formatter_class,
+    )
+    subparsers_signature = parser_signature.add_subparsers()
+
+    # 'signature_list' command
+    parser_signature_list = subparsers_signature.add_parser(
+        "list",
+        help="List all available signatures",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_signature_list.set_defaults(func=signature_list_handler)
+    parser_signature_list.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    parser_signature_info = subparsers_signature.add_parser(
+        "info",
+        help="Get information about a signature",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_signature_info.set_defaults(func=signature_info_handler)
+    parser_signature_info.add_argument(
+        "signature_id", type=str, help="The signature id"
+    )
+    parser_signature_info.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    group = parser_signature_info.add_mutually_exclusive_group()
+
+    # Ajoutez les arguments au groupe
+    group.add_argument(
+        "--attack",
+        help="The signature_id is a reference to an attack name. Cannot be used with --scenario.",
+        action="store_true",
+    )
+    group.add_argument(
+        "--scenario",
+        help="The signature_id is a reference to a scenario name. Cannot be used with --attack.",
+        action="store_true",
+    )
+
+    parser_signature.set_defaults(func=lambda _: parser_signature.print_help())
+
+
+def print_signatures_of_scenario(signatures: List[Signature]) -> None:
+    for sig in signatures:
+        print_signature(sig)
+
+
+def print_signature(signature) -> None:
+
+    print(f"[+] \033[1m{signature.attack_reference_id}\033[0m")
+    print("  [+] available implementations:")
+    for i in signature.implementations:
+        print(f"  [>] {i.implementation_type.value}")
+        for k in i.signatures:
+            print(f"    [+] Format \033[1m{k.signature_type.value}\033[0m")
+            print(f"    {k.signature}")
+
+
+def signature_info_handler(args: Any) -> None:
+    # Parameters
+    signature_id = args.signature_id
+    signatures = []
+    signature = None
+
+    try:
+        if args.attack is True:
+            # It's an attack ID
+            signature = scenario_api.fetch_signature_by_attack_id(signature_id)
+        elif args.scenario is True:
+            # It's a scenario ID
+            signatures = scenario_api.fetch_signatures_by_scenario_id(signature_id)
+        else:
+            # It's a signature ID
+            signature = scenario_api.fetch_signature_by_signature_id(signature_id)
+    except Exception as e:
+        print(colored(f"Error when fetching signature {signature_id}: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json and args.scenario is False and signature is not None:
+        print(signature.json())
+        return
+    elif args.json and args.scenario is True and signatures is not None:
+        json_data = [sig.json() for sig in signatures]
+
+        # Now json_data is a list of JSON strings, but if you want it in the final JSON format, use json.dumps
+        json_str = json.dumps(json_data)
+
+        # Print the final JSON string
+        print(json_str)
+        return
+
+    if args.scenario is True:
+        print_signatures_of_scenario(signatures)
+    else:
+        print_signature(signature)

mantis_api_client/cli_parser/topology_parser.py

@@ -0,0 +1,255 @@
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import sys
+from typing import Any
+
+from mantis_scenario_model.lab_config_model import LabConfig
+from mantis_scenario_model.node import VirtualMachine
+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
+
+
+#
+# 'topology_list_handler' handler
+#
+def topology_list_handler(args: Any) -> None:
+    try:
+        topologies = sorted(scenario_api.fetch_topologies(), key=lambda x: x.name)
+    except Exception as e:
+        print(colored(f"Error when fetching topologies: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(topologies, default=pydantic_encoder))
+        return
+
+    print("[+] Available topologies:")
+    for topology in topologies:
+        print(f"  [+] {topology.name}")
+
+
+#
+# 'topology_info_handler' handler
+#
+def topology_info_handler(args: Any) -> None:
+    # Parameters
+    topology_name = args.topology_name
+
+    try:
+        topology = scenario_api.fetch_topology_by_name(topology_name)
+    except Exception as e:
+        print(colored(f"Error when fetching topology {topology_name}: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json:
+        print(topology.json())
+        return
+
+    print("[+] Topology information:")
+    print(f"  [+] \033[1mName\033[0m: {topology.name}")
+    print("  [+] \033[1mNodes\033[0m:")
+    for node in topology.nodes:
+        if type(node) is VirtualMachine:
+            print(f"       [+] {node.type}: {node.name} ({node.basebox_id})")
+        else:
+            print(f"       [+] {node.type}: {node.name}")
+    print("  [+] \033[1mLinks\033[0m:")
+    for link in topology.links:
+        print(
+            f"        [+] {link.switch.name}({link.switch.type}) -- {link.params.ip} -- {link.node.name}({link.node.type})"
+        )
+
+
+def _topology_create_or_run_lab(args: Any, do_run: bool = True):
+    # Parameters
+    topology_name = args.topology_name
+    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 topology
+    try:
+        topology = scenario_api.fetch_topology_by_name(topology_name)
+        print(f"[+] Going to execute topology: {topology.name}")
+
+        if do_run:
+            lab_id = scenario_api.run_lab_topology(
+                topology,
+                lab_config,
+                workspace_id,
+                args.public,
+            )
+        else:
+            lab_id = scenario_api.create_lab_topology(
+                topology,
+                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 topology {topology_name}: '{e}'", "red"))
+        sys.exit(1)
+    finally:
+        if do_run and args.destroy_after_scenario:
+            print("[+] Stopping lab...")
+            scenario_api.stop_lab(lab_id)
+
+
+#
+# 'topology_create_lab_handler' handler
+#
+def topology_create_lab_handler(args: Any) -> None:
+    _topology_create_or_run_lab(args, do_run=False)
+
+
+#
+# 'topology_run_lab_handler' handler
+#
+def topology_run_lab_handler(args: Any) -> None:
+    _topology_create_or_run_lab(args, do_run=True)
+
+
+def add_topology_parser(
+    root_parser: argparse.ArgumentParser,
+    subparsers: Any,
+) -> None:
+    # -------------------
+    # --- Scenario API options (topologies)
+    # -------------------
+
+    parser_topology = subparsers.add_parser(
+        "topology",
+        help="Scenario API related commands (topology)",
+        formatter_class=root_parser.formatter_class,
+    )
+    subparsers_topology = parser_topology.add_subparsers()
+
+    # 'topology_list' command
+    parser_topology_list = subparsers_topology.add_parser(
+        "list",
+        help="List all available topologies",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_topology_list.set_defaults(func=topology_list_handler)
+    parser_topology_list.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'topology_info' command
+    parser_topology_info = subparsers_topology.add_parser(
+        "info",
+        help="Get information about a topology",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_topology_info.set_defaults(func=topology_info_handler)
+    parser_topology_info.add_argument(
+        "topology_name", type=str, help="The topology name"
+    )
+    parser_topology_info.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'topology_create_lab' command
+    parser_topology_create_lab = subparsers_topology.add_parser(
+        "create",
+        help="Create a lab with a specific topology",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_topology_create_lab.set_defaults(func=topology_create_lab_handler)
+    parser_topology_create_lab.add_argument(
+        "topology_name", type=str, help="The topology name"
+    )
+    parser_topology_create_lab.add_argument(
+        "--destroy",
+        action="store_true",
+        dest="destroy_after_scenario",
+        help="Do not keep the lab up after topology execution (False by default)",
+    )
+    parser_topology_create_lab.add_argument(
+        "--workspace",
+        dest="workspace_id",
+        help="The group ID that have ownership on lab",
+    )
+    parser_topology_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_topology_create_lab.add_argument(
+        "--public",
+        action="store_true",
+        help="Create a public link to access this lab",
+    )
+
+    # 'topology_run_lab' command
+    parser_topology_run_lab = subparsers_topology.add_parser(
+        "run",
+        help="Create and run a lab with a specific topology",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_topology_run_lab.set_defaults(func=topology_run_lab_handler)
+    parser_topology_run_lab.add_argument(
+        "topology_name", type=str, help="The topology name"
+    )
+    parser_topology_run_lab.add_argument(
+        "--destroy",
+        action="store_true",
+        dest="destroy_after_scenario",
+        help="Do not keep the lab up after topology execution (False by default)",
+    )
+    parser_topology_run_lab.add_argument(
+        "--workspace",
+        dest="workspace_id",
+        help="The group ID that have ownership on lab",
+    )
+    parser_topology_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_topology_run_lab.add_argument(
+        "--public",
+        action="store_true",
+        help="Create a public link to access this lab",
+    )
+
+    parser_topology.set_defaults(func=lambda _: parser_topology.print_help())