Updated OTP Generate with new workflow
This commit is contained in:
+79
-12
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user