Implemented OTP Activities and OTP Revoke

flows/otp.py

@@ -16,11 +16,16 @@
 
 
 import logging
+import os
+import pandas as pd
 
 from services.agenthandler import selectAgents
 from services.API import AirlockAPIWrapper
 from utils.selector import Selector
+from utils.configmanager import load_env
 from utils.utils import colorText, get_sanitized_input
+from datetime import datetime
+from services.agenthandler import selectAgents
 
 logger = logging.getLogger(__name__)
 
@@ -57,17 +62,109 @@ def generate(api: AirlockAPIWrapper):
     return otp_dict
 
 def otp_activities_by_agent(api: AirlockAPIWrapper):
-    agents = selectAgents(api)
-    otp_dict = {}
-    for agent in agents:
-        otp_info = api.otp_find_by_agent(agent.agentid)
-        otp_dict[agent.hostname] = otp_info
+    activeagents = api.otp_find_active()
+    awaitingagents = api.otp_find_awaiting()
+    enforcedagents = api.otp_find_enforced()
+    revokedagents = api.otp_find_revoked()
+   
+   
+    # Add a 'status' column to each DataFrame
+    activeagents['status'] = 'active'
+    awaitingagents['status'] = 'awaiting'
+    enforcedagents['status'] = 'enforced'
+    revokedagents['status'] = 'revoked'
+
+    # Combine all into one DataFrame
+    combined_agents = pd.concat([activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True)
+    combined_agents = combined_agents.sort_values(by='otpid', ascending=False)
+
+    #Optionally, select specific hosts
+    user_input = get_sanitized_input("\nWould you like to search for a specific device? (y/n): ").strip().lower()
+    if user_input == 'y':
+        agentnames = []
+        agents = selectAgents(api)
+        for agent in agents:
+            agentnames.append(agent.hostname)
+
+        combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)]
+
+    #Present and select rows
+    selected_rows = Selector.select_dataframe_with_mode(
+        combined_agents,
+        columns=['otpid', 'hostname', 'status','purpose','granted'],
+        header="OTP Sessions"
+    )
+    combined_df = pd.DataFrame()
+
+    for row in selected_rows:
+        otpid = row['otpid']
+        hostname = row['hostname']
+        result = api.otp_get_activities(otpid)
+        result['hostname'] = hostname
+        if not result.empty:
+            logger.info(f"Activities for {hostname} (otpid: {otpid}):\n{result}")
+            combined_df = pd.concat([combined_df, result], ignore_index=True)
+        else:
+            logger.info(f"No activities found for {hostname} (otpid: {otpid})")
+
+    user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower()
+    if user_input == 'y':
+        working_dir = load_env("WORKING_DIR")
+        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        filename = f"otp_activities_{timestamp}.csv"
+        file_path = os.path.join(str(working_dir), filename)
+
+        combined_df.to_csv(file_path, index=False)
+        logging.info(f"Exported Data to {file_path}")
+
+        print(
+            colorText(
+                f"\n✅ OTP Activity exported to: {working_dir}\\{filename}",
+                "green",
+            )
+        )
+    else:
+        logging.debug("User declined to export the DataFrame.")
+
 
-    return otp_dict
 
 def revoke(api: AirlockAPIWrapper):
-    otp_dict = otp_activities_by_agent(api)
-    list_to_revoke = [entry["otpid"] for entry in otp_dict]
-    if otp_dict and list_to_revoke:
-        for revokee in list_to_revoke:
-            api.otp_revoke(revokee)
+
+    activeagents = api.otp_find_active()
+    awaitingagents = api.otp_find_awaiting()
+    
+    activeagents['status'] = 'active'
+    awaitingagents['status'] = 'awaiting'
+
+    combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True)
+    combined_agents = combined_agents.sort_values(by='otpid', ascending=False)
+
+ # Combine all into one DataFrame
+    combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True)
+    combined_agents = combined_agents.sort_values(by='otpid', ascending=False)
+
+    #Optionally, select specific hosts
+    user_input = get_sanitized_input("\nWould you like to search for a specific device? (y/n): ").strip().lower()
+    if user_input == 'y':
+        agentnames = []
+        agents = selectAgents(api)
+        for agent in agents:
+            agentnames.append(agent.hostname)
+
+        combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)]
+
+    #Present and select rows
+    selected_rows = Selector.select_dataframe_with_mode(
+        combined_agents,
+        columns=['otpid', 'hostname', 'status','purpose','granted'],
+        header="OTP Sessions"
+    )
+    combined_df = pd.DataFrame()
+
+    for row in selected_rows:
+        otpid = row['otpid']
+        hostname = row['hostname']
+        result = api.otp_revoke(otpid)
+        logger.info(f"{hostname} (otpid: {otpid}):\n{result}")
+     
+