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

Post Black Linting

Browse Source
This commit is contained in:
Zarithas
2025-11-06 11:04:59 -05:00
parent f33b041ac0
commit b538f12e9a
20 changed files with 1106 additions and 618 deletions
Download Patch File Download Diff File Expand all files Collapse all files
services/API.py
+38 -34
View File
@@ -23,7 +23,6 @@ import requests
logger = logging.getLogger(__name__)
class AirlockAPIWrapper:
"""
A wrapper class for interacting with the Airlock API.
@@ -141,25 +140,25 @@ class AirlockAPIWrapper:
payload = {"status": "0"}
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}
result = self._post("/v1/otp/usage", payload)
return pd.DataFrame(result["response"]["otpusage"])
def otp_generate(self, agentid: str, duration: int, purpose: str) -> str:
"""Generate a new OTP for an agent."""
payload = {
@@ -175,7 +174,7 @@ class AirlockAPIWrapper:
payload = {"otpid": otpid}
result = self._post("/v1/otp/activities", payload)
return pd.DataFrame(result["response"]["otpactivities"])
def otp_revoke(self, otpid: str) -> dict:
"""
Revoke an active OTP.
@@ -186,18 +185,17 @@ class AirlockAPIWrapper:
"""
payload = {"otpid": otpid}
return self._post("/v1/otp/revoke", payload)
def otp_validate(self, otpcode: str) -> dict:
"""
Validate an OTP code.
Parameters:
- otpcode (str): The OTP code to validate.
Returns:
- dict: JSON response indicating validity.
"""
payload = {"otpcode": otpcode}
return self._post("/v1/otp/validate", payload)
def otp_validate(self, otpcode: str) -> dict:
"""
Validate an OTP code.
Parameters:
- otpcode (str): The OTP code to validate.
Returns:
- dict: JSON response indicating validity.
"""
payload = {"otpcode": otpcode}
return self._post("/v1/otp/validate", payload)
# Policy Management
def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict:
@@ -225,7 +223,7 @@ class AirlockAPIWrapper:
payload = {"groupid": groupid}
result = self._post("/v1/group/agents", payload)
return pd.DataFrame(result["response"]["agents"])
def policy_list_allowlists(self, groupid: str) -> pd.DataFrame:
"""List allowlists assigned to a specific policy group."""
payload = {"groupid": groupid}
@@ -236,31 +234,37 @@ class AirlockAPIWrapper:
"""Set audit mode for a policy group. 1=Audit, 0=Enforcement"""
payload = {"groupid": groupid, "auditmode": auditmode}
return self._post("/v1/group/settings/auditmode", payload)
def policy_set_script_custom(self,
groupid: str,
script_custom: int,
scripts_audit: List[str],
scripts_disabled: List[str],
scripts_respect: List[str],
) -> dict:
def policy_set_script_custom(
self,
groupid: str,
script_custom: int,
scripts_audit: List[str],
scripts_disabled: List[str],
scripts_respect: List[str],
) -> dict:
"""Set audit mode for a policy group. 1=Audit, 0=Enforcement"""
payload = {"groupid": groupid,
"script_custom": script_custom,
"scripts_audit": scripts_audit,
"scripts_disabled": scripts_disabled,
"scripts_respect": scripts_respect
}
payload = {
"groupid": groupid,
"script_custom": script_custom,
"scripts_audit": scripts_audit,
"scripts_disabled": scripts_disabled,
"scripts_respect": scripts_respect,
}
return self._post("/v1/group/settings/script_custom", payload)
# Execution History
def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> str:
def history_logging(
self, type: List[str], checkpoint: str, policy: List[str]
) -> str:
"""Retrieve execution history logs."""
payload = {"type": type, "checkpoint": checkpoint, "policy": policy}
result = self._post("/v1/logging/exechistories", payload)
return result["response"]["exechistories"]
def history_execution(self, today: str, date_selected: str, agent_name: str) -> List[Dict]:
def history_execution(
self, today: str, date_selected: str, agent_name: str
) -> List[Dict]:
"""
Retrieve execution history logs.
services/agenthandler.py
+93 -28
View File
@@ -47,7 +47,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
print(colorText("No agents selected or invalid history range.", "red"))
return
historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d")
historical_date = (datetime.now() - timedelta(days=history_days)).strftime(
"%Y-%m-%d"
)
today = datetime.now().strftime("%Y-%m-%d")
all_history = []
@@ -56,7 +58,11 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
try:
exechistory = api.history_execution(today, historical_date, agent.hostname)
except Exception as e:
print(colorText(f"❌ Error retrieving history for {agent.hostname}: {e}", "red"))
print(
colorText(
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
)
)
continue
if isinstance(exechistory, list):
@@ -76,7 +82,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
print(colorText(f"{key}: {value}", "green"))
print("\n")
else:
print(colorText(f"No execution history found for {agent.hostname}.", "yellow"))
print(
colorText(f"No execution history found for {agent.hostname}.", "yellow")
)
if outputjson:
print(json.dumps(all_history, indent=2))
@@ -92,6 +100,7 @@ def findAllAgents(api):
return agents
def findAgents(api, return_dataframe):
agents = selectAgents(api)
working_dir = load_env("WORKING_DIR")
@@ -113,8 +122,14 @@ def findAgents(api, return_dataframe):
print(agent_df)
logging.debug("Displayed DataFrame to console.")
user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower()
if user_input == 'y':
user_input = (
get_sanitized_input(
"\nWould you like to export the results to a CSV file? (y/n): "
)
.strip()
.lower()
)
if user_input == "y":
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = f"agentsearch_{timestamp}.csv"
file_path = os.path.join(str(working_dir), filename)
@@ -131,17 +146,27 @@ def findAgents(api, return_dataframe):
else:
logging.debug("User declined to export the DataFrame.")
def collect_device_names() -> List[str]:
print(colorText("🔍 Device Search", "cyan"))
print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
print(colorText("When you're done, press Enter twice (Three times if you have a single device).\n", "cyan"))
print(
colorText(
"Enter the device hostnames you'd like to search for, one per line.", "cyan"
)
)
print(
colorText(
"When you're done, press Enter twice (Three times if you have a single device).\n",
"cyan",
)
)
print(colorText("Example:", "cyan"))
print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan"))
print(colorText("Paste or type your device names below:", "white"))
device_input_lines = []
empty_line_count = 0
valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')
valid_line_pattern = re.compile(r"^[a-zA-Z0-9_\- ]+$")
while True:
line = get_sanitized_input("")
@@ -158,7 +183,12 @@ def collect_device_names() -> List[str]:
if valid_line_pattern.match(stripped_line):
device_input_lines.append(stripped_line)
else:
print(colorText(f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow"))
print(
colorText(
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
"yellow",
)
)
return [name for name in device_input_lines if name]
@@ -168,10 +198,13 @@ def choose_match_type() -> bool:
return get_sanitized_input("").strip().lower() in ["y", "yes"]
def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool) -> List['Agent']:
def match_agents(
device_names: List[str], agents: List["Agent"], use_exact: bool
) -> List["Agent"]:
if use_exact:
return [
agent for agent in agents
agent
for agent in agents
if agent.hostname.lower() in [name.lower() for name in device_names]
]
else:
@@ -180,23 +213,38 @@ def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool
return [agent for agent in agents if regex.search(agent.hostname)]
def show_unmatched(device_names: List[str], matched_agents: List['Agent'], use_exact: bool):
def show_unmatched(
device_names: List[str], matched_agents: List["Agent"], use_exact: bool
):
if use_exact:
unmatched = [name for name in device_names if not any(agent.hostname.lower() == name.lower() for agent in matched_agents)]
unmatched = [
name
for name in device_names
if not any(
agent.hostname.lower() == name.lower() for agent in matched_agents
)
]
else:
unmatched = [name for name in device_names if not any(re.search(re.escape(name), agent.hostname, re.IGNORECASE) for agent in matched_agents)]
unmatched = [
name
for name in device_names
if not any(
re.search(re.escape(name), agent.hostname, re.IGNORECASE)
for agent in matched_agents
)
]
if unmatched:
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
def enrich_agents(agents: List['Agent'], policies: List['Policy']):
def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
for agent in agents:
agent.enrich_with_policies(policies)
def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']:
def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
device_names = collect_device_names()
if not device_names:
logger.debug("No device names entered")
@@ -227,11 +275,11 @@ def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']:
if idx < len(matched_agents):
line += f"{matched_agents[idx].hostname:<30}"
logger.info(line)
matched_agents = Selector.select_with_mode(
matched_agents,
label_func=lambda agent: agent.hostname,
header="Matched Devices:"
header="Matched Devices:",
)
if not matched_agents:
@@ -263,11 +311,17 @@ def moveAgentToRelatedPolicy(
if agent.groupid in policy_relationship_map:
target_policy = policy_relationship_map[agent.groupid]
elif agent.groupid in policy_relationship_map.values():
logger.debug(f"Agent {agent.hostname} is already in an audit group. No action needed.")
print(f"Agent {agent.hostname} is already in an audit group. No action needed.")
logger.debug(
f"Agent {agent.hostname} is already in an audit group. No action needed."
)
print(
f"Agent {agent.hostname} is already in an audit group. No action needed."
)
return
else:
logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.")
logger.warning(
f"Error: No corresponding audit policy found for groupid: {agent.groupid}."
)
return
elif mode == "enforcement":
@@ -275,10 +329,14 @@ def moveAgentToRelatedPolicy(
if agent.groupid in inverse_map:
target_policy = inverse_map[agent.groupid]
elif agent.groupid in inverse_map.values():
logger.info(f"Agent {agent.hostname} is already in an enforcement group. No action needed.")
logger.info(
f"Agent {agent.hostname} is already in an enforcement group. No action needed."
)
return
else:
logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.")
logger.warning(
f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}."
)
return
else:
@@ -299,25 +357,32 @@ def toggleEnforcement(api: AirlockAPIWrapper):
devices = selectAgents(api)
for device in devices:
print(device.hostname)
confirm = Selector.confirm("Would you like to continue with these devices? Y/N: ")
confirm = Selector.confirm(
"Would you like to continue with these devices? Y/N: "
)
if direction and devices and confirm:
for device in devices:
result = moveAgentToRelatedPolicy(api,device, str(direction).lower())
result = moveAgentToRelatedPolicy(api, device, str(direction).lower())
logger.info(f"{device.hostname}: result: {result}")
get_sanitized_input("Press enter to continue")
def moveAgents(api: AirlockAPIWrapper):
devices = selectAgents(api)
for device in devices:
print(device.hostname)
confirm_devices = Selector.confirm("Would you like to continue with these devices? Y/N: ")
confirm_devices = Selector.confirm(
"Would you like to continue with these devices? Y/N: "
)
if devices and confirm_devices:
policies = selectPolicies(api, False)
confirm_move = Selector.confirm(f"Would you like to move these devices to {policies[0].name}?")
confirm_move = Selector.confirm(
f"Would you like to move these devices to {policies[0].name}?"
)
if confirm_move:
for device in devices:
result = api.agent_move(device.agentid, policies[0].groupid)
logger.info(f"{device.hostname}: result: {result}")
else:
logger.info("Exiting without change")
get_sanitized_input("Press enter to continue")
get_sanitized_input("Press enter to continue")
services/policyhandler.py
+11 -8
View File
@@ -34,7 +34,6 @@ from utils.utils import areYouSure, colorText, get_sanitized_input
logger = logging.getLogger(__name__)
def pullPolicyExechistories(
api: AirlockAPIWrapper,
policy: Policy,
@@ -72,7 +71,7 @@ def pullPolicyExechistories(
) as pbar:
while True:
histories = api.history_logging(
type=type, checkpoint=checkpoint, policy= [policy.name]
type=type, checkpoint=checkpoint, policy=[policy.name]
)
# Ensure histories is a list of dictionaries
@@ -98,13 +97,17 @@ def pullPolicyExechistories(
# Update checkpoint on last item
if index == len(histories) - 1:
checkpoint = history_item["checkpoint"] # pyright: ignore[reportArgumentType]
checkpoint = history_item[
"checkpoint"
] # pyright: ignore[reportArgumentType]
filebar.desc = f"Checkpoint Progress: {checkpoint}"
break
try:
history_date = datetime.datetime.strptime(
history_item["datetime"].replace(" +0000 UTC", ""), # pyright: ignore[reportArgumentType]
history_item["datetime"].replace(
" +0000 UTC", ""
), # pyright: ignore[reportArgumentType]
"%Y-%m-%dT%H:%M:%SZ",
).date()
except ValueError:
@@ -178,7 +181,7 @@ def pullPolicyExechistories(
def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
import airlock_libs
executionhist_policy = pd.DataFrame()
exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days)
if exehist is not None:
@@ -205,7 +208,7 @@ def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
executionhist_policy = executionhist_policy.sort_values(
by=["sha256", "filename"]
)
logger.debug( f"Staging of Execution history for policy: {policy} is complete")
logger.debug(f"Staging of Execution history for policy: {policy} is complete")
print(
colorText(
f"Staging of Execution history for policy: {policy} is complete",
@@ -237,10 +240,10 @@ def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
for enforcement_policy, audit_policy in policy_relationship_map.items():
api.policy_clone(enforcement_policy, audit_policy)
api.policy_set_auditmode(audit_policy, "1")
def confirmUpdateAfromE(api: AirlockAPIWrapper):
areYouSure()
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
if confirmation.strip() == "I AGREE":
updateAuditPoliciesFromEnforcementPolices(api)
updateAuditPoliciesFromEnforcementPolices(api)
services/security.py
+25 -12
View File
@@ -28,8 +28,8 @@ import keyring
# Constants
KDF_ITERATIONS = 200_000
SALT_SIZE = 16 # 128-bit Salt
NONCE_SIZE = 12 # AES-GCM
SALT_SIZE = 16 # 128-bit Salt
NONCE_SIZE = 12 # AES-GCM
KEY_SIZE = 32 # AES-256
logger = logging.getLogger(__name__)
@@ -49,9 +49,11 @@ def configure_keyring_backend():
system = platform.system()
if system == "Windows":
import keyring.backends.Windows
keyring.set_keyring(keyring.backends.Windows.WinVaultKeyring())
elif system == "Linux":
import keyring.backends.kwallet
keyring.set_keyring(keyring.backends.kwallet.DBusKeyring())
else:
raise EnvironmentError(f"Unsupported OS: {system}")
@@ -68,8 +70,9 @@ def store_api_key(service: str, username: str, api_key: str, password: str):
b64 = base64.b64encode(blob).decode()
keyring.set_password(service, username, b64)
logger.debug(f"API key for service '{service}' and user '{username}' stored successfully.")
logger.debug(
f"API key for service '{service}' and user '{username}' stored successfully."
)
print("\n✅ API key stored securely.")
print("The program will now exit. Press Enter to continue...")
@@ -90,8 +93,8 @@ def retrieve_api_key(service: str, username: str, password: str) -> str:
raise ValueError("No stored secret for this service/username.")
blob = base64.b64decode(b64)
salt = blob[:SALT_SIZE]
nonce = blob[SALT_SIZE:SALT_SIZE + NONCE_SIZE]
ct = blob[SALT_SIZE + NONCE_SIZE:]
nonce = blob[SALT_SIZE : SALT_SIZE + NONCE_SIZE]
ct = blob[SALT_SIZE + NONCE_SIZE :]
key = _derive_key(password.encode(), salt)
aesgcm = AESGCM(key)
pt = aesgcm.decrypt(nonce, ct, associated_data=None)
@@ -124,7 +127,9 @@ def getAPI(USERNAME, SERVICE_NAME):
if api_key_exists(SERVICE_NAME, USERNAME):
for attempt in range(1, 4):
password = getpass(f"Attempt {attempt}/3 - Enter password to unlock your API key: ")
password = getpass(
f"Attempt {attempt}/3 - Enter password to unlock your API key: "
)
try:
apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password)
logging.debug("API key successfully retrieved.")
@@ -134,9 +139,15 @@ def getAPI(USERNAME, SERVICE_NAME):
logging.error("Failed to retrieve API key after 3 incorrect attempts.")
raise ValueError("Failed to retrieve API key after 3 incorrect attempts.")
else:
logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.")
api_key = getpass(f"No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip()
print("Please exit and relaunch program after saving your credential to avoid errors")
logging.warning(
f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'."
)
api_key = getpass(
f"No API key found. Please enter your API key for '{SERVICE_NAME}': "
).strip()
print(
"Please exit and relaunch program after saving your credential to avoid errors"
)
while True:
password = getpass("Create a password to encrypt your API key: ")
@@ -155,7 +166,9 @@ def getAPI(USERNAME, SERVICE_NAME):
logging.error(f"Failed to store API key: {e}")
break
else:
logging.warning("Password does not meet complexity requirements. Try again.")
logging.warning(
"Password does not meet complexity requirements. Try again."
)
class APIKeyManager:
@@ -169,4 +182,4 @@ class APIKeyManager:
def get(cls) -> str:
if cls._api_key is None:
raise ValueError("API key not loaded. Call APIKeyManager.load() first.")
return cls._api_key
return cls._api_key