fix(cli config):

msoedov · Jan 4, 2025 · 7a7ee4f · 7a7ee4f
1 parent 64f7f4b
commit 7a7ee4f
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 7 deletions.
diff --git a/agentic_security/__main__.py b/agentic_security/__main__.py
@@ -33,12 +33,14 @@ def ci(self):
         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
@@ -254,16 +254,31 @@ def entrypoint(self):
             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 @@ def generate_default_cfg(self):
 [modules.AgenticBackend]
 dataset_name = "AgenticBackend"
 [modules.AgenticBackend.opts]
-port = 9094
+port = $PORT
 modules = ["encoding"]
 
 [thresholds]
@@ -289,7 +304,11 @@ def generate_default_cfg(self):
 high = 0.5
 
 
-"""
+""".replace(
+                    "$HOST", host
+                ).replace(
+                    "$PORT", str(port)
+                )
             )
 
         logger.info(