From 64f7f4b2d536a6033d8c422e47b16ca2872b3b8c Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 4 Jan 2025 19:32:17 +0200
Subject: [PATCH] feat(Update CLI):

---
 agentic_security/__main__.py | 40 +++++++++++++++++++++++++++---------
 agentic_security/lib.py      |  4 ++++
 agentic_security/test_lib.py |  6 +++---
 3 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/agentic_security/__main__.py b/agentic_security/__main__.py
index 5b308ba..9574db6 100644
--- a/agentic_security/__main__.py
+++ b/agentic_security/__main__.py
@@ -8,28 +8,48 @@ from agentic_security.app import app
 from agentic_security.lib import AgenticSecurity
 
 
-class T:
-    def server(self, port=8718, host="0.0.0.0"):
+class CLI:
+    def server(self, port: int = 8718, host: str = "0.0.0.0"):
+        """
+        Launch the Agentic Security server.
+
+        Args:
+            port (int): Port number for the server to listen on. Default is 8718.
+            host (str): Host address for the server. Default is "0.0.0.0".
+        """
         sys.path.append(os.path.dirname("."))
         config = uvicorn.Config(
             app, port=port, host=host, log_level="info", reload=True
         )
         server = uvicorn.Server(config)
         server.run()
-        return
 
-    def headless(self):
+    s = server
+
+    def ci(self):
+        """
+        Run Agentic Security in CI mode.
+        """
         sys.path.append(os.path.dirname("."))
         AgenticSecurity().entrypoint()
 
-
-def entrypoint():
-    fire.Fire(T().server)
+    def init(self):
+        """
+        Generate the default CI configuration file.
+        """
+        sys.path.append(os.path.dirname("."))
+        AgenticSecurity().generate_default_cfg()
 
 
-def ci_entrypoint():
-    fire.Fire(T().headless)
+def main():
+    """
+    Entry point for the CLI. Default behavior launches the server,
+    while subcommands allow CI or configuration generation.
+    """
+    fire.Fire(
+        CLI,
+    )
 
 
 if __name__ == "__main__":
-    ci_entrypoint()
+    main()
diff --git a/agentic_security/lib.py b/agentic_security/lib.py
index 6ac437f..373ac2d 100644
--- a/agentic_security/lib.py
+++ b/agentic_security/lib.py
@@ -291,3 +291,7 @@ high = 0.5
 
 """
             )
+
+        logger.info(
+            f"Default configuration generated successfully to {self.default_path}."
+        )
diff --git a/agentic_security/test_lib.py b/agentic_security/test_lib.py
index e970be1..a8afddb 100644
--- a/agentic_security/test_lib.py
+++ b/agentic_security/test_lib.py
@@ -5,8 +5,9 @@ import subprocess
 import tempfile
 import time
 
-import agentic_security.test_spec_assets as test_spec_assets
 import pytest
+
+import agentic_security.test_spec_assets as test_spec_assets
 from agentic_security.lib import AgenticSecurity
 
 
@@ -125,7 +126,6 @@ class TestAS:
 
 
 class TestEntrypointCI:
-
     def test_generate_default_cfg_to_tmp_path(self):
         """
         Test that the `generate_default_cfg` method generates a valid default config file in a temporary path.
@@ -145,7 +145,7 @@ class TestEntrypointCI:
             assert os.path.exists(temp_path), f"{temp_path} file should be generated."
 
             # Validate the contents of the generated config file
-            with open(temp_path, "r") as f:
+            with open(temp_path) as f:
                 generated_content = f.read()
                 assert (
                     "maxBudget = 1000000" in generated_content