aiverify-moonshot 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

moonshot/integrations/cli/redteam/attack_module.py

@@ -1,40 +1,49 @@
+from ast import literal_eval
+
 import cmd2
 from rich.console import Console
 from rich.table import Table
 
 from moonshot.api import api_delete_attack_module, api_get_all_attack_module_metadata
+from moonshot.integrations.cli.utils.process_data import filter_data
 
 console = Console()
 
 
-def list_attack_modules() -> None:
+# ------------------------------------------------------------------------------
+# CLI Functions
+# ------------------------------------------------------------------------------
+def list_attack_modules(args) -> list | None:
     """
     Retrieves and prints the metadata of all attack modules.
+
+    Args:
+        args: A namespace object from argparse. It should have an optional attribute:
+        find (str): Optional field to find attack module(s) with a keyword.
+        pagination (str): Optional field to paginate attack modules.
+
+    Returns:
+         list | None: A list of AttackModule or None if there is no result.
     """
-    print("Listing attack modules may take a while...")
-    attack_module_metadata_list = api_get_all_attack_module_metadata()
-
-    if attack_module_metadata_list:
-        table = Table(
-            title="Attack Module List",
-            show_lines=True,
-            expand=True,
-            header_style="bold",
-        )
-        table.add_column("No.", width=2)
-        table.add_column("Details", justify="left", width=98)
-
-        for attack_module_index, attack_module_data in enumerate(
-            attack_module_metadata_list, 1
-        ):
-            attack_module_data_str = ""
-            for k, v in attack_module_data.items():
-                attack_module_data_str += f"[blue]{k.capitalize()}:[/blue] {v}\n\n"
-            table.add_row(str(attack_module_index), attack_module_data_str)
+    try:
+        print("Listing attack modules may take a while...")
+        attack_module_metadata_list = api_get_all_attack_module_metadata()
+        keyword = args.find.lower() if args.find else ""
+        pagination = literal_eval(args.pagination) if args.pagination else ()
 
-        console.print(table)
-    else:
-        console.print("[red]There are no attack modules found.[/red]", style="bold")
+        if attack_module_metadata_list:
+            filtered_attack_modules_list = filter_data(
+                attack_module_metadata_list, keyword, pagination
+            )
+            if filtered_attack_modules_list:
+                _display_attack_modules(filtered_attack_modules_list)
+                return filtered_attack_modules_list
+
+        console.print("[red]There are no attack modules found.[/red]")
+        return None
+
+    except Exception as e:
+        print(f"[list_attack_modules]: {str(e)}")
 
 
 def delete_attack_module(args) -> None:
@@ -59,6 +68,41 @@ def delete_attack_module(args) -> None:
         print(f"[delete_attack_module]: {str(e)}")
 
 
+# ------------------------------------------------------------------------------
+# Helper functions: Display on cli
+# ------------------------------------------------------------------------------
+def _display_attack_modules(attack_modules: list) -> None:
+    """
+    Display a list of attack modules.
+
+    This function takes a list of attack modules and displays them in a table format. If the list is empty, it prints a
+    message indicating that no attack modules were found.
+
+    Args:
+        attack_modules (list): A list of attack modules.
+
+    Returns:
+        None
+    """
+    table = Table(
+        title="Attack Module List",
+        show_lines=True,
+        expand=True,
+        header_style="bold",
+    )
+    table.add_column("No.", width=2)
+    table.add_column("Details", justify="left", width=98)
+
+    for idx, attack_module_data in enumerate(attack_modules, 1):
+        attack_module_data_str = ""
+        for k, v in attack_module_data.items():
+            if k != "idx":
+                attack_module_data_str += f"[blue]{k.capitalize()}:[/blue] {v}\n\n"
+        idx = attack_module_data.get("idx", idx)
+        table.add_row(str(idx), attack_module_data_str)
+    console.print(table)
+
+
 # Delete attack module arguments
 delete_attack_module_args = cmd2.Cmd2ArgumentParser(
     description="Delete an attack module.",
@@ -68,3 +112,25 @@ delete_attack_module_args = cmd2.Cmd2ArgumentParser(
 delete_attack_module_args.add_argument(
     "attack_module", type=str, help="The ID of the attack module to delete"
 )
+
+# List attack modules arguments
+list_attack_modules_args = cmd2.Cmd2ArgumentParser(
+    description="List all attack modules.",
+    epilog='Example:\n list_attack_modules -f "text"',
+)
+
+list_attack_modules_args.add_argument(
+    "-f",
+    "--find",
+    type=str,
+    help="Optional field to find attack module(s) with keyword",
+    nargs="?",
+)
+
+list_attack_modules_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate attack module(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)

moonshot/integrations/cli/redteam/context_strategy.py

@@ -1,4 +1,5 @@
 import argparse
+from ast import literal_eval
 
 import cmd2
 from rich.console import Console
@@ -10,6 +11,7 @@ from moonshot.api import (
     api_update_context_strategy,
 )
 from moonshot.integrations.cli.active_session_cfg import active_session
+from moonshot.integrations.cli.utils.process_data import filter_data
 from moonshot.src.redteaming.session.session import Session
 
 console = Console()
@@ -32,12 +34,12 @@ def use_context_strategy(args: argparse.Namespace) -> None:
 
     # Check if current session exists. If it does, update context strategy and number of previous prompts
     if active_session:
-        active_session["context_strategy"] = new_context_strategy_name
-        active_session["cs_num_of_prev_prompts"] = num_of_prev_prompts
         try:
             api_update_context_strategy(
                 active_session["session_id"], new_context_strategy_name
             )
+            active_session["context_strategy"] = new_context_strategy_name
+            active_session["cs_num_of_prev_prompts"] = num_of_prev_prompts
             print(
                 f"Updated session: {active_session['session_id']}. "
                 f"Context Strategy: {active_session['context_strategy']}."
@@ -52,30 +54,36 @@ def use_context_strategy(args: argparse.Namespace) -> None:
         )
 
 
-def list_context_strategies() -> None:
+def list_context_strategies(args) -> list | None:
     """
     List all context strategies available.
+
+    Args:
+        args: A namespace object from argparse. It should have an optional attribute:
+        find (str): Optional field to find context strategies with a keyword.
+        pagination (str): Optional field to paginate context strategies.
+
+    Returns:
+        list | None: A list of ContextStrategy or None if there is no result.
     """
-    context_strategy_metadata_list = api_get_all_context_strategy_metadata()
-    if context_strategy_metadata_list:
-        table = Table(
-            title="Context Strategy List",
-            show_lines=True,
-            expand=True,
-            header_style="bold",
-        )
-        table.add_column("No.", justify="left", width=2)
-        table.add_column("Context Strategy Information", justify="left", width=98)
-        for context_strategy_index, context_strategy_data in enumerate(
-            context_strategy_metadata_list, 1
-        ):
-            context_strategy_data_str = ""
-            for k, v in context_strategy_data.items():
-                context_strategy_data_str += f"[blue]{k.capitalize()}:[/blue] {v}\n\n"
-            table.add_row(str(context_strategy_index), context_strategy_data_str)
-        console.print(table)
-    else:
-        console.print("[red]There are no context strategies found.[/red]", style="bold")
+    try:
+        context_strategy_metadata_list = api_get_all_context_strategy_metadata()
+        keyword = args.find.lower() if args.find else ""
+        pagination = literal_eval(args.pagination) if args.pagination else ()
+
+        if context_strategy_metadata_list:
+            filtered_context_strategies_list = filter_data(
+                context_strategy_metadata_list, keyword, pagination
+            )
+            if filtered_context_strategies_list:
+                _display_context_strategies(filtered_context_strategies_list)
+                return filtered_context_strategies_list
+
+        console.print("[red]There are no context strategies found.[/red]")
+        return None
+
+    except Exception as e:
+        print(f"[list_context_strategies]: {str(e)}")
 
 
 def clear_context_strategy() -> None:
@@ -118,6 +126,36 @@ def delete_context_strategy(args) -> None:
         print(f"[delete_context_strategy]: {str(e)}")
 
 
+def _display_context_strategies(context_strategies: list) -> None:
+    """
+    Display a list of context strategies.
+
+    This function takes a list of context strategies and displays them in a table format. If the list is empty,
+    it prints a message indicating that no attack modules were found.
+
+    Args:
+        context_strategies (list): A list of context strategies.
+
+    Returns:
+        None
+    """
+    table = Table(
+        title="Context Strategy List",
+        show_lines=True,
+        expand=True,
+        header_style="bold",
+    )
+    table.add_column("No.", justify="left", width=2)
+    table.add_column("Context Strategy Information", justify="left", width=98)
+    for idx, context_strategy_data in enumerate(context_strategies, 1):
+        context_strategy_data_str = ""
+        for k, v in context_strategy_data.items():
+            if k != "idx":
+                context_strategy_data_str += f"[blue]{k.capitalize()}:[/blue] {v}\n\n"
+        table.add_row(str(idx), context_strategy_data_str)
+    console.print(table)
+
+
 # Use context strategy arguments
 use_context_strategy_args = cmd2.Cmd2ArgumentParser(
     description="Use a context strategy.",
@@ -145,3 +183,25 @@ delete_context_strategy_args = cmd2.Cmd2ArgumentParser(
 delete_context_strategy_args.add_argument(
     "context_strategy", type=str, help="The ID of the context strategy to delete"
 )
+
+# List context strategies arguments
+list_context_strategies_args = cmd2.Cmd2ArgumentParser(
+    description="List all context strategies.",
+    epilog='Example:\n list_context_strategies -f "previous_prompt"',
+)
+
+list_context_strategies_args.add_argument(
+    "-f",
+    "--find",
+    type=str,
+    help="Optional field to find context strategies with keyword",
+    nargs="?",
+)
+
+list_context_strategies_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate context strategies(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)

moonshot/integrations/cli/redteam/prompt_template.py

@@ -20,11 +20,11 @@ def use_prompt_template(args: argparse.Namespace) -> None:
 
     # Check if current session exists
     if active_session:
-        active_session["prompt_template"] = new_prompt_template_name
         try:
             api_update_prompt_template(
                 active_session["session_id"], new_prompt_template_name
             )
+            active_session["prompt_template"] = new_prompt_template_name
             print(
                 f"Updated session: {active_session['session_id']}. "
                 f"Prompt Template: {active_session['prompt_template']}."

moonshot/integrations/cli/redteam/redteam.py

@@ -6,12 +6,14 @@ from moonshot.integrations.cli.redteam.attack_module import (
     delete_attack_module,
     delete_attack_module_args,
     list_attack_modules,
+    list_attack_modules_args,
 )
 from moonshot.integrations.cli.redteam.context_strategy import (
     clear_context_strategy,
     delete_context_strategy,
     delete_context_strategy_args,
     list_context_strategies,
+    list_context_strategies_args,
     use_context_strategy,
     use_context_strategy_args,
 )
@@ -21,16 +23,30 @@ from moonshot.integrations.cli.redteam.prompt_template import (
     use_prompt_template_args,
 )
 from moonshot.integrations.cli.redteam.session import (
+    add_bookmark,
+    add_bookmark_args,
     automated_rt_session_args,
+    delete_bookmark,
+    delete_bookmark_args,
     delete_session,
     delete_session_args,
     end_session,
+    export_bookmarks,
+    export_bookmarks_args,
+    list_bookmarks,
+    list_bookmarks_args,
     list_sessions,
+    list_sessions_args,
     new_session,
     new_session_args,
     run_attack_module,
+    show_prompts,
+    use_bookmark,
+    use_bookmark_args,
     use_session,
     use_session_args,
+    view_bookmark,
+    view_bookmark_args,
 )
 
 
@@ -46,8 +62,9 @@ class RedTeamCommandSet(cmd2.CommandSet):
     def do_end_session(self, _: cmd2.Statement) -> None:
         end_session()
 
-    def do_list_sessions(self, _: cmd2.Statement) -> None:
-        list_sessions()
+    @cmd2.with_argparser(list_sessions_args)
+    def do_list_sessions(self, args: argparse.Namespace) -> None:
+        list_sessions(args)
 
     @cmd2.with_argparser(use_session_args)
     def do_use_session(self, args: argparse.Namespace) -> None:
@@ -60,8 +77,9 @@ class RedTeamCommandSet(cmd2.CommandSet):
     def do_clear_prompt_template(self, _: cmd2.Statement) -> None:
         clear_prompt_template()
 
-    def do_list_context_strategies(self, _: cmd2.Statement) -> None:
-        list_context_strategies()
+    @cmd2.with_argparser(list_context_strategies_args)
+    def do_list_context_strategies(self, args: argparse.Namespace) -> None:
+        list_context_strategies(args)
 
     @cmd2.with_argparser(use_context_strategy_args)
     def do_use_context_strategy(self, args: argparse.Namespace) -> None:
@@ -74,8 +92,9 @@ class RedTeamCommandSet(cmd2.CommandSet):
     def do_run_attack_module(self, args: argparse.Namespace) -> None:
         run_attack_module(args)
 
-    def do_list_attack_modules(self, _: cmd2.Statement) -> None:
-        list_attack_modules()
+    @cmd2.with_argparser(list_attack_modules_args)
+    def do_list_attack_modules(self, args: argparse.Namespace) -> None:
+        list_attack_modules(args)
 
     @cmd2.with_argparser(delete_session_args)
     def do_delete_session(self, args: argparse.Namespace) -> None:
@@ -88,3 +107,30 @@ class RedTeamCommandSet(cmd2.CommandSet):
     @cmd2.with_argparser(delete_attack_module_args)
     def do_delete_attack_module(self, args: argparse.Namespace) -> None:
         delete_attack_module(args)
+
+    @cmd2.with_argparser(add_bookmark_args)
+    def do_add_bookmark(self, args: argparse.Namespace) -> None:
+        add_bookmark(args)
+
+    @cmd2.with_argparser(use_bookmark_args)
+    def do_use_bookmark(self, args: argparse.Namespace) -> None:
+        use_bookmark(args)
+
+    @cmd2.with_argparser(delete_bookmark_args)
+    def do_delete_bookmark(self, args: argparse.Namespace) -> None:
+        delete_bookmark(args)
+
+    @cmd2.with_argparser(list_bookmarks_args)
+    def do_list_bookmarks(self, args: argparse.Namespace) -> None:
+        list_bookmarks(args)
+
+    @cmd2.with_argparser(view_bookmark_args)
+    def do_view_bookmark(self, args: argparse.Namespace) -> None:
+        view_bookmark(args)
+
+    @cmd2.with_argparser(export_bookmarks_args)
+    def do_export_bookmarks(self, args: argparse.Namespace) -> None:
+        export_bookmarks(args)
+
+    def do_show_prompts(self, _: cmd2.Statement) -> None:
+        show_prompts()