mantis_api_client 4.9.0__14-py3-none-any.whl

mantis_api_client/__init__.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+import os
+import threading
+from pathlib import Path
+
+# Component version
+__version__ = "4.8.0"
+
+# Remove incompatible env variable
+os.environ.pop("USE_PERMISSIONS", None)
+
+# Get component full version from file generated at build time
+current_file_dir = Path(__file__).resolve().parent
+fullversion_file = Path(current_file_dir, "fullversion.txt")
+if os.path.isfile(fullversion_file):
+    __fullversion__ = open(fullversion_file, "r").read().strip()
+else:
+    __fullversion__ = __version__
+
+shutil_make_archive_lock = threading.Lock()

mantis_api_client/cli_parser/account_parser.py

@@ -0,0 +1,347 @@
+# -*- coding: utf-8 -*-
+import argparse
+import getpass
+import os
+import subprocess
+import sys
+from http.server import BaseHTTPRequestHandler
+from http.server import HTTPServer
+from pathlib import Path
+from threading import Thread
+from typing import Any
+from typing import List
+from urllib.parse import parse_qs
+from urllib.parse import urlparse
+
+import requests
+from jose import jwt
+from rich.console import Console
+from rich.prompt import Prompt
+
+import mantis_api_client.dataset_api as dataset_api
+import mantis_api_client.oidc
+import mantis_api_client.scenario_api as scenario_api
+from mantis_api_client.config import mantis_api_client_config
+from mantis_api_client.oidc import oidc_client
+from mantis_api_client.utils import colored
+
+
+class Version:
+    def __init__(self, str_vers: str) -> None:
+        try:
+            self.major, self.minor, self.patch = str_vers.split(".")
+        except Exception as e:
+            raise Exception(
+                "Bad version format for '{}': 'X.Y.Z' expected. Error: {}".format(
+                    str_vers, e
+                )
+            )
+
+
+#
+# 'status' related functions
+#
+def status_handler(args: Any) -> None:  # noqa: C901
+    """Get platform status."""
+
+    exit_code = 0
+
+    client_version = mantis_api_client.__version__
+    client_vers = Version(str_vers=client_version)
+    client_fullversion = mantis_api_client.__fullversion__
+    print(
+        f"[+] mantis_api_client version: {client_version} ({client_fullversion})".format(
+            client_version
+        )
+    )
+
+    active_profile_domain = oidc_client.get_active_profile_domain(raise_exc=False)
+    if active_profile_domain:
+        print(f"[+] Authenticated to {active_profile_domain}")
+    else:
+        print(
+            colored(
+                "[+] Not authenticated, you need to execute 'mantis account login'",
+                "red",
+            )
+        )
+        sys.exit(1)
+    print("[+] APIs status")
+
+    # Dataset API
+    print("  [+] Dataset API")
+    print("    [+] address: {}".format(mantis_api_client_config.dataset_api_url))
+    try:
+        dataset_api_version = dataset_api.get_version()
+        dataset_vers = Version(str_vers=dataset_api_version)
+    except requests.exceptions.ConnectionError:
+        exit_code = 1
+        print("    [-] API status: " + colored("not running !", "white", "on_red"))
+    else:
+        print("    [+] API status: " + colored("OK", "grey", "on_green"))
+        print("    [+] version: {}".format(dataset_api_version))
+        if dataset_vers.major != client_vers.major:
+            exit_code = 1
+            print(
+                "    [-] "
+                + colored(
+                    "Error: Dataset API major version ({}) mismatchs with mantis_api_client major version ({})".format(
+                        dataset_vers.major, client_vers.major
+                    ),
+                    "white",
+                    "on_red",
+                )
+            )
+
+    # Scenario API
+    print("  [+] Scenario API")
+    print("    [+] address: {}".format(mantis_api_client_config.scenario_api_url))
+    try:
+        scenario_api_version = scenario_api.get_version()
+        scenario_vers = Version(str_vers=scenario_api_version)
+    except requests.exceptions.ConnectionError:
+        exit_code = 1
+        print("    [-] API status: " + colored("not running !", "white", "on_red"))
+    else:
+        print("    [+] API status: " + colored("OK", "grey", "on_green"))
+        print("    [+] version: {}".format(scenario_api_version))
+        if scenario_vers.major != client_vers.major:
+            exit_code = 1
+            print(
+                "    [-] "
+                + colored(
+                    "Error: Scenario API major version ({}) mismatchs with mantis_api_client major version ({})".format(
+                        scenario_vers.major, client_vers.major
+                    ),
+                    "white",
+                    "on_red",
+                )
+            )
+
+    if exit_code != 0:
+        sys.exit(exit_code)
+
+
+#
+# 'info' related functions
+#
+def info_handler(args: Any) -> None:  # noqa: C901
+    """Get personal info."""
+
+    active_profile_domain = oidc_client.get_active_profile_domain(raise_exc=False)
+    if not active_profile_domain:
+        print("[+] Not authenticated")
+        return
+
+    console = Console(highlight=False)
+    rprint = console.print
+    active_tokens = oidc_client.get_active_tokens()
+    access_token = jwt.get_unverified_claims(active_tokens["access_token"])
+    id_token = jwt.get_unverified_claims(active_tokens["id_token"])
+    default_group = oidc_client.get_default_group()
+    groups_comma_sep = ", ".join(id_token["groups"])
+    if default_group:
+        # mark default group bold
+        groups_comma_sep = groups_comma_sep.replace(
+            default_group, f"[bold]{default_group}[/bold]"
+        )
+    sorted_scopes = ", ".join(sorted(access_token["scope"].split()))
+
+    print(f"[+] Connected to {active_profile_domain}")
+    print(f"  [+] Username: {id_token['preferred_username']}")
+    print(f"  [+] Email: {id_token.get('email', 'N/A')}")
+    rprint(rf"  \[+] Group membership: {groups_comma_sep}")
+    print(f"  [+] Scopes: {sorted_scopes}")
+
+
+#
+# 'login_handler' handler
+#
+def login_handler(args: Any) -> None:
+    # Parameters
+    oidc_domain = args.domain
+    username = args.username
+    if args.password_stdin:
+        password = sys.stdin.read().rstrip()
+    elif args.password_fd:
+        with os.fdopen(args.password_fd) as f:
+            password = f.read().rstrip()
+    elif args.password_file:
+        password = args.password_file.read_text().rstrip()
+    elif args.username:
+        password = getpass.getpass()
+    else:
+        password = None
+
+    scope = "openid offline_access groups profile email"
+    redirect_uri = mantis_api_client_config.oidc.redirect_uri
+    redirect_auto = redirect_uri.startswith("http")
+    code_placeholder: List[str] = []
+    if redirect_auto:
+        thread = _init_create_callback_request_handler_thread(code_placeholder)
+        thread.start()
+    if username and password:
+        token = oidc_client.token(
+            oidc_domain, username, password, redirect_uri=redirect_uri, scope=scope
+        )
+    else:
+        auth_url = oidc_client.auth_url(
+            oidc_domain,
+            redirect_uri=redirect_uri,
+            scope=scope,
+        )
+        subprocess.run(["xdg-open", auth_url])
+        print(f"A web browser should been opened for {oidc_domain!r}")
+        if redirect_auto:
+            timeout = mantis_api_client_config.oidc.redirect_url_timeout
+            print(f"Waiting callback for {timeout}s")
+            thread.join(timeout)
+            if thread.is_alive():
+                print("Timeout exceeded, exiting")
+                print("Access NOT granted")
+                exit(1)
+            code = code_placeholder[0]
+        else:
+            code = input("Paste the code displayed in the webpage: ")
+        token = oidc_client.token(
+            oidc_domain,
+            grant_type="authorization_code",
+            code=code,
+            redirect_uri=redirect_uri,
+        )
+    print("Access granted")
+
+    claims = jwt.get_unverified_claims(token["access_token"])
+    selected_group = args.group
+    claimed_groups = claims.get("groups")
+    if claimed_groups and len(claimed_groups) == 1 and selected_group is None:
+        selected_group = claimed_groups[0]
+    if selected_group not in claimed_groups:
+        if selected_group:
+            print(f"You are not member of group '{selected_group}'")
+        selected_group = Prompt.ask(
+            "Select a group to activate", choices=claimed_groups
+        )
+    print(f"Group `{selected_group}` activated")
+    oidc_client.configure_profile(oidc_domain, token["refresh_token"], selected_group)
+
+
+#
+# 'logout_handler' handler
+#
+def logout_handler(args: Any) -> None:
+    oidc_client.configure_profile(None)
+
+
+def _init_create_callback_request_handler_thread(code_placeholder: list) -> Thread:
+    class CallbackHTTPRequestHandler(BaseHTTPRequestHandler):
+        def do_GET(self) -> None:
+            nonlocal code_placeholder
+            query = urlparse(self.path).query
+            if "error" in query:
+                self.send_response(404)
+            else:
+                code = parse_qs(query)["code"][0]
+                code_placeholder.append(code)
+                self.send_response(200)
+            self.send_header("Content-Type", "text/html")
+            self.end_headers()
+            self.wfile.write(
+                "<script>window.close()</script> Authorization {}. You may close this window.\r\n".format(
+                    "failed" if "error" in query else "succeeded"
+                ).encode()
+            )
+
+        def log_request(*args, **kwargs):
+            # do nothing
+            pass
+
+    def serve_one_redirect_callback() -> None:
+        with HTTPServer(
+            (
+                mantis_api_client_config.oidc.redirect_url.host,
+                mantis_api_client_config.oidc.redirect_url.port,
+            ),
+            CallbackHTTPRequestHandler,
+        ) as server:
+            server.handle_request()
+
+    return Thread(target=serve_one_redirect_callback, daemon=True)
+
+
+def add_account_parser(root_parser: argparse.ArgumentParser, subparsers: Any) -> None:
+    # -------------------
+    # --- login options
+    # -------------------
+
+    parser_account = subparsers.add_parser(
+        "account",
+        help="Authentication actions for M&NTIS CLI.",
+        formatter_class=root_parser.formatter_class,
+    )
+    subparsers_account = parser_account.add_subparsers()
+
+    # 'status' command
+    parser_status = subparsers_account.add_parser(
+        "status",
+        help="Get platform status",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_status.set_defaults(func=status_handler)
+
+    # 'info' command
+    parser_info = subparsers_account.add_parser(
+        "info", help="Get personal info", formatter_class=root_parser.formatter_class
+    )
+    parser_info.set_defaults(func=info_handler)
+
+    parser_login = subparsers_account.add_parser(
+        "login",
+        help="Log into a M&ntis account",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_login.add_argument(
+        "--domain",
+        help="The M&ntis cluster SSO domain (default: %(default)s)",
+        default="mantis-platform.io",
+    )
+    parser_login.add_argument(
+        "--username",
+        "-u",
+        help="Your M&ntis cluster SSO username",
+    )
+    parser_login_mex_group = parser_login.add_mutually_exclusive_group()
+    parser_login_mex_group.add_argument(
+        "--password-stdin",
+        action="store_true",
+        help="Read your M&ntis cluster SSO password from stdin",
+    )
+    parser_login_mex_group.add_argument(
+        "--password-fd",
+        type=int,
+        help="Read your M&ntis cluster SSO password from a descriptor",
+    )
+    parser_login_mex_group.add_argument(
+        "--password-file",
+        type=Path,
+        help="Read your M&ntis cluster SSO password from a file",
+    )
+    parser_login.add_argument(
+        "-g",
+        "--group",
+        help="Select the group name that will be used for contextual commands like `mantis scenario run`",
+    )
+    parser_login.set_defaults(func=login_handler)
+
+    # -------------------
+    # --- logout options
+    # -------------------
+
+    parser_logout = subparsers_account.add_parser(
+        "logout",
+        help="Log out from your M&ntis account",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_logout.set_defaults(func=logout_handler)
+
+    parser_account.set_defaults(func=lambda _: parser_account.print_help())

mantis_api_client/cli_parser/attack_parser.py

@@ -0,0 +1,229 @@
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import sys
+import traceback
+from typing import Any
+
+from mantis_scenario_model.scenario_run_config_model import ScenarioRunConfig
+from mantis_scenario_model.unit_attack_model import Empty
+from pydantic.json import pydantic_encoder
+from ruamel.yaml import YAML
+
+from mantis_api_client import scenario_api
+from mantis_api_client.oidc import oidc_client
+from mantis_api_client.utils import colored
+from mantis_api_client.utils import wait_lab
+
+
+#
+# 'attack_list_handler' handler
+#
+def attack_list_handler(args: Any) -> None:
+    try:
+        attacks = scenario_api.fetch_attacks()
+    except Exception as e:
+        print(colored(f"Error when fetching attacks: '{e}'", "red"))
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(attacks, default=pydantic_encoder))
+        return
+
+    width = 35
+    print(f"[+] Available attacks ({len(attacks)}):")
+    name = "NAME"
+    print(f"  [+] \033[1m{name: <{width}}- DESCRIPTION\033[0m")
+
+    for attack in attacks:
+        if not isinstance(attack.mitre_data.subtechnique, Empty):
+            mitre_print = attack.mitre_data.subtechnique.id
+        else:
+            mitre_print = attack.mitre_data.technique.id
+
+        print("  [+] ", end="")
+        print(f"{attack.name: <{width}}", end="")
+        print(f"- {attack.title} ({mitre_print})")
+
+
+#
+# 'attack_info_handler' handler
+#
+def attack_info_handler(args: Any) -> None:
+    # Parameters
+    attack_name = args.attack_name
+
+    try:
+        attack = scenario_api.fetch_attack_by_name(attack_name)
+    except Exception as e:
+        print(
+            colored(
+                f"Error when fetching attack {attack_name}: '{e}'", "red", "on_white"
+            )
+        )
+        sys.exit(1)
+
+    if args.json:
+        print(attack.json())
+        return
+
+    print("[+] Attack information:")
+    print(f"  [+] \033[1mID\033[0m: {attack.worker_id}")
+    print(f"  [+] \033[1mName\033[0m: {attack.name}")
+    print(f"  [+] \033[1mDescription\033[0m: {attack.description}")
+    print("  [+] \033[1mAvailable scenario config\033[0m:")
+    for scenario_config in attack.scenario_config:
+        print(
+            f"    [+] {scenario_config['name']} (Topology: {scenario_config['file']})"
+        )
+
+
+#
+# 'attack_run_handler' handler
+#
+def attack_run_handler(args: Any) -> None:
+    # Parameters
+    attack_name = args.attack_name
+    scenario_config_name = args.scenario_config_name
+    scenario_run_config_path = args.scenario_run_config_path
+
+    if not args.group_id:
+        try:
+            group_id = oidc_client.get_default_group(raise_exc=True)
+        except Exception as e:
+            print(colored(f"Error when fetching attacks: '{e}'", "red"))
+            sys.exit(1)
+    else:
+        group_id = args.group_id
+
+    # Safety checks
+    try:
+        attack = scenario_api.fetch_attack_by_name(attack_name)
+    except Exception as e:
+        print(colored(f"{e}", "red", "on_white"))
+        sys.exit(1)
+    print(f"[+] Going to execute attack: {attack.name}")
+
+    # Safety checks
+    if len(attack.scenario_config) == 0:
+        print(
+            f"Cannot run attack '{attack_name}', because it does not have any unit scenario"
+        )
+        sys.exit(-1)
+
+    if scenario_config_name is None:
+        print(
+            "Needed argument --scenario_config_name, in order to choose the unit scenario to run"
+        )
+        print("Available unit scenarios:")
+        for available_scenario_config in attack.scenario_config:
+            print(f"  [+] {available_scenario_config['name']}")
+        sys.exit(-1)
+
+    for available_scenario_config in attack.scenario_config:
+        if available_scenario_config["name"] == scenario_config_name:
+            break
+    else:
+        print(f"Select scenario config '{scenario_config_name}' is not available")
+        print("Available unit scenarios:")
+        for available_scenario_config in attack.scenario_config:
+            print(f"  [+] {available_scenario_config['name']}")
+        sys.exit(-1)
+
+    # Manage scenario run configuration
+    if scenario_run_config_path is None:
+        scenario_run_config_dict = {}
+    else:
+        with open(scenario_run_config_path, "r") as fd:
+            yaml_content = fd.read()
+        loader = YAML(typ="rt")
+        scenario_run_config_dict = loader.load(yaml_content)
+    scenario_run_config = ScenarioRunConfig(**scenario_run_config_dict)
+
+    # Launch topology
+    try:
+        lab_id = scenario_api.run_attack(
+            attack, scenario_config_name, scenario_run_config, group_id
+        )
+
+        print(f"[+] Scenario lab ID: {lab_id}")
+
+        wait_lab(lab_id)
+
+        print("[+] Scenario ended")
+    except Exception as e:
+        print(colored(f"Error when running attack {attack_name}: '{e}'", "red"))
+        print(traceback.format_exc())
+        sys.exit(1)
+    finally:
+        if args.destroy_after_scenario:
+            print("[+] Stopping lab...")
+            scenario_api.stop_lab(lab_id)
+
+
+def add_attack_parser(root_parser: argparse.ArgumentParser, subparsers: Any) -> None:
+    # --------------------
+    # --- Scenario API options (attack)
+    # --------------------
+
+    parser_attack = subparsers.add_parser(
+        "attack",
+        help="Scenario API related commands (attacks)",
+        formatter_class=root_parser.formatter_class,
+    )
+    subparsers_attack = parser_attack.add_subparsers()
+
+    # 'attack_list' command
+    parser_attack_list = subparsers_attack.add_parser(
+        "list",
+        help="List all available attacks",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_attack_list.set_defaults(func=attack_list_handler)
+    parser_attack_list.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'attack_info' command
+    parser_attack_info = subparsers_attack.add_parser(
+        "info",
+        help="Get information about an attack",
+        formatter_class=root_parser.formatter_class,
+    )
+    parser_attack_info.set_defaults(func=attack_info_handler)
+    parser_attack_info.add_argument("attack_name", type=str, help="The attack name")
+    parser_attack_info.add_argument(
+        "--json", help="Return JSON result.", action="store_true"
+    )
+
+    # 'attack_run' command
+    parser_attack_run = subparsers_attack.add_parser(
+        "run", help="Run a specific attack", formatter_class=root_parser.formatter_class
+    )
+    parser_attack_run.set_defaults(func=attack_run_handler)
+    parser_attack_run.add_argument("attack_name", type=str, help="The attack name")
+    parser_attack_run.add_argument(
+        "--destroy",
+        action="store_true",
+        dest="destroy_after_scenario",
+        help="Do not keep the lab up after scenario execution (False by default)",
+    )
+    parser_attack_run.add_argument(
+        "--scenario_config_name",
+        dest="scenario_config_name",
+        help="Allows to define the unit scenario to run for a unit attack",
+    )
+    parser_attack_run.add_argument(
+        "--group_id",
+        dest="group_id",
+        help="The group ID that have ownership on lab",
+    )
+    parser_attack_run.add_argument(
+        "--scenario_run_config",
+        action="store",
+        required=False,
+        dest="scenario_run_config_path",
+        help="Input path of a YAML configuration for the scenario run",
+    )
+
+    parser_attack.set_defaults(func=lambda _: parser_attack.print_help())