diff --git a/flows/otp.py b/flows/otp.py index b651b70..d0f5e0d 100644 --- a/flows/otp.py +++ b/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) \ No newline at end of file + + 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}") + + diff --git a/services/API.py b/services/API.py index 5742d0a..c618a79 100644 --- a/services/API.py +++ b/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} diff --git a/utils/menus.py b/utils/menus.py index d7ed3bf..f6c24b9 100644 --- a/utils/menus.py +++ b/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: ") diff --git a/utils/selector.py b/utils/selector.py index eb7fcaa..b500efe 100644 --- a/utils/selector.py +++ b/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_selected_items(selected, label_func) - + Selector._display_choices(remaining_items, label_func, num_columns=num_columns, header=header) + 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'.") \ No newline at end of file + 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] \ No newline at end of file