Updated OTP Generate with new workflow

utils/tui.py

@@ -18,9 +18,12 @@ from textual.widgets import (
     Tabs,
 )
 
-from flows.otp import otp_activities_by_agent, otp_generate, otp_revoke
+from flows.otp import otp_activities_by_agent, otp_revoke
 from flows.prepPolicy import menu_policy_enforce
 from flows.quietAgent import findQuietAgents
+from models.agent import Agent
+from models.policy import Policy
+from screens.otpworkflowscreen import OTPWorkflowScreen
 from services.agenthandler import findAgents, moveAgents, toggleEnforcement
 from services.API import AirlockAPIWrapper
 from services.policyhandler import confirmUpdateAfromE
@@ -28,6 +31,7 @@ from utils.configmanager import load_env
 from utils.setup import get_base_directory, load_user_config
 from utils.utils import open_directory
 from widgets.multiagentselector import MultiAgentSelector
+from widgets.OTP_generate import OTPGenerator
 from widgets.policytreewidget import PolicyTreeWidget
 from widgets.themeselector import ThemeSelector
 
@@ -107,7 +111,7 @@ class MainMenuScreen(Screen):
             ("🔀 - Move - Other", "move_other_button"),
         ],
         "otp": [
-            ("🔐 - Generate OTPs", "otp_generate_button"),
+            ("🎫 - Generate OTPs", "otp_generate_button"),
             ("📊 - OTP Activities By Agent", "otp_activities_button"),
             ("❌ - Revoke OTPs", "otp_revoke_button"),
         ],
@@ -117,8 +121,9 @@ class MainMenuScreen(Screen):
         ],
     }
 
-    def __init__(self) -> None:
+    def __init__(self, api: AirlockAPIWrapper) -> None:
         super().__init__()
+        self.api = api
         self.extras = load_env("EXTRAS")
         wd = load_env("WORKING_DIR") or os.getcwd()
         if not os.path.isdir(wd):
@@ -144,7 +149,6 @@ class MainMenuScreen(Screen):
             Tab("OTP", id="otp"),
             Tab("Directory", id="dir"),
             Tab("Settings", id="settings"),
-            Tab("Multi Select", id="multi_select"),
         ]
 
         if self.extras == "POLICYPREP":
@@ -207,8 +211,6 @@ class MainMenuScreen(Screen):
             content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
         elif tab_id == "settings":
             content.mount(ThemeSelector())
-        elif tab_id == "multi_select":
-            content.mount(MultiAgentSelector(self.app.devices.to_dict("records")))
         else:
             content.mount(Static(f"Unknown tab: {tab_id}"))
 
@@ -235,6 +237,30 @@ class MainMenuScreen(Screen):
         _PENDING_JOB = ("restart",)
         self.app.exit()
 
+    def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
+        """Handle OTP generation request from the workflow."""
+        global _PENDING_JOB
+
+        # Log what we received
+        logger.info(
+            "OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d",
+            len(message.devices),
+            message.requestor,
+            message.reasoning,
+            message.duration,
+        )
+
+        # Set up the job to run the OTP generation
+        _PENDING_JOB = (
+            "otp_workflow",
+            message.devices,
+            message.requestor,
+            message.reasoning,
+            message.duration,
+        )
+
+        self.app.exit()
+
     def on_directory_tree_file_selected(
         self, event: DirectoryTree.FileSelected
     ) -> None:
@@ -268,7 +294,10 @@ class MainMenuScreen(Screen):
             case "move_other_button":
                 _PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {})
             case "otp_generate_button":
-                _PENDING_JOB = ("legacy", otp_generate, (self.app.api,), {})
+                # NEW: Push OTP workflow screen instead of legacy function
+                self.app.push_screen(OTPWorkflowScreen(self.app.devices))
+                event.stop()
+                return  # Don't exit the app
             case "otp_activities_button":
                 _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {})
             case "otp_revoke_button":
@@ -314,16 +343,20 @@ class Loxide(App):
 
         # Add error handling for API calls
         try:
-            self.policies = api.policy_find_all()
-            self.devices = api.agent_find_all()
+            self.policies = [
+                Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()
+            ]
+            self.devices = [
+                Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()
+            ]
         except Exception as exc:
             logger.error("Failed to load policies/devices: %s", exc)
             self.policies = None
             self.devices = None
 
-    def on_mount(self) -> None:
+    def on_mount(self, api: AirlockAPIWrapper) -> None:
         self.theme = self._textual_theme
-        self.push_screen(MainMenuScreen())
+        self.push_screen(MainMenuScreen(api))
 
     def action_quit(self) -> None:
         global _PENDING_JOB
@@ -410,10 +443,44 @@ def run_Loxide(api: AirlockAPIWrapper) -> None:
 
         if job[0] == "multi_agent_action":
             # Handle multi-agent selection
-            # TODO: Implement actual multi-agent action handling
             logger.info("Multi-agent action with selected agents: %s", job[1])
             continue
 
+        # NEW: Handle OTP workflow
+        if job[0] == "otp_workflow":
+            _, devices, requestor, reasoning, duration = job
+
+            # Call your OTP generation with the parameters
+            def otp_generate_with_params():
+
+                print(f"\n{'='*60}")
+                print("OTP GENERATION")
+                print(f"{'='*60}")
+                print(f"Requestor: {requestor}")
+                print(f"Reasoning: {reasoning}")
+                print(f"Duration: {duration} minutes")
+                print(f"\nGenerating OTPs for {len(devices)} devices:")
+                print(f"{'='*60}\n")
+
+                # Call your actual OTP generation function
+                # You'll need to adapt otp_generate to accept these parameters
+                # For now, this is a placeholder showing the structure
+                for device in devices:
+                    print(f"Device: {device}")
+                    print(f"  Requestor: {requestor}")
+                    print(f"  Reason: {reasoning}")
+                    print(f"  Duration: {duration} minutes")
+                    # TODO: Actually call your API to generate OTP
+                    # result = api.generate_otp(device, requestor, reasoning, duration)
+                    print()
+
+                print(f"{'='*60}")
+                print("OTP Generation Complete!")
+                print(f"{'='*60}")
+
+            _run_legacy_job(otp_generate_with_params, (), {})
+            continue
+
         break