fix(prowler-check-kreator): fix index list problems

util/prowler_check_kreator/lib/llms/gemini.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 
 import google.generativeai as genai
 
@@ -67,7 +68,11 @@ def _generate_content(self, prompt_parts: list) -> str:
 
  def _prepare_check_prompt(self, check_name: str, check_reference: str) -> list:
  """Prepare the prompt for generating the check."""
- class_name = check_reference.split("(")[0].split("class ")[1]
+
+ class_name = re.search(
+ r"class\s+([A-Za-z_][A-Za-z0-9_]*)\s*\((.*?)\)\s*:", check_reference
+ ).group(1)
+
  prompt_parts = [
  f"Your task is to create a new security check called '{check_name}' for Prowler (a Cloud Security tool).",
  "The control is a Python class that inherits from the Check class and has only one method called execute.",

util/prowler_check_kreator/prowler_check_kreator.py

@@ -83,7 +83,7 @@ def __init__(self, provider: str, check_name: str):
  # Let the user to use the model that he wants
  self._model = None
  supported_models = [
- "gemini-1.5-flash (default)",
+ "gemini-1.5-flash",
  "gemini-1.5-pro",
  "gemini-1.0-pro",
  ]
@@ -361,8 +361,8 @@ def _write_test_file(self) -> None:
  sys.exit(prowler_check_creator.kreate_check())
 
  except ValueError as e:
- print(f"Error: {e}")
+ print(f"Error: {e} (line {sys.exc_info()[-1].tb_lineno})")
  sys.exit(1)
  except Exception as e:
- print(f"Error: {e}")
+ print(f"Error: {e} (line {sys.exc_info()[-1].tb_lineno})")
  sys.exit(1)