RustImplementation #23

Merged
mysticmomba merged 118 commits from RustImplementation into master 2025-11-04 18:13:24 -05:00
4 changed files with 218 additions and 31 deletions
Showing only changes of commit 595f19ab27 - Show all commits
+108 -11
View File
@@ -16,11 +16,16 @@
import logging import logging
import os
import pandas as pd
from services.agenthandler import selectAgents from services.agenthandler import selectAgents
from services.API import AirlockAPIWrapper from services.API import AirlockAPIWrapper
from utils.selector import Selector from utils.selector import Selector
from utils.configmanager import load_env
from utils.utils import colorText, get_sanitized_input from utils.utils import colorText, get_sanitized_input
from datetime import datetime
from services.agenthandler import selectAgents
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -57,17 +62,109 @@ def generate(api: AirlockAPIWrapper):
return otp_dict return otp_dict
def otp_activities_by_agent(api: AirlockAPIWrapper): 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): 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}")
+12
View File
@@ -142,6 +142,18 @@ class AirlockAPIWrapper:
result = self._post("/v1/otp/usage", payload) result = self._post("/v1/otp/usage", payload)
return pd.DataFrame(result["response"]["otpusage"]) 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: def otp_find_by_agent(self, agentid) -> pd.DataFrame:
"""Find OTP by agent.""" """Find OTP by agent."""
payload = {"agentid": agentid} payload = {"agentid": agentid}
+2 -2
View File
@@ -55,7 +55,7 @@ def menu_main(api: AirlockAPIWrapper):
clear_screen() clear_screen()
displayIntro() displayIntro()
# Add Settings, and give option to change working dir # 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("2. 🎫 - OTP", "yellow"))
print(colorText("3. 🔄 - Move to Audit/Enforcement", "yellow")) print(colorText("3. 🔄 - Move to Audit/Enforcement", "yellow"))
print(colorText("4. 🔍 - Device Search", "yellow")) print(colorText("4. 🔍 - Device Search", "yellow"))
@@ -269,7 +269,7 @@ def menu_policymanagment(api: AirlockAPIWrapper):
def menu_settings(): def menu_settings():
while True: while True:
print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan")) 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("2. Sub-option B","cyan"))
print(colorText("B. 🔙 - Back", "yellow")) print(colorText("B. 🔙 - Back", "yellow"))
choice = get_sanitized_input("Enter your choice: ") choice = get_sanitized_input("Enter your choice: ")
+96 -18
View File
@@ -1,11 +1,10 @@
import logging import logging
from typing import Any, Callable, List, Optional, Union from typing import Any, Callable, List, Optional, Union
import pandas as pd
from utils.utils import colorText, get_sanitized_input from utils.utils import colorText, get_sanitized_input
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class Selector: class Selector:
@staticmethod @staticmethod
def _get_sorted_items(items: List[Any], label_func: Callable[[Any], str]) -> List[Any]: def _get_sorted_items(items: List[Any], label_func: Callable[[Any], str]) -> List[Any]:
@@ -18,6 +17,11 @@ class Selector:
num_columns: int = 4, num_columns: int = 4,
header: str = "Available Choices:" header: str = "Available Choices:"
) -> None: ) -> 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 rows = (len(items) + num_columns - 1) // num_columns
print(f"\n{header}") print(f"\n{header}")
for row in range(rows): for row in range(rows):
@@ -39,7 +43,6 @@ class Selector:
if not selected: if not selected:
print(" (none)") print(" (none)")
return return
sorted_selected = sorted(selected, key=lambda item: label_func(item).lower()) sorted_selected = sorted(selected, key=lambda item: label_func(item).lower())
rows = (len(sorted_selected) + num_columns - 1) // num_columns rows = (len(sorted_selected) + num_columns - 1) // num_columns
for row in range(rows): for row in range(rows):
@@ -72,7 +75,8 @@ class Selector:
label_func: Callable[[Any], str], label_func: Callable[[Any], str],
allow_multiple: bool = False, allow_multiple: bool = False,
prompt_each: bool = False, prompt_each: bool = False,
header: str = "Available Choices:" header: str = "Available Choices:",
num_columns: int = 4
) -> Union[Optional[Any], List[Any]]: ) -> Union[Optional[Any], List[Any]]:
if not items: if not items:
logger.warning("No items available for selection.") logger.warning("No items available for selection.")
@@ -84,11 +88,9 @@ class Selector:
if allow_multiple: if allow_multiple:
while True: 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() 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": if choice == "q":
break break
elif choice == "r": elif choice == "r":
@@ -96,10 +98,8 @@ class Selector:
remaining_items = full_sorted_items.copy() remaining_items = full_sorted_items.copy()
print(colorText("🔄 Selections reset.", "yellow")) print(colorText("🔄 Selections reset.", "yellow"))
continue continue
indices = Selector._parse_selection_input(choice, len(remaining_items)) indices = Selector._parse_selection_input(choice, len(remaining_items))
newly_selected = [] newly_selected = []
for index in indices: for index in indices:
item = remaining_items[index - 1] item = remaining_items[index - 1]
if item not in selected: if item not in selected:
@@ -109,13 +109,10 @@ class Selector:
logger.info(f"Selected: {label_func(item)}") logger.info(f"Selected: {label_func(item)}")
else: else:
logger.warning("Item already selected.") 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] remaining_items = [item for item in remaining_items if item not in newly_selected]
return selected if selected else None return selected if selected else None
else: 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: try:
choice = int(get_sanitized_input("Select one item by number: ")) choice = int(get_sanitized_input("Select one item by number: "))
if 1 <= choice <= len(full_sorted_items): if 1 <= choice <= len(full_sorted_items):
@@ -136,7 +133,6 @@ class Selector:
) -> List[Any]: ) -> List[Any]:
print(colorText("Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white")) print(colorText("Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white"))
mode = get_sanitized_input("").strip().lower() mode = get_sanitized_input("").strip().lower()
if mode == "a": if mode == "a":
return items return items
selected = Selector._select_from_list( selected = Selector._select_from_list(
@@ -146,10 +142,8 @@ class Selector:
prompt_each=False, prompt_each=False,
header=header header=header
) )
if not selected: if not selected:
return items return items
if mode == "i": if mode == "i":
print(colorText(f"✅ Included {len(selected)} item(s).", "green")) print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
return selected return selected
@@ -237,4 +231,88 @@ class Selector:
logger.info("User declined action.") logger.info("User declined action.")
return False return False
else: 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]