Major refactor: security enhancements, modularization, config integration, reduced Parquet reliance
- Migrated codebase to class-based architecture for better modularity and maintainability - Introduced system_config.json for centralized configuration (required for runtime) - Added structured working directories for improved file organization - Significantly reduced reliance on Parquet; replaced with alternative data handling - Implemented security improvements across modules - Several TODOs remain in the main script for future enhancements - Linter formatting affected readability in some files (e.g., utils); cleanup is on the agenda
This commit is contained in:
+239
@@ -0,0 +1,239 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import pandas as pd
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
||||
class AirlockAPIWrapper:
|
||||
"""
|
||||
A wrapper class for interacting with the Airlock API.
|
||||
Provides methods for managing agents, policies, hashes, OTPs, and execution history.
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: str, api_key: str):
|
||||
"""
|
||||
Initialize the API wrapper.
|
||||
|
||||
Parameters:
|
||||
- base_url (str): Base URL of the Airlock API.
|
||||
- api_key (str): API key for authentication.
|
||||
"""
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.headers = {"X-APIKey": self.api_key}
|
||||
|
||||
def _post(self, endpoint: str, payload: Optional[dict] = None) -> dict:
|
||||
"""
|
||||
Internal method to send POST requests to the API.
|
||||
|
||||
Parameters:
|
||||
- endpoint (str): API endpoint.
|
||||
- payload (dict, optional): Request payload.
|
||||
|
||||
Returns:
|
||||
- dict: JSON response from the API.
|
||||
"""
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
data = json.dumps(payload or {})
|
||||
try:
|
||||
logger.debug(f"POST Request to {url} with payload: {payload}")
|
||||
response = requests.post(url, headers=self.headers, data=data, verify=False)
|
||||
response.raise_for_status()
|
||||
logger.debug(f"Response received from {url}")
|
||||
return response.json()
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"API request failed: {e}")
|
||||
raise
|
||||
|
||||
# Allowlist Management
|
||||
def allowlist_find_all(self) -> pd.DataFrame:
|
||||
"""
|
||||
Retrieve all applications in the allowlist.
|
||||
|
||||
Returns:
|
||||
- pd.DataFrame: DataFrame containing allowlisted applications.
|
||||
"""
|
||||
result = self._post("/v1/application", {})
|
||||
return pd.DataFrame(result["response"]["applications"])
|
||||
|
||||
# Agent Management
|
||||
def agent_find_all(self) -> pd.DataFrame:
|
||||
"""Retrieve all agents."""
|
||||
result = self._post("/v1/agent/find", {})
|
||||
return pd.DataFrame(result["response"]["agents"])
|
||||
|
||||
def agent_find_by_hostname(self, hostname: str) -> pd.DataFrame:
|
||||
"""Find agents by hostname."""
|
||||
payload = {"hostname": hostname}
|
||||
result = self._post("/v1/agent/find", payload)
|
||||
return pd.DataFrame(result["response"]["agents"])
|
||||
|
||||
def agent_find_by_id(self, agentid: str) -> pd.DataFrame:
|
||||
"""Find agents by agent ID."""
|
||||
payload = {"agentid": agentid}
|
||||
result = self._post("/v1/agent/find", payload)
|
||||
return pd.DataFrame(result["response"]["agents"])
|
||||
|
||||
def agent_find_by_status(self, status: int) -> pd.DataFrame:
|
||||
"""Find agents by status (0 = Offline, 1 = Online, 3 = Safemode)."""
|
||||
payload = {"status": status}
|
||||
result = self._post("/v1/agent/find", payload)
|
||||
return pd.DataFrame(result["response"]["agents"])
|
||||
|
||||
def agent_find_by_username(self, username: str) -> pd.DataFrame:
|
||||
"""Find agents by username."""
|
||||
payload = {"username": username}
|
||||
result = self._post("/v1/agent/find", payload)
|
||||
return pd.DataFrame(result["response"]["agents"])
|
||||
|
||||
def agent_move(self, agentid: str, groupid: str) -> dict:
|
||||
"""Move an agent to a different group."""
|
||||
payload = {"agentid": agentid, "groupid": groupid}
|
||||
return self._post("/v1/agent/move", payload)
|
||||
|
||||
def agents_find_by_group(self, groupid: str) -> pd.DataFrame:
|
||||
"""Find agents by group ID."""
|
||||
payload = {"groupid": groupid}
|
||||
result = self._post("/v1/agent/find", payload)
|
||||
return pd.DataFrame(result["response"]["agents"])
|
||||
|
||||
# Hash Management
|
||||
def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
|
||||
"""Add hashes to the allowlist for a specific application."""
|
||||
payload = {"applicationid": applicationid, "hashes": hashes}
|
||||
return self._post("/v1/hash/application/add", payload)
|
||||
|
||||
def hash_query(self, hashes: List[str]) -> pd.DataFrame:
|
||||
"""Query information about specific hashes."""
|
||||
payload = {"hashes": hashes}
|
||||
result = self._post("/v1/hash/query", payload)
|
||||
return pd.DataFrame(result["response"]["results"])
|
||||
|
||||
# OTP Management
|
||||
def otp_find_active(self) -> pd.DataFrame:
|
||||
"""Find active OTPs."""
|
||||
payload = {"status": "1"}
|
||||
result = self._post("/v1/otp/usage", payload)
|
||||
return pd.DataFrame(result["response"]["otpusage"])
|
||||
|
||||
def otp_find_awaiting(self) -> pd.DataFrame:
|
||||
"""Find OTPs that are awaiting activation."""
|
||||
payload = {"status": "0"}
|
||||
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 = {
|
||||
"duration": str(duration),
|
||||
"agentid": str(agentid),
|
||||
"purpose": purpose,
|
||||
}
|
||||
result = self._post("/v1/otp/retrieve", payload)
|
||||
return result["response"]["otpcode"]
|
||||
|
||||
def otp_get_activities(self, otpid: str) -> pd.DataFrame:
|
||||
"""Retrieve activities associated with a specific OTP."""
|
||||
payload = {"otpid": otpid}
|
||||
result = self._post("/v1/otp/activities", payload)
|
||||
return pd.DataFrame(result["response"]["otpactivities"])
|
||||
|
||||
# Policy Management
|
||||
def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict:
|
||||
"""Add path exclusions to a policy group."""
|
||||
payload = {"groupid": groupid, "path": paths}
|
||||
return self._post("/v1/group/path/add", payload)
|
||||
|
||||
def policy_add_publishers(self, groupid: str, publishers: List[str]) -> dict:
|
||||
"""Add publishers to a policy group."""
|
||||
payload = {"groupid": groupid, "publisher": publishers}
|
||||
return self._post("/v1/group/publisher/add", payload)
|
||||
|
||||
def policy_clone(self, source_groupid: str, target_groupid: str) -> dict:
|
||||
"""Clone a policy from one group to another."""
|
||||
payload = {"groupid": source_groupid, "targetgroupid": target_groupid}
|
||||
return self._post("/v1/group/assign", payload)
|
||||
|
||||
def policy_find_all(self) -> pd.DataFrame:
|
||||
"""Retrieve all policy groups."""
|
||||
result = self._post("/v1/group")
|
||||
return pd.DataFrame(result["response"]["groups"])
|
||||
|
||||
def policy_list_agents(self, groupid: str) -> pd.DataFrame:
|
||||
"""List agents assigned to a specific policy group."""
|
||||
payload = {"groupid": groupid}
|
||||
result = self._post("/v1/group/agents", payload)
|
||||
return pd.DataFrame(result["response"]["agents"])
|
||||
|
||||
def policy_set_auditmode(self, groupid: str, auditmode: str) -> dict:
|
||||
"""Set audit mode for a policy group. 1=Audit, 0=Enforcement"""
|
||||
payload = {"groupid": groupid, "auditmode": auditmode}
|
||||
return self._post("/v1/group/settings/auditmode", payload)
|
||||
|
||||
# Execution History
|
||||
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]:
|
||||
"""
|
||||
Retrieve execution history logs.
|
||||
|
||||
"datefrom":"", //(Optional) Datefrom is for date range search, formatted as "YYYY-MM-DD"
|
||||
"dateto":"", //(Optional) Dateto is for date range search, formatted as "YYYY-MM-DD"
|
||||
"category":"", //(Optional) Category for filtering type
|
||||
"hostname":"", //(Optional) Hostname to filter
|
||||
"username":"admin", //(Optional) Username to filter
|
||||
"netdomain":"", //(Optional) Domain (or group) to filter
|
||||
"filename":"", //(Optional) Filename to filter
|
||||
"ppolicy":"", //(Optional) Parent Policy name to filter
|
||||
"policyname":"", //(Optional) Policy name to filter
|
||||
"policyver":"", //(Optional) Policy version to filter (e.g. "v95")
|
||||
"commandline":"", //(Optional) Commandline to filter
|
||||
"publisher":"", //(Optional) Publisher to filter
|
||||
"pprocess":"", //(Optional) Parent Process to filter
|
||||
"sha256":"", //(Optional) SHA256 hash to filter
|
||||
"contains":["hostname"], //(Optional) Contains is an array for wildcard searches on a filter
|
||||
"limit":"5" //(Optional) Limit the amount of results returned, default set to 50
|
||||
|
||||
"""
|
||||
|
||||
payload = {"datefrom": date_selected, "dateto": today, "hostname": agent_name}
|
||||
result = self._post("/v1/getexechistory", payload)
|
||||
return result["response"]["exechistory"]
|
||||
|
||||
|
||||
"""
|
||||
from services.API import AirlockAPIWrapper
|
||||
|
||||
|
||||
api = AirlockAPIWrapper(base_url="https://airlock.example.com/api", api_key="your_api_key_here")
|
||||
|
||||
#Example: Get all agents
|
||||
|
||||
agents_df = api.agent_find_all()
|
||||
print("All Agents:")
|
||||
print(agents_df)
|
||||
"""
|
||||
@@ -0,0 +1,232 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime, timedelta
|
||||
from typing import List
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from models.agent import Agent
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.selector import Selector
|
||||
from utils.utils import colorText, load_env, load_env_json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||||
agents = selectAgents(api)
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
|
||||
if not agents or not history_days:
|
||||
print(colorText("No agents selected or invalid history range.", "red"))
|
||||
return
|
||||
|
||||
historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d")
|
||||
today = datetime.now().strftime("%Y-%m-%d")
|
||||
|
||||
all_history = []
|
||||
|
||||
for agent in agents:
|
||||
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"))
|
||||
continue
|
||||
|
||||
if isinstance(exechistory, list):
|
||||
for block in exechistory:
|
||||
record = {
|
||||
"Command": block.get("commandline", "N/A"),
|
||||
"Date": block.get("datetime", "N/A"),
|
||||
"Filename": block.get("filename", "N/A"),
|
||||
"Policy Name": block.get("policyname", "N/A"),
|
||||
"Hostname": block.get("hostname", "N/A"),
|
||||
"Hash": block.get("sha256", "N/A"),
|
||||
}
|
||||
all_history.append(record)
|
||||
|
||||
if not outputjson:
|
||||
for key, value in record.items():
|
||||
print(colorText(f"{key}: {value}", "green"))
|
||||
print("\n")
|
||||
else:
|
||||
print(colorText(f"No execution history found for {agent.hostname}.", "yellow"))
|
||||
|
||||
if outputjson:
|
||||
print(json.dumps(all_history, indent=2))
|
||||
|
||||
|
||||
def findAllAgents(api):
|
||||
# Step 1: Load data from API
|
||||
policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()]
|
||||
agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()]
|
||||
|
||||
# Step 2: Create groupid → groupname map
|
||||
groupid_to_name = {policy.groupid: policy.name for policy in policies}
|
||||
|
||||
# Step 3: Enrich agents
|
||||
for agent in agents:
|
||||
agent.enrich(groupid_to_name)
|
||||
|
||||
return agents
|
||||
|
||||
|
||||
|
||||
def findAgents(api, return_dataframe):
|
||||
agents = selectAgents(api)
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
if agents:
|
||||
agent_dicts = [asdict(agent) for agent in agents]
|
||||
agent_df = pd.DataFrame(agent_dicts)
|
||||
|
||||
if return_dataframe:
|
||||
logging.debug("Returning DataFrame to caller.")
|
||||
return agent_df
|
||||
else:
|
||||
print(agent_df)
|
||||
logging.debug("Displayed DataFrame to console.")
|
||||
|
||||
# Ask user if they want to export
|
||||
user_input = 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(working_dir, filename)
|
||||
|
||||
agent_df.to_csv(file_path, index=False)
|
||||
logging.info(f"Exported DataFrame to {file_path}")
|
||||
|
||||
print(
|
||||
colorText(
|
||||
f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
logging.debug("User declined to export the DataFrame.")
|
||||
else:
|
||||
logging.warning("No agents found.")
|
||||
print("No agents matched the criteria.")
|
||||
|
||||
def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
|
||||
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.\n", "cyan"))
|
||||
print(colorText("Example:", "cyan"))
|
||||
print(colorText("H00000", "cyan"))
|
||||
print(colorText("UTN00000", "cyan"))
|
||||
print(colorText("i-hSuperSecretServer", "cyan"))
|
||||
print(colorText("u-hVenderBroke\n", "cyan"))
|
||||
|
||||
print(colorText("Paste or type your device names below:", "white"))
|
||||
|
||||
device_input_lines = []
|
||||
empty_line_count = 0
|
||||
|
||||
while True:
|
||||
line = input()
|
||||
if line.strip() == "":
|
||||
empty_line_count += 1
|
||||
if empty_line_count == 2:
|
||||
break
|
||||
else:
|
||||
empty_line_count = 0
|
||||
device_input_lines.append(line.strip())
|
||||
|
||||
device_names = [name for name in device_input_lines if name]
|
||||
if not device_names:
|
||||
logger.debug("No device names entered")
|
||||
print(colorText("⚠️ No device names entered.", "red"))
|
||||
return []
|
||||
|
||||
# Build regex pattern
|
||||
pattern = "|".join(map(re.escape, device_names))
|
||||
regex = re.compile(pattern, re.IGNORECASE)
|
||||
|
||||
# Fetch agents
|
||||
agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()]
|
||||
matched_agents = [agent for agent in agents if regex.search(agent.hostname)]
|
||||
matched_agents.sort(key=lambda agent: agent.hostname.lower())
|
||||
|
||||
# Show unmatched
|
||||
unmatched = [name for name in device_names if not any(regex.search(agent.hostname) for agent in agents)]
|
||||
if unmatched:
|
||||
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
|
||||
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
|
||||
|
||||
if not matched_agents:
|
||||
logger.debug("❌ No matching devices found.")
|
||||
print(colorText("❌ No matching devices found.", "red"))
|
||||
else:
|
||||
logger.debug(f"✅ Found {len(matched_agents)} matching device(s).")
|
||||
print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))
|
||||
|
||||
return matched_agents
|
||||
|
||||
def moveAgentToRelatedPolicy(
|
||||
api: AirlockAPIWrapper,
|
||||
agent: Agent,
|
||||
mode: str = "audit",
|
||||
):
|
||||
"""
|
||||
Moves an agent between audit and enforcement policies based on the mode.
|
||||
|
||||
Args:
|
||||
api: AirlockAPIWrapper instance.
|
||||
agent: Agent object.
|
||||
policy_relationship_map: Dict mapping enforcement → audit.
|
||||
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
|
||||
"""
|
||||
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
|
||||
if mode == "audit":
|
||||
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.")
|
||||
return
|
||||
else:
|
||||
logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.")
|
||||
return
|
||||
|
||||
elif mode == "enforcement":
|
||||
inverse_map = {v: k for k, v in policy_relationship_map.items()}
|
||||
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.")
|
||||
return
|
||||
else:
|
||||
logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.")
|
||||
return
|
||||
|
||||
else:
|
||||
logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.")
|
||||
return
|
||||
|
||||
api.agent_move(agent.agentid, target_policy)
|
||||
@@ -0,0 +1,236 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
|
||||
import datetime
|
||||
import gc
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pandas as pd
|
||||
import tqdm
|
||||
from bson import ObjectId
|
||||
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.utils import colorText, load_env, load_env_json
|
||||
from services.setup import get_base_directory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
||||
def pullPolicyExechistories(
|
||||
api: AirlockAPIWrapper,
|
||||
policy: Policy,
|
||||
type: list,
|
||||
days,
|
||||
outputjson: bool,
|
||||
):
|
||||
|
||||
file_path = f"{get_base_directory()}\\cache\\chunkinator.json"
|
||||
|
||||
# Ensure the file exists
|
||||
if not os.path.exists(file_path):
|
||||
with open(file_path, "w") as file:
|
||||
json.dump({"error": "Success", "response": {"exechistories": []}}, file)
|
||||
logger.debug(f"File '{file_path}' has been created.")
|
||||
else:
|
||||
logger.debug(f"File '{file_path}' already exists.")
|
||||
|
||||
checkpoint = str(skipback(days))
|
||||
json_output = {"error": "Success", "response": {"exechistories": []}}
|
||||
|
||||
with tqdm.tqdm(
|
||||
file=sys.stdout,
|
||||
leave=True,
|
||||
total=10000,
|
||||
desc=f"Checkpoint Progress: {checkpoint}",
|
||||
colour="blue",
|
||||
initial=1,
|
||||
) as filebar:
|
||||
with tqdm.tqdm(
|
||||
file=sys.stdout,
|
||||
leave=True,
|
||||
total=100,
|
||||
desc=f"Total of {policy} Complete: ",
|
||||
) as pbar:
|
||||
while True:
|
||||
histories = api.history_logging(
|
||||
type=type, checkpoint=checkpoint, policy= [policy.name]
|
||||
)
|
||||
|
||||
# Ensure histories is a list of dictionaries
|
||||
if not isinstance(histories, list) or not all(
|
||||
isinstance(h, dict) for h in histories
|
||||
):
|
||||
logger.error(
|
||||
"Unexpected response format from API. Expected list of dictionaries."
|
||||
)
|
||||
break
|
||||
|
||||
filebar.total = len(histories)
|
||||
|
||||
if not histories:
|
||||
break
|
||||
|
||||
for index, history_item in enumerate(histories):
|
||||
if (
|
||||
"checkpoint" not in history_item
|
||||
or "datetime" not in history_item
|
||||
):
|
||||
continue # Skip malformed entries
|
||||
|
||||
# Update checkpoint on last item
|
||||
if index == len(histories) - 1:
|
||||
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]
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
).date()
|
||||
except ValueError:
|
||||
continue # Skip if date format is invalid
|
||||
|
||||
if (
|
||||
datetime.date.today() - datetime.timedelta(days=days)
|
||||
) <= history_date:
|
||||
json_output["response"]["exechistories"].append(history_item)
|
||||
|
||||
filebar.update(1)
|
||||
filebar.refresh()
|
||||
|
||||
# Deduplicate entries
|
||||
seen = {}
|
||||
if os.path.exists(file_path):
|
||||
with open(file_path, "r") as file:
|
||||
existing_data = json.load(file)
|
||||
combined = (
|
||||
existing_data["response"]["exechistories"]
|
||||
+ json_output["response"]["exechistories"]
|
||||
)
|
||||
else:
|
||||
combined = json_output["response"]["exechistories"]
|
||||
|
||||
for entry in combined:
|
||||
key = (
|
||||
entry.get("sha256"),
|
||||
entry.get("filename"),
|
||||
entry.get("hostname"),
|
||||
)
|
||||
seen[key] = entry
|
||||
|
||||
deduplicated = list(seen.values())
|
||||
with open(file_path, "w") as file:
|
||||
json.dump(
|
||||
{
|
||||
"error": "Success",
|
||||
"response": {"exechistories": deduplicated},
|
||||
},
|
||||
file,
|
||||
)
|
||||
|
||||
json_output["response"]["exechistories"].clear()
|
||||
|
||||
# Update progress bar based on last valid item
|
||||
try:
|
||||
last_date = datetime.datetime.strptime(
|
||||
history_item["datetime"].replace(" +0000 UTC", ""), # type: ignore
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
).date()
|
||||
date_diff = datetime.date.today() - last_date
|
||||
percentage_diff = (
|
||||
((days + 10) - date_diff.days) / (days + 10)
|
||||
) * 100
|
||||
pbar.n = round(percentage_diff)
|
||||
pbar.set_description_str(f"Total of {policy} Complete: ")
|
||||
pbar.refresh()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
filebar.n = 1
|
||||
|
||||
# Final output
|
||||
with open(file_path, "r") as file:
|
||||
final_output = json.load(file)
|
||||
os.remove(file_path)
|
||||
|
||||
return json.dumps(final_output) if outputjson else None
|
||||
|
||||
|
||||
def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
|
||||
executionhist_policy = pd.DataFrame()
|
||||
exehist = pullPolicyExechistories(api, policy, type, days, True)
|
||||
if exehist is not None:
|
||||
data = json.loads(exehist)
|
||||
executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
|
||||
if not executionhist_policy.empty:
|
||||
executionhist_policy = executionhist_policy[
|
||||
[
|
||||
"datetime",
|
||||
"sha256",
|
||||
"publisher",
|
||||
"filename",
|
||||
"hostname",
|
||||
"username",
|
||||
"pprocess",
|
||||
"gprocess",
|
||||
"commandline",
|
||||
]
|
||||
]
|
||||
executionhist_policy["policy"] = policy # Add policy column here
|
||||
executionhist_policy = executionhist_policy.drop_duplicates(
|
||||
subset=["sha256", "filename", "hostname"]
|
||||
)
|
||||
executionhist_policy = executionhist_policy.sort_values(
|
||||
by=["sha256", "filename"]
|
||||
)
|
||||
logger.debug( f"Staging of Execution history for policy: {policy} is complete")
|
||||
print(
|
||||
colorText(
|
||||
f"Staging of Execution history for policy: {policy} is complete",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
del data
|
||||
del exehist
|
||||
gc.collect()
|
||||
return executionhist_policy
|
||||
|
||||
|
||||
def skipback(days):
|
||||
"""
|
||||
Generate a MongoDB ObjectId for a given number of days ago from today.
|
||||
"""
|
||||
adjusted_days = days
|
||||
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
|
||||
days=adjusted_days
|
||||
)
|
||||
timestamp = int(date_days_ago.timestamp())
|
||||
hex_timestamp = format(timestamp, "08x")
|
||||
objectid_hex = hex_timestamp + "0000000000000000"
|
||||
return ObjectId(objectid_hex)
|
||||
|
||||
|
||||
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
|
||||
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
for enforcement_policy, audit_policy in policy_relationship_map.items():
|
||||
api.policy_clone(enforcement_policy, audit_policy)
|
||||
api.policy_set_auditmode(audit_policy, "1")
|
||||
@@ -0,0 +1,352 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List
|
||||
|
||||
import schedule
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# TODO - move this File where all jobs are persisted
|
||||
JOBS_FILE = f"{os.getenv('WORKING_DIR')}\\scheduling\\jobs.json"
|
||||
|
||||
# Ensure directory exists
|
||||
os.makedirs(os.path.dirname(JOBS_FILE), exist_ok=True)
|
||||
|
||||
# Registry of functions that can be scheduled
|
||||
FUNCTION_MAP: Dict[str, Callable] = {}
|
||||
|
||||
# -------------------------------
|
||||
# Function Registration
|
||||
# -------------------------------
|
||||
|
||||
|
||||
def register_function(name: str, func: Callable):
|
||||
"""
|
||||
Register a function so it can be called by name later.
|
||||
Example:
|
||||
register_function("say_hello", say_hello)
|
||||
"""
|
||||
FUNCTION_MAP[name] = func
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Persistence Helpers
|
||||
# -------------------------------
|
||||
|
||||
|
||||
def load_jobs() -> List[Dict[str, Any]]:
|
||||
"""Load jobs from the JSON file, or return [] if none exist."""
|
||||
if not os.path.exists(JOBS_FILE):
|
||||
return []
|
||||
with open(JOBS_FILE, "r") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _atomic_save(path: str, data: Any):
|
||||
"""Write JSON atomically to avoid partial writes."""
|
||||
tmp = f"{path}.tmp"
|
||||
with open(tmp, "w") as f:
|
||||
json.dump(data, f, indent=4)
|
||||
os.replace(tmp, path)
|
||||
|
||||
|
||||
def save_jobs(jobs: List[Dict[str, Any]]):
|
||||
"""Save jobs to the JSON file (overwrite)."""
|
||||
_atomic_save(JOBS_FILE, jobs)
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Uniqueness Helpers
|
||||
# -------------------------------
|
||||
|
||||
|
||||
def job_in_store(job_id: str) -> bool:
|
||||
"""Check if a job id exists in the persisted JSON file."""
|
||||
return any(j.get("id") == job_id for j in load_jobs())
|
||||
|
||||
|
||||
def job_in_scheduler(job_id: str) -> bool:
|
||||
"""
|
||||
Check if a job with this tag exists in the in-memory scheduler.
|
||||
Uses schedule.get_jobs(tag=...) if available, otherwise scans tags.
|
||||
"""
|
||||
try:
|
||||
jobs = schedule.get_jobs(tag=job_id) # schedule >= 1.2.0
|
||||
return len(jobs) > 0
|
||||
except TypeError:
|
||||
# Fallback for older versions
|
||||
return any(job_id in getattr(j, "tags", set()) for j in schedule.jobs)
|
||||
|
||||
|
||||
def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool:
|
||||
"""
|
||||
Ensure the job_id is unique across persistence and in-memory schedule.
|
||||
on_conflict:
|
||||
- "error": raise ValueError if exists.
|
||||
- "skip" : print and return False.
|
||||
- "replace": remove existing (in-memory + JSON), then continue.
|
||||
"""
|
||||
exists = job_in_store(job_id) or job_in_scheduler(job_id)
|
||||
if not exists:
|
||||
return True
|
||||
if on_conflict == "error":
|
||||
raise ValueError(f"Job id '{job_id}' already exists.")
|
||||
elif on_conflict == "skip":
|
||||
logger.info(f"Job '{job_id}' already exists. Skipping creation.")
|
||||
return False
|
||||
elif on_conflict == "replace":
|
||||
# Clear from scheduler
|
||||
schedule.clear(job_id)
|
||||
# Remove from persistence
|
||||
jobs = [j for j in load_jobs() if j.get("id") != job_id]
|
||||
save_jobs(jobs)
|
||||
return True
|
||||
else:
|
||||
raise ValueError(f"Unsupported on_conflict policy: {on_conflict}")
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Internal scheduling (no persistence)
|
||||
# -------------------------------
|
||||
|
||||
|
||||
def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None):
|
||||
args = args or []
|
||||
kwargs = kwargs or {}
|
||||
|
||||
def job_wrapper():
|
||||
"""Executes the job once, then removes it."""
|
||||
if func_name not in FUNCTION_MAP:
|
||||
logger.error(f"Function '{func_name}' is not registered.")
|
||||
return
|
||||
FUNCTION_MAP[func_name](*args, **kwargs)
|
||||
# Remove from persistence
|
||||
jobs = load_jobs()
|
||||
jobs = [j for j in jobs if j["id"] != job_id]
|
||||
save_jobs(jobs)
|
||||
# Clear from in-memory schedule
|
||||
schedule.clear(job_id)
|
||||
|
||||
delay_seconds = run_at_timestamp - time.time()
|
||||
if delay_seconds <= 0:
|
||||
logger.info(f"Job {job_id} scheduled in the past. Skipping.")
|
||||
return
|
||||
# Schedule via schedule library
|
||||
schedule.every(int(delay_seconds)).seconds.do(job_wrapper).tag(job_id)
|
||||
|
||||
|
||||
def _schedule_recurring(
|
||||
job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None
|
||||
):
|
||||
args = args or []
|
||||
kwargs = kwargs or {}
|
||||
|
||||
def job_wrapper():
|
||||
if func_name not in FUNCTION_MAP:
|
||||
logger.error(f"Function '{func_name}' is not registered.")
|
||||
return
|
||||
FUNCTION_MAP[func_name](*args, **kwargs)
|
||||
|
||||
if unit == "seconds":
|
||||
schedule.every(interval).seconds.do(job_wrapper).tag(job_id)
|
||||
elif unit == "minutes":
|
||||
schedule.every(interval).minutes.do(job_wrapper).tag(job_id)
|
||||
elif unit == "hours":
|
||||
schedule.every(interval).hours.do(job_wrapper).tag(job_id)
|
||||
elif unit == "days":
|
||||
schedule.every(interval).days.do(job_wrapper).tag(job_id)
|
||||
else:
|
||||
raise ValueError(f"Unsupported unit: {unit}")
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Public APIs (with uniqueness + persistence)
|
||||
# -------------------------------
|
||||
|
||||
|
||||
def run_once_job(
|
||||
job_id: str,
|
||||
func_name: str,
|
||||
run_at_timestamp: float,
|
||||
args=None,
|
||||
kwargs=None,
|
||||
*,
|
||||
replace: bool = False,
|
||||
persist: bool = True,
|
||||
):
|
||||
"""
|
||||
Schedule a job to run once at a specific timestamp.
|
||||
replace: if True, replace existing job with same id; otherwise print and skip.
|
||||
persist: if False, do not write to JSON (used by reload_jobs()).
|
||||
"""
|
||||
if persist:
|
||||
policy = "replace" if replace else "skip"
|
||||
if not ensure_unique(job_id, on_conflict=policy):
|
||||
return
|
||||
elif job_in_scheduler(job_id):
|
||||
schedule.clear(job_id)
|
||||
_schedule_once(job_id, func_name, run_at_timestamp, args, kwargs)
|
||||
if persist:
|
||||
jobs = [j for j in load_jobs() if j["id"] != job_id]
|
||||
jobs.append(
|
||||
{
|
||||
"id": job_id,
|
||||
"type": "once",
|
||||
"run_at": run_at_timestamp,
|
||||
"function": func_name,
|
||||
"args": args or [],
|
||||
"kwargs": kwargs or {},
|
||||
}
|
||||
)
|
||||
save_jobs(jobs)
|
||||
|
||||
|
||||
def recurring_job(
|
||||
job_id: str,
|
||||
func_name: str,
|
||||
interval: int,
|
||||
unit: str,
|
||||
args=None,
|
||||
kwargs=None,
|
||||
*,
|
||||
replace: bool = False,
|
||||
persist: bool = True,
|
||||
):
|
||||
"""
|
||||
Schedule a recurring job.
|
||||
replace: if True, replace existing job with same id; otherwise print and skip.
|
||||
persist: if False, do not write to JSON (used by reload_jobs()).
|
||||
"""
|
||||
if persist:
|
||||
policy = "replace" if replace else "skip"
|
||||
if not ensure_unique(job_id, on_conflict=policy):
|
||||
return
|
||||
elif job_in_scheduler(job_id):
|
||||
schedule.clear(job_id)
|
||||
_schedule_recurring(job_id, func_name, interval, unit, args, kwargs)
|
||||
if persist:
|
||||
jobs = [j for j in load_jobs() if j["id"] != job_id]
|
||||
jobs.append(
|
||||
{
|
||||
"id": job_id,
|
||||
"type": "recurring",
|
||||
"interval": interval,
|
||||
"unit": unit,
|
||||
"function": func_name,
|
||||
"args": args or [],
|
||||
"kwargs": kwargs or {},
|
||||
}
|
||||
)
|
||||
save_jobs(jobs)
|
||||
|
||||
|
||||
def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float = 1.0):
|
||||
"""
|
||||
Find all jobs whose ID contains the given PID substring and reschedule them to run sooner.
|
||||
"""
|
||||
jobs = load_jobs()
|
||||
matched_jobs = [job for job in jobs if pid_substring in job.get("id", "")]
|
||||
|
||||
if not matched_jobs:
|
||||
logger.info(f"No jobs found containing PID substring '{pid_substring}'.")
|
||||
return
|
||||
|
||||
logger.info(f"Found {len(matched_jobs)} job(s) containing '{pid_substring}':")
|
||||
|
||||
for job in matched_jobs:
|
||||
job_id = job["id"]
|
||||
logger.debug(f" - Prioritizing job: {job_id}")
|
||||
|
||||
# Clear existing job from scheduler
|
||||
schedule.clear(job_id)
|
||||
|
||||
# Reschedule based on job type
|
||||
if job["type"] == "once":
|
||||
run_once_job(
|
||||
job_id,
|
||||
job["function"],
|
||||
time.time() + new_delay_seconds,
|
||||
job.get("args"),
|
||||
job.get("kwargs"),
|
||||
replace=True,
|
||||
persist=True,
|
||||
)
|
||||
elif job["type"] == "recurring":
|
||||
recurring_job(
|
||||
job_id,
|
||||
job["function"],
|
||||
job["interval"],
|
||||
job["unit"],
|
||||
job.get("args"),
|
||||
job.get("kwargs"),
|
||||
replace=True,
|
||||
persist=True,
|
||||
)
|
||||
else:
|
||||
logger.warning(f"Unknown job type for job '{job_id}'")
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Reload Saved Jobs
|
||||
# -------------------------------
|
||||
|
||||
|
||||
def reload_jobs():
|
||||
"""Reload jobs from JSON and reschedule them (no re-persist)."""
|
||||
jobs = load_jobs()
|
||||
for job in jobs:
|
||||
if job["type"] == "once":
|
||||
if job["run_at"] > time.time():
|
||||
run_once_job(
|
||||
job["id"],
|
||||
job["function"],
|
||||
job["run_at"],
|
||||
job.get("args"),
|
||||
job.get("kwargs"),
|
||||
persist=False,
|
||||
)
|
||||
elif job["type"] == "recurring":
|
||||
recurring_job(
|
||||
job["id"],
|
||||
job["function"],
|
||||
job["interval"],
|
||||
job["unit"],
|
||||
job.get("args"),
|
||||
job.get("kwargs"),
|
||||
persist=False,
|
||||
)
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Scheduler Loop
|
||||
# -------------------------------
|
||||
|
||||
|
||||
def start_scheduler():
|
||||
"""
|
||||
Start the scheduler loop (blocking).
|
||||
Call this once in main to begin.
|
||||
"""
|
||||
try:
|
||||
while True:
|
||||
schedule.run_pending()
|
||||
time.sleep(0.5)
|
||||
except KeyboardInterrupt:
|
||||
logger.critical("Scheduler stopped.")
|
||||
@@ -0,0 +1,148 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
from getpass import getpass
|
||||
|
||||
import keyring
|
||||
from cryptography.hazmat.primitives import hashes
|
||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
||||
|
||||
# Constants
|
||||
KDF_ITERATIONS = 200_000
|
||||
SALT_SIZE = 16 # 128-bit Salt
|
||||
NONCE_SIZE = 12 # AES-GCM
|
||||
KEY_SIZE = 32 # AES-256
|
||||
|
||||
|
||||
def _derive_key(password: bytes, salt: bytes) -> bytes:
|
||||
kdf = PBKDF2HMAC(
|
||||
algorithm=hashes.SHA256(),
|
||||
length=KEY_SIZE,
|
||||
salt=salt,
|
||||
iterations=KDF_ITERATIONS,
|
||||
)
|
||||
return kdf.derive(password)
|
||||
|
||||
|
||||
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}")
|
||||
|
||||
|
||||
def store_api_key(service: str, username: str, api_key: str, password: str):
|
||||
configure_keyring_backend()
|
||||
salt = os.urandom(SALT_SIZE)
|
||||
key = _derive_key(password.encode(), salt)
|
||||
aesgcm = AESGCM(key)
|
||||
nonce = os.urandom(NONCE_SIZE)
|
||||
ct = aesgcm.encrypt(nonce, api_key.encode(), associated_data=None)
|
||||
blob = salt + nonce + ct
|
||||
b64 = base64.b64encode(blob).decode()
|
||||
keyring.set_password(service, username, b64)
|
||||
|
||||
|
||||
def retrieve_api_key(service: str, username: str, password: str) -> str:
|
||||
configure_keyring_backend()
|
||||
b64 = keyring.get_password(service, username)
|
||||
if b64 is None:
|
||||
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:]
|
||||
key = _derive_key(password.encode(), salt)
|
||||
aesgcm = AESGCM(key)
|
||||
pt = aesgcm.decrypt(nonce, ct, associated_data=None)
|
||||
return pt.decode()
|
||||
|
||||
|
||||
def api_key_exists(service: str, username: str) -> bool:
|
||||
configure_keyring_backend()
|
||||
return keyring.get_password(service, username) is not None
|
||||
|
||||
|
||||
def check_password_complexity(password: str) -> bool:
|
||||
if len(password) < 12:
|
||||
return False
|
||||
if not re.search(r"[A-Z]", password):
|
||||
return False
|
||||
if not re.search(r"[a-z]", password):
|
||||
return False
|
||||
if not re.search(r"[0-9]", password):
|
||||
return False
|
||||
if not re.search(r"[^A-Za-z0-9]", password):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def getAPI(USERNAME, SERVICE_NAME):
|
||||
logging.debug(
|
||||
f"Checking for stored API key for user '{USERNAME}' in service '{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: ")
|
||||
try:
|
||||
apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password)
|
||||
logging.debug("API key successfully retrieved.")
|
||||
return apikey
|
||||
except Exception as e:
|
||||
logging.warning(f"Attempt {attempt} failed: {str(e)}")
|
||||
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 = input(f"🔑 No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip()
|
||||
|
||||
while True:
|
||||
password = getpass("🔐 Create a password to encrypt your API key: ")
|
||||
if check_password_complexity(password):
|
||||
try:
|
||||
store_api_key(SERVICE_NAME, USERNAME, api_key, password)
|
||||
logging.info("API key stored securely.")
|
||||
break
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to store API key: {e}")
|
||||
break
|
||||
else:
|
||||
print("❌ Password does not meet complexity requirements. Try again.")
|
||||
return api_key
|
||||
|
||||
class APIKeyManager:
|
||||
_api_key = None
|
||||
|
||||
@classmethod
|
||||
def load(cls, service: str, username: str, password: str):
|
||||
cls._api_key = retrieve_api_key(service, username, password)
|
||||
|
||||
@classmethod
|
||||
def get(cls) -> str:
|
||||
if cls._api_key is None:
|
||||
raise ValueError("API key not loaded. Call APIKeyManager.load() first.")
|
||||
return cls._api_key
|
||||
@@ -0,0 +1,121 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
|
||||
import logging
|
||||
from typing import Any, List, Optional, Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Selector:
|
||||
@staticmethod
|
||||
def select_objects(
|
||||
objects: List[Any],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False
|
||||
) -> Union[Optional[Any], List[Any]]:
|
||||
if not objects:
|
||||
logger.warning("No objects available for selection.")
|
||||
return None
|
||||
|
||||
# Sort objects alphabetically by their 'name' attribute
|
||||
sorted_objects = sorted(objects, key=lambda obj: getattr(obj, "name", str(obj)).lower())
|
||||
|
||||
# Display in 4 columns with extra spacing
|
||||
num_columns = 4
|
||||
rows = (len(sorted_objects) + num_columns - 1) // num_columns
|
||||
print("\nAvailable Choices:")
|
||||
for row in range(rows):
|
||||
line = ""
|
||||
for col in range(num_columns):
|
||||
idx = row + col * rows
|
||||
if idx < len(sorted_objects):
|
||||
obj = sorted_objects[idx]
|
||||
name = getattr(obj, "name", str(obj))
|
||||
line += f"{idx + 1}: {name:<30}"
|
||||
print(line)
|
||||
|
||||
selected = []
|
||||
|
||||
if allow_multiple:
|
||||
while True:
|
||||
choice = input("Select an object by number (or Q to finish): ").strip().lower()
|
||||
if choice == "q":
|
||||
break
|
||||
try:
|
||||
index = int(choice)
|
||||
if 1 <= index <= len(sorted_objects):
|
||||
obj = sorted_objects[index - 1]
|
||||
if obj not in selected:
|
||||
selected.append(obj)
|
||||
if prompt_each:
|
||||
logger.info(f"Selected: {getattr(obj, 'name', str(obj))}")
|
||||
else:
|
||||
logger.warning("Object already selected.")
|
||||
else:
|
||||
logger.warning("Selection out of range. Try again.")
|
||||
except ValueError:
|
||||
logger.warning("Invalid input. Enter a number or 'Q' to quit.")
|
||||
return selected if selected else None
|
||||
else:
|
||||
try:
|
||||
choice = int(input("Select one object by number: "))
|
||||
if 1 <= choice <= len(sorted_objects):
|
||||
selected_obj = sorted_objects[choice - 1]
|
||||
logger.info(f"Selected: {getattr(selected_obj, 'name', str(selected_obj))}")
|
||||
return selected_obj
|
||||
else:
|
||||
logger.warning("Selection out of range.")
|
||||
except ValueError:
|
||||
logger.warning("Invalid input.")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def select_value(
|
||||
prompt: str,
|
||||
value_type: type = int,
|
||||
valid_range: Optional[tuple] = None,
|
||||
allow_quit: bool = False
|
||||
) -> Optional[Any]:
|
||||
while True:
|
||||
user_input = input(prompt).strip().lower()
|
||||
if allow_quit and user_input == "q":
|
||||
logger.info("User opted to quit value selection.")
|
||||
return None
|
||||
try:
|
||||
value = value_type(user_input)
|
||||
if valid_range:
|
||||
min_val, max_val = valid_range
|
||||
if not (min_val <= value <= max_val):
|
||||
logger.warning(f"Value out of range ({min_val}–{max_val}).")
|
||||
continue
|
||||
logger.info(f"User selected value: {value}")
|
||||
return value
|
||||
except ValueError:
|
||||
logger.warning(f"Invalid input. Expected a {value_type.__name__}.")
|
||||
|
||||
@staticmethod
|
||||
def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
|
||||
while True:
|
||||
response = input(prompt).strip().lower()
|
||||
if response in ["y", "yes"]:
|
||||
logger.info("User confirmed action.")
|
||||
return True
|
||||
elif response in ["n", "no"]:
|
||||
logger.info("User declined action.")
|
||||
return False
|
||||
else:
|
||||
logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
|
||||
@@ -0,0 +1,178 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
|
||||
|
||||
import json
|
||||
import logging
|
||||
import logging.handlers
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv, set_key
|
||||
|
||||
PROTECTED_KEYS = [
|
||||
"APPNAME",
|
||||
"PATH_EXCLUSION_CONST",
|
||||
"MIN_FILES_FOR_PATH",
|
||||
"VT_THREAT_TOLERANCE",
|
||||
"POLICY_MAP_ENF_AUD"
|
||||
]
|
||||
|
||||
def get_base_directory() -> Path:
|
||||
system = platform.system()
|
||||
home = Path.home()
|
||||
if system == 'Windows':
|
||||
return Path(os.getenv('APPDATA', home / 'AppData' / 'Roaming')) / "AirlockTools"
|
||||
elif system == 'Darwin':
|
||||
return home / 'Library' / 'Application Support' / "AirlockTools"
|
||||
else:
|
||||
return home / '.local' / 'share' / "AirlockTools"
|
||||
|
||||
def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
|
||||
log_file = log_dir / "airlocktools.log"
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(getattr(logging, log_level.upper(), logging.DEBUG))
|
||||
|
||||
if not logger.handlers:
|
||||
file_handler = logging.handlers.RotatingFileHandler(
|
||||
log_file, maxBytes=5_000_000, backupCount=5, encoding='utf-8'
|
||||
)
|
||||
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
|
||||
logger.addHandler(file_handler)
|
||||
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setFormatter(logging.Formatter('%(levelname)s - %(message)s'))
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
if platform.system() == "Windows":
|
||||
try:
|
||||
event_handler = logging.handlers.NTEventLogHandler("AirlockTools")
|
||||
event_handler.setLevel(logging.CRITICAL)
|
||||
event_handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s'))
|
||||
logger.addHandler(event_handler)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not attach Windows Event Log handler: {e}")
|
||||
|
||||
logger.debug("Logging configured.")
|
||||
|
||||
def get_system_config_path() -> Path:
|
||||
base_path = Path(getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__))))
|
||||
return base_path.parent / "system_config.json"
|
||||
|
||||
def load_system_config() -> dict:
|
||||
try:
|
||||
config_path = get_system_config_path()
|
||||
with open(config_path, "r") as f:
|
||||
return json.load(f)
|
||||
except FileNotFoundError:
|
||||
logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
|
||||
return {
|
||||
"APPNAME": "AirlockTools",
|
||||
"LOG_LEVEL": "DEBUG",
|
||||
"PATH_EXCLUSION_CONST": 4,
|
||||
"MIN_FILES_FOR_PATH": 4,
|
||||
"VT_THREAT_TOLERANCE": 4,
|
||||
"POLICY_MAP_ENF_AUD": {
|
||||
"enforced_id": "audit_id"
|
||||
}
|
||||
}
|
||||
|
||||
def load_user_config(config_dir: Path) -> dict:
|
||||
user_config_path = config_dir / "user_config.json"
|
||||
if not user_config_path.exists():
|
||||
default_user_config = {
|
||||
"URL": "",
|
||||
"LOG_LEVEL": "INFO"
|
||||
}
|
||||
with open(user_config_path, "w") as f:
|
||||
json.dump(default_user_config, f, indent=4)
|
||||
logging.debug(f"Created user config at {user_config_path}")
|
||||
with open(user_config_path, "r") as f:
|
||||
return json.load(f)
|
||||
|
||||
def write_config_to_env(config: dict, env_path: Path):
|
||||
for key, value in config.items():
|
||||
try:
|
||||
serialized = json.dumps(value) if isinstance(value, (list, dict)) else str(value)
|
||||
set_key(env_path, key, serialized)
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to write {key} to .env: {e}")
|
||||
|
||||
def setup() -> Path:
|
||||
base_dir = get_base_directory()
|
||||
dirs = {
|
||||
'config': base_dir / 'config',
|
||||
'cache': base_dir / 'cache',
|
||||
'logs': base_dir / 'logs',
|
||||
}
|
||||
|
||||
for name, path in dirs.items():
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
logging.debug(f"{name.capitalize()} directory ensured at: {path}")
|
||||
|
||||
system_config = load_system_config()
|
||||
configure_logging(dirs['logs'], system_config.get("LOG_LEVEL", "DEBUG"))
|
||||
|
||||
env_path = base_dir / ".env"
|
||||
if not env_path.exists():
|
||||
env_path.touch()
|
||||
load_dotenv(dotenv_path=env_path, override=True)
|
||||
|
||||
working_dir = Path(os.getenv("WORKING_DIR") or (base_dir / "data"))
|
||||
working_dir.mkdir(parents=True, exist_ok=True)
|
||||
set_key(env_path, "WORKING_DIR", str(working_dir))
|
||||
os.environ["WORKING_DIR"] = str(working_dir)
|
||||
logging.debug(f"Working directory set to: {working_dir}")
|
||||
|
||||
folders_structure = {
|
||||
"Approved": [],
|
||||
"Needs_Review": ["Review_First", "Review_Second", "HTML"],
|
||||
"Preflight": ["HTML"],
|
||||
"Archived": ["HTML"],
|
||||
"Scheduling": []
|
||||
}
|
||||
|
||||
for folder_name, subfolders in folders_structure.items():
|
||||
folder_path = working_dir / folder_name
|
||||
folder_path.mkdir(parents=True, exist_ok=True)
|
||||
logging.debug(f"'{folder_name}' folder ensured at: {folder_path}")
|
||||
for subfolder in subfolders:
|
||||
subfolder_path = folder_path / subfolder
|
||||
subfolder_path.mkdir(parents=True, exist_ok=True)
|
||||
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
|
||||
|
||||
user_config = load_user_config(dirs['config'])
|
||||
merged_config = {**system_config, **user_config}
|
||||
|
||||
for key in PROTECTED_KEYS:
|
||||
merged_config[key] = system_config.get(key, "")
|
||||
|
||||
# ✅ URL resolution order: system_config → .env → user prompt
|
||||
url = system_config.get("URL")
|
||||
if not url:
|
||||
url = os.getenv("URL")
|
||||
if not url:
|
||||
url = input("🌐 Enter the service URL (e.g., https://example.com/api): ").strip()
|
||||
|
||||
merged_config["URL"] = url
|
||||
set_key(env_path, "URL", url)
|
||||
os.environ["URL"] = url
|
||||
logging.debug(f"Service URL set to: {url}")
|
||||
|
||||
write_config_to_env(merged_config, env_path)
|
||||
|
||||
return working_dir
|
||||
Reference in New Issue
Block a user