From 7a7ee4f0de7f45c2a39248f932065e01a7343cd3 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 4 Jan 2025 19:40:56 +0200
Subject: [PATCH] fix(cli config):

---
 agentic_security/__main__.py |  6 ++++--
 agentic_security/lib.py      | 29 ++++++++++++++++++++++++-----
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/agentic_security/__main__.py b/agentic_security/__main__.py
index 9574db6..210aded 100644
--- a/agentic_security/__main__.py
+++ b/agentic_security/__main__.py
@@ -33,12 +33,14 @@ class CLI:
         sys.path.append(os.path.dirname("."))
         AgenticSecurity().entrypoint()
 
-    def init(self):
+    def init(self, host: str = "0.0.0.0", port: int = 8718):
         """
         Generate the default CI configuration file.
         """
         sys.path.append(os.path.dirname("."))
-        AgenticSecurity().generate_default_cfg()
+        AgenticSecurity().generate_default_cfg(host, port)
+
+    i = init
 
 
 def main():
diff --git a/agentic_security/lib.py b/agentic_security/lib.py
index 373ac2d..d2db2af 100644
--- a/agentic_security/lib.py
+++ b/agentic_security/lib.py
@@ -254,16 +254,31 @@ class AgenticSecurity(CfgMixin):
             exit(1)
 
         self.load_config(self.default_path)
-        print(f"Configuration: {self.config}")
+        logger.info("Configuration loaded successfully.")
+        print(self.config)
+        datasets = list(self.get_config_value("modules").values())
+        for d in datasets:
+            d["selected"] = True
+        self.scan(
+            llmSpec=self.get_config_value("general.llmSpec"),
+            maxBudget=self.get_config_value("general.maxBudget"),
+            datasets=datasets,
+            max_th=self.get_config_value("general.max_th"),
+            optimize=self.get_config_value("general.optimize"),
+            enableMultiStepAttack=self.get_config_value(
+                "general.enableMultiStepAttack"
+            ),
+        )
 
-    def generate_default_cfg(self):
+    def generate_default_cfg(self, host: str = "0.0.0.0", port: int = 8718):
+        # Accept host / port as parameters
         with open(self.default_path, "w") as f:
             f.write(
                 """
 [general]
 # General configuration for the security scan
 llmSpec = \"""
-POST http://0.0.0.0:9094/v1/self-probe
+POST http://$HOST:$PORT/v1/self-probe
 Authorization: Bearer XXXXX
 Content-Type: application/json
 
@@ -279,7 +294,7 @@ enableMultiStepAttack = false # Enable multi-step attack simulations
 [modules.AgenticBackend]
 dataset_name = "AgenticBackend"
 [modules.AgenticBackend.opts]
-port = 9094
+port = $PORT
 modules = ["encoding"]
 
 [thresholds]
@@ -289,7 +304,11 @@ medium = 0.3
 high = 0.5
 
 
-"""
+""".replace(
+                    "$HOST", host
+                ).replace(
+                    "$PORT", str(port)
+                )
             )
 
         logger.info(