mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Implemented OTP Activities and OTP Revoke

Browse Source
This commit is contained in:
Zarithas
2025-10-27 12:14:01 -04:00
parent c52767b3a3
commit 595f19ab27
4 changed files with 218 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
flows/otp.py
+108 -11
View File
@@ -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}")
services/API.py
+12
View File
@@ -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
+2 -2
View File
@@ -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
+96 -18
View File
@@ -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'.")
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]