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)
+    activeagents = api.otp_find_active()
-    otp_dict = {}
+    awaitingagents = api.otp_find_awaiting()
-    for agent in agents:
+    enforcedagents = api.otp_find_enforced()
-        otp_info = api.otp_find_by_agent(agent.agentid)
+    revokedagents = api.otp_find_revoked()
-        otp_dict[agent.hostname] = otp_info
+   
+   
+    # 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]
+    activeagents = api.otp_find_active()
-    if otp_dict and list_to_revoke:
+    awaitingagents = api.otp_find_awaiting()
-        for revokee in list_to_revoke:
+    
-            api.otp_revoke(revokee)
+    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}")
+     
+

services/API.py

@@ -142,6 +142,18 @@ class AirlockAPIWrapper:
         result = self._post("/v1/otp/usage", payload)
         return pd.DataFrame(result["response"]["otpusage"])
     
+    def otp_find_enforced(self) -> pd.DataFrame:
+        """Find OTPs that are awaiting activation."""
+        payload = {"status": "2"}
+        result = self._post("/v1/otp/usage", payload)
+        return pd.DataFrame(result["response"]["otpusage"])
+    
+    def otp_find_revoked(self) -> pd.DataFrame:
+        """Find OTPs that are awaiting activation."""
+        payload = {"status": "3"}
+        result = self._post("/v1/otp/usage", payload)
+        return pd.DataFrame(result["response"]["otpusage"])
+    
     def otp_find_by_agent(self, agentid) -> pd.DataFrame:
         """Find OTP by agent."""
         payload = {"agentid": agentid}

utils/menus.py

@@ -55,7 +55,7 @@ def menu_main(api: AirlockAPIWrapper):
         clear_screen()
         displayIntro()
         # Add Settings, and give option to change working dir
-        print(colorText("1. ✅   - Move Device(s) to local approval", "yellow"))
+        print(colorText("1. ✅  - Move Device(s) to local approval (Placeholder)", "yellow"))
         print(colorText("2. 🎫  - OTP", "yellow"))
         print(colorText("3. 🔄  - Move to Audit/Enforcement", "yellow"))
         print(colorText("4. 🔍  - Device Search", "yellow"))
@@ -269,7 +269,7 @@ def menu_policymanagment(api: AirlockAPIWrapper):
 def menu_settings():
     while True:
         print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
-        print(colorText("This Feature is still in development", "cyan"))
+        print(colorText("This Feature is still in development, if you have ideas for options you would like to see, let us know.", "cyan"))
         # print(colorText("2. Sub-option B","cyan"))
         print(colorText("B. 🔙  - Back", "yellow"))
         choice = get_sanitized_input("Enter your choice: ")

utils/selector.py

@@ -1,11 +1,10 @@
 import logging
 from typing import Any, Callable, List, Optional, Union
-
+import pandas as pd
 from utils.utils import colorText, get_sanitized_input
 
 logger = logging.getLogger(__name__)
 
-
 class Selector:
     @staticmethod
     def _get_sorted_items(items: List[Any], label_func: Callable[[Any], str]) -> List[Any]:
@@ -18,6 +17,11 @@ class Selector:
         num_columns: int = 4,
         header: str = "Available Choices:"
     ) -> None:
+        # Force single column if items are DataFrame rows
+        
+        if items and isinstance(items[0], (pd.Series, dict)):
+            num_columns = 1
+
         rows = (len(items) + num_columns - 1) // num_columns
         print(f"\n{header}")
         for row in range(rows):
@@ -39,7 +43,6 @@ class Selector:
         if not selected:
             print(" (none)")
             return
-
         sorted_selected = sorted(selected, key=lambda item: label_func(item).lower())
         rows = (len(sorted_selected) + num_columns - 1) // num_columns
         for row in range(rows):
@@ -72,7 +75,8 @@ class Selector:
         label_func: Callable[[Any], str],
         allow_multiple: bool = False,
         prompt_each: bool = False,
-        header: str = "Available Choices:"
+        header: str = "Available Choices:",
+        num_columns: int = 4
     ) -> Union[Optional[Any], List[Any]]:
         if not items:
             logger.warning("No items available for selection.")
@@ -84,11 +88,9 @@ class Selector:
 
         if allow_multiple:
             while True:
-                Selector._display_choices(remaining_items, label_func, header=header)
+                Selector._display_choices(remaining_items, label_func, num_columns=num_columns, header=header)
-                Selector._display_selected_items(selected, label_func)
+                Selector._display_selected_items(selected, label_func, num_columns=num_columns)
-
                 choice = get_sanitized_input("Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: ").strip().lower()
-
                 if choice == "q":
                     break
                 elif choice == "r":
@@ -96,10 +98,8 @@ class Selector:
                     remaining_items = full_sorted_items.copy()
                     print(colorText("🔄 Selections reset.", "yellow"))
                     continue
-
                 indices = Selector._parse_selection_input(choice, len(remaining_items))
                 newly_selected = []
-
                 for index in indices:
                     item = remaining_items[index - 1]
                     if item not in selected:
@@ -109,13 +109,10 @@ class Selector:
                             logger.info(f"Selected: {label_func(item)}")
                     else:
                         logger.warning("Item already selected.")
-
-                # Remove newly selected items from remaining list
                 remaining_items = [item for item in remaining_items if item not in newly_selected]
-
             return selected if selected else None
         else:
-            Selector._display_choices(full_sorted_items, label_func, header=header)
+            Selector._display_choices(full_sorted_items, label_func,  num_columns=num_columns, header=header)
             try:
                 choice = int(get_sanitized_input("Select one item by number: "))
                 if 1 <= choice <= len(full_sorted_items):
@@ -136,7 +133,6 @@ class Selector:
     ) -> List[Any]:
         print(colorText("Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white"))
         mode = get_sanitized_input("").strip().lower()
-
         if mode == "a":
             return items
         selected = Selector._select_from_list(
@@ -146,10 +142,8 @@ class Selector:
             prompt_each=False,
             header=header
         )
-
         if not selected:
             return items
-
         if mode == "i":
             print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
             return selected
@@ -237,4 +231,88 @@ class Selector:
                 logger.info("User declined action.")
                 return False
             else:
-                logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
+                logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
+
+    @staticmethod
+    def select_dataframe_rows(
+        df: pd.DataFrame,
+        columns: Optional[List[str]] = None,
+        allow_multiple: bool = False,
+        prompt_each: bool = False,
+        header: str = "Available Rows:"
+    ) -> List[pd.Series]:
+        if df.empty:
+            print("DataFrame is empty.")
+            return []
+
+        if columns:
+            df = df[columns]
+
+        items = [row for _, row in df.iterrows()]
+        label_func = lambda row: str(row.to_dict())
+
+        result = Selector._select_from_list(
+            items,
+            label_func=label_func,
+            allow_multiple=allow_multiple,
+            prompt_each=prompt_each,
+            header=header
+        )
+
+        if isinstance(result, pd.Series):
+            return [result]
+        elif isinstance(result, list):
+            return result
+        else:
+            return []
+
+    @staticmethod
+    def select_dataframe_with_mode(
+        df: pd.DataFrame,
+        columns: Optional[List[str]] = None,
+        header: str = "Available Rows:"
+    ) -> List[pd.Series]:
+        if df.empty:
+            print("⚠️ DataFrame is empty.")
+            return []
+
+        # Filter columns if specified
+        if columns:
+            df = df[columns]
+
+        items = df.to_dict("records")
+        label_func = lambda row: " | ".join(str(row[col]) for col in df.columns)
+
+        # Show rows first
+        print(colorText(header, "cyan"))
+        for i, row in enumerate(items):
+            print(f"{i}: {label_func(row)}")
+
+        # Prompt for mode once
+        print(colorText("\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white"))
+        mode = get_sanitized_input("").strip().lower()
+
+        if mode == "a":
+            return [pd.Series(row) for row in items]
+
+        # Prompt for selection only once
+        selected = Selector._select_from_list(
+            items,
+            label_func=label_func,
+            allow_multiple=True,
+            prompt_each=False,
+            header=header
+        )
+
+        if not selected:
+            return [pd.Series(row) for row in items]
+
+        if mode == "i":
+            print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
+            return [pd.Series(row) for row in selected]
+        elif mode == "e":
+            print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow"))
+            return [pd.Series(row) for row in items if row not in selected]
+        else:
+            print(colorText("⚠️ Invalid mode. Returning all rows.", "yellow"))
+            return [pd.Series(row) for row in items]