RC 1.1.1
This commit is contained in:
+1
-1
@@ -10,5 +10,5 @@ jobs.json
|
||||
securitytest.py
|
||||
*.toml
|
||||
system_config.json
|
||||
Devel_unused/
|
||||
Development/
|
||||
AirlockTools_client*/
|
||||
@@ -0,0 +1,46 @@
|
||||
import sys
|
||||
import os
|
||||
import logging
|
||||
|
||||
|
||||
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
|
||||
sys.path.append(project_root)
|
||||
|
||||
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.security import getAPI
|
||||
|
||||
logger = logging.getLogger()
|
||||
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
def main():
|
||||
try:
|
||||
url = "https://172.17.22.240:3129"
|
||||
username = os.getenv("USERNAME")
|
||||
|
||||
if not url:
|
||||
raise ValueError("Missing URL in environment variables.")
|
||||
if not username:
|
||||
raise ValueError("Missing USERNAME in environment variables.")
|
||||
|
||||
logger.debug(f"Retrieved URL: {url}")
|
||||
logger.debug(f"Retrieved Username: {username}")
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"Configuration error: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
|
||||
if username:
|
||||
api = AirlockAPIWrapper(
|
||||
base_url= url,
|
||||
api_key = getAPI(username, "AirlockTools"), # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
source = ""
|
||||
target = ""
|
||||
response = api.policy_clone(source, target)
|
||||
print(response)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,44 @@
|
||||
import sys
|
||||
import os
|
||||
import logging
|
||||
|
||||
|
||||
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', ".."))
|
||||
sys.path.append(project_root)
|
||||
|
||||
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.security import getAPI
|
||||
|
||||
logger = logging.getLogger()
|
||||
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
def main():
|
||||
try:
|
||||
url = "https://172.17.22.240:3129"
|
||||
username = os.getenv("USERNAME")
|
||||
|
||||
if not url:
|
||||
raise ValueError("Missing URL in environment variables.")
|
||||
if not username:
|
||||
raise ValueError("Missing USERNAME in environment variables.")
|
||||
|
||||
logger.debug(f"Retrieved URL: {url}")
|
||||
logger.debug(f"Retrieved Username: {username}")
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"Configuration error: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
|
||||
if username:
|
||||
api = AirlockAPIWrapper(
|
||||
base_url= url,
|
||||
api_key = getAPI(username, "AirlockTools"), # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
response = api.policy_find_all()
|
||||
response.to_csv("All_Policies.csv", index=False)
|
||||
|
||||
print(response)
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,75 @@
|
||||
import sys
|
||||
import os
|
||||
import logging
|
||||
|
||||
|
||||
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
|
||||
sys.path.append(project_root)
|
||||
|
||||
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.security import getAPI
|
||||
|
||||
logger = logging.getLogger()
|
||||
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
def main():
|
||||
try:
|
||||
url = "https://172.17.22.240:3129"
|
||||
username = os.getenv("USERNAME")
|
||||
|
||||
if not url:
|
||||
raise ValueError("Missing URL in environment variables.")
|
||||
if not username:
|
||||
raise ValueError("Missing USERNAME in environment variables.")
|
||||
|
||||
logger.debug(f"Retrieved URL: {url}")
|
||||
logger.debug(f"Retrieved Username: {username}")
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"Configuration error: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
|
||||
if username:
|
||||
api = AirlockAPIWrapper(
|
||||
base_url= url,
|
||||
api_key = getAPI(username, "AirlockTools"), # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
"""
|
||||
Available script types are:
|
||||
|
||||
"batch",
|
||||
"powershell",
|
||||
"command",
|
||||
"vbscript",
|
||||
"javascript"
|
||||
,"windowsinstaller",
|
||||
"htmlapplication",
|
||||
"javaapplication",
|
||||
"windowsscriptcomponent",
|
||||
"compiledhtml",
|
||||
"shellscript",
|
||||
"dylib",
|
||||
"python"
|
||||
|
||||
"""
|
||||
groupid = "5aebf6a0-1d67-47b4-9c5f-2866ffca5671" #AT Testing
|
||||
script_custom = 1
|
||||
scripts_audit = [
|
||||
"batch",
|
||||
"powershell",
|
||||
"command",
|
||||
"vbscript",
|
||||
"javascript",
|
||||
"windowsinstaller",
|
||||
"javaapplication",
|
||||
"python"
|
||||
]
|
||||
|
||||
scripts_disabled = ["compiledhtml", "htmlapplication", "shellscript", "dylib","windowsscriptcomponent"]
|
||||
scripts_respect = []
|
||||
|
||||
response = api.policy_set_script_custom(groupid, script_custom, scripts_audit, scripts_disabled, scripts_respect)
|
||||
print(response)
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,39 @@
|
||||
import os
|
||||
import json
|
||||
from utils.configmanager import (
|
||||
load_protected_config,
|
||||
get_protected_value,
|
||||
get_protected_json,
|
||||
PROTECTED_KEYS
|
||||
)
|
||||
from utils.setup import setup
|
||||
|
||||
def test_protected_config():
|
||||
print("🔒 Testing protected config loading...")
|
||||
protected = load_protected_config()
|
||||
assert isinstance(protected, dict), "Protected config should be a dictionary"
|
||||
for key in PROTECTED_KEYS:
|
||||
assert key in protected, f"Missing protected key: {key}"
|
||||
print(f"✔ {key} = {protected[key]}")
|
||||
|
||||
def test_json_parsing():
|
||||
print("\n🧪 Testing JSON parsing for POLICY_MAP_ENF_AUD...")
|
||||
policy = get_protected_json("POLICY_MAP_ENF_AUD")
|
||||
assert isinstance(policy, dict), "POLICY_MAP_ENF_AUD should be a dictionary"
|
||||
print(f"✔ POLICY_MAP_ENF_AUD = {json.dumps(policy, indent=2)}")
|
||||
|
||||
def test_setup_env():
|
||||
print("\n⚙️ Running setup() to validate environment setup...")
|
||||
working_dir = setup()
|
||||
assert working_dir.exists(), "Working directory should exist"
|
||||
print(f"✔ Working directory: {working_dir}")
|
||||
|
||||
print("\n🌍 Checking .env values (excluding protected)...")
|
||||
for key in os.environ:
|
||||
if key not in PROTECTED_KEYS:
|
||||
print(f"🔧 {key} = {os.environ[key]}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_protected_config()
|
||||
test_json_parsing()
|
||||
test_setup_env()
|
||||
@@ -0,0 +1,170 @@
|
||||
"""
|
||||
datastore.py
|
||||
|
||||
Utility module for saving/loading objects to/from JSON and inserting/updating/reading them in SQLite.
|
||||
|
||||
Supports any class with:
|
||||
- a `to_dict()` method
|
||||
- a constructor accepting `**kwargs`
|
||||
- optionally, a `from_dict()` method
|
||||
|
||||
Author: Brandon Wickline, James Brotosky
|
||||
License: GNU Affero General Public License v3.0
|
||||
"""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import logging
|
||||
from typing import Type, List, TypeVar
|
||||
|
||||
# Setup logger
|
||||
logger = logging.getLogger(__name__)
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
# Generic type variable for typing
|
||||
T = TypeVar('T')
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# JSON SAVE
|
||||
# -------------------------------------------------------------------
|
||||
def save_to_json(obj_list: List[T], filepath: str) -> None:
|
||||
"""
|
||||
Save a list of objects to a JSON file using their `to_dict()` method.
|
||||
|
||||
Args:
|
||||
obj_list (List[T]): List of objects to serialize.
|
||||
filepath (str): Path to the output JSON file.
|
||||
|
||||
Example:
|
||||
save_to_json(policies, "policies.json")
|
||||
"""
|
||||
try:
|
||||
with open(filepath, 'w', encoding='utf-8') as f:
|
||||
json.dump([obj.to_dict() for obj in obj_list], f, ensure_ascii=False, indent=4) # pyright: ignore[reportAttributeAccessIssue]
|
||||
logger.info(f"Saved {len(obj_list)} objects to {filepath}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save to {filepath}: {e}")
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# JSON LOAD
|
||||
# -------------------------------------------------------------------
|
||||
def load_from_json(cls: Type[T], filepath: str) -> List[T]:
|
||||
"""
|
||||
Load a list of objects from a JSON file and instantiate them using the class constructor
|
||||
or a `from_dict()` method if available.
|
||||
|
||||
Args:
|
||||
cls (Type[T]): Class type to instantiate.
|
||||
filepath (str): Path to the input JSON file.
|
||||
|
||||
Returns:
|
||||
List[T]: List of instantiated objects.
|
||||
|
||||
Example:
|
||||
loaded_agents = load_from_json(Agent, "agents.json")
|
||||
"""
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
data_list = json.load(f)
|
||||
logger.info(f"Loaded {len(data_list)} records from {filepath}")
|
||||
if hasattr(cls, "from_dict"):
|
||||
return [cls.from_dict(data) for data in data_list] # pyright: ignore[reportAttributeAccessIssue]
|
||||
return [cls(**data) for data in data_list]
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load from {filepath}: {e}")
|
||||
return []
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# SQLITE INSERT OR UPDATE
|
||||
# -------------------------------------------------------------------
|
||||
def insert_or_update_objects_to_sqlite(obj_list: List[T], table_name: str, db_path: str, primary_key: str) -> None:
|
||||
"""
|
||||
Insert or update a list of objects into a SQLite table.
|
||||
Uses `ON CONFLICT(primary_key) DO UPDATE` for upsert behavior.
|
||||
|
||||
Args:
|
||||
obj_list (List[T]): List of objects with `to_dict()` method.
|
||||
table_name (str): Name of the SQLite table.
|
||||
db_path (str): Path to the SQLite database file.
|
||||
primary_key (str): Field name to use as the primary key.
|
||||
|
||||
Example:
|
||||
insert_or_update_objects_to_sqlite(hashes, "hashes", "data_store.db", primary_key="sha256")
|
||||
"""
|
||||
if not obj_list:
|
||||
logger.warning("No objects to insert or update into SQLite.")
|
||||
return
|
||||
|
||||
try:
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
sample_dict = obj_list[0].to_dict() # pyright: ignore[reportAttributeAccessIssue]
|
||||
columns = ', '.join(sample_dict.keys())
|
||||
placeholders = ', '.join(['?'] * len(sample_dict))
|
||||
update_clause = ', '.join([f"{key}=excluded.{key}" for key in sample_dict.keys() if key != primary_key])
|
||||
|
||||
# Create table if it doesn't exist
|
||||
create_stmt = f"""
|
||||
CREATE TABLE IF NOT EXISTS {table_name} (
|
||||
{', '.join([f"{key} TEXT" for key in sample_dict.keys()])},
|
||||
PRIMARY KEY ({primary_key})
|
||||
)
|
||||
"""
|
||||
cursor.execute(create_stmt)
|
||||
|
||||
# Insert or update each object
|
||||
for obj in obj_list:
|
||||
values = tuple(str(v) if v is not None else "" for v in obj.to_dict().values()) # pyright: ignore[reportAttributeAccessIssue]
|
||||
insert_stmt = f"""
|
||||
INSERT INTO {table_name} ({columns}) VALUES ({placeholders})
|
||||
ON CONFLICT({primary_key}) DO UPDATE SET {update_clause}
|
||||
"""
|
||||
cursor.execute(insert_stmt, values)
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
logger.info(f"Inserted or updated {len(obj_list)} records into {table_name} table in {db_path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to insert or update into SQLite: {e}")
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# SQLITE READ
|
||||
# -------------------------------------------------------------------
|
||||
def read_objects_from_sqlite(cls: Type[T], table_name: str, db_path: str) -> List[T]:
|
||||
"""
|
||||
Generic function to read rows from a SQLite table and convert them into class instances.
|
||||
|
||||
Args:
|
||||
cls (Type[T]): The class to instantiate (e.g., Agent, Policy, Hash).
|
||||
table_name (str): The name of the table to query.
|
||||
db_path (str): Path to the SQLite database file.
|
||||
|
||||
Returns:
|
||||
List[T]: List of class instances.
|
||||
|
||||
Example:
|
||||
agents = read_objects_from_sqlite(Agent, "agents", "data_store.db")
|
||||
"""
|
||||
try:
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(f"SELECT * FROM {table_name}")
|
||||
rows = cursor.fetchall()
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
objects = []
|
||||
for row in rows:
|
||||
data = dict(zip(columns, row))
|
||||
if hasattr(cls, "from_dict"):
|
||||
obj = cls.from_dict(data) # pyright: ignore[reportAttributeAccessIssue]
|
||||
else:
|
||||
obj = cls(**data)
|
||||
objects.append(obj)
|
||||
|
||||
conn.close()
|
||||
logger.info(f"Read {len(objects)} records from {table_name} table in {db_path}")
|
||||
return objects
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to read from {table_name} in {db_path}: {e}")
|
||||
return []
|
||||
@@ -0,0 +1,291 @@
|
||||
# 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 logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
|
||||
import dotenv
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from models.agent import Agent
|
||||
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_json, load_env, load_env_json
|
||||
from utils.setup import get_base_directory
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
def getLocalApprovals(api: AirlockAPIWrapper):
|
||||
base_dir = get_base_directory
|
||||
result = api.otp_find_awaiting()
|
||||
local_approval = pd.DataFrame(result["response"]["otpusage"])
|
||||
if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"):
|
||||
previous_run = pd.read_parquet(f"{base_dir}\\cache\\newest_local_approval.parquet")
|
||||
previous_run.to_parquet(
|
||||
f"{base_dir}\\cache\\last_local_approval.parquet", index=False
|
||||
)
|
||||
os.remove(f"{base_dir}\\cache\\newest_local_approval.parquet")
|
||||
|
||||
# Only keep rows presumably created by the generate local approval function
|
||||
local_approval = local_approval[
|
||||
local_approval["purpose"].str.startswith("🎫 Local Approval 🎫")
|
||||
]
|
||||
|
||||
local_approval["batchid"] = local_approval["purpose"].apply(
|
||||
lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1)
|
||||
)
|
||||
|
||||
if not local_approval.empty:
|
||||
local_approval.to_parquet(
|
||||
f"{base_dir}\\cache\\newest_local_approval.parquet", index=False
|
||||
)
|
||||
|
||||
return local_approval
|
||||
|
||||
|
||||
def scheduleAddingLAHashes(api: AirlockAPIWrapper):
|
||||
|
||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")
|
||||
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
|
||||
pups = load_env_json("PUPS", "[]")
|
||||
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type = int)
|
||||
|
||||
try:
|
||||
register_function("add_hash", returnFromLocalApproval)
|
||||
register_function("move_device", moveAgentToRelatedPolicy)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to register functions: {e}")
|
||||
return
|
||||
|
||||
try:
|
||||
approvals_df = getNewLocalApprovals(api)
|
||||
if approvals_df.empty:
|
||||
logger.debug("No new local approvals found. Nothing to schedule.")
|
||||
return
|
||||
batches = approvals_df.groupby("batchid")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to retrieve or group local approvals: {e}")
|
||||
return
|
||||
|
||||
for batchid, batch_df in batches:
|
||||
try:
|
||||
duration_minutes = int(batch_df["duration"].iloc[0])
|
||||
start_time = datetime.datetime.now()
|
||||
run_time = start_time + datetime.timedelta(minutes=duration_minutes)
|
||||
early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95))
|
||||
|
||||
early_timestamp = early_time.timestamp()
|
||||
run_timestamp = run_time.timestamp()
|
||||
|
||||
# Schedule add_hash job
|
||||
try:
|
||||
run_once_job(
|
||||
f"add_hash_{batchid}",
|
||||
"add_hash",
|
||||
early_timestamp,
|
||||
[
|
||||
api,
|
||||
batch_df,
|
||||
policy_relationship_map,
|
||||
bad_publisher_list,
|
||||
pups,
|
||||
threat_tolerance_constant,
|
||||
],
|
||||
None,
|
||||
)
|
||||
logger.debug(f"Scheduled add_hash for batch {batchid} at {early_time}")
|
||||
except Exception:
|
||||
logger.debug("Failed to schedule add_hash for batch {batchid}: {e}")
|
||||
|
||||
# Schedule move_device jobs
|
||||
devices = batch_df["agentid"].drop_duplicates().tolist()
|
||||
agents = []
|
||||
|
||||
for device in devices:
|
||||
rows = api.agent_find_by_hostname(device).iterrows()
|
||||
agents += [Agent(**row["data"]) for _, row in rows]
|
||||
|
||||
for agent in agents:
|
||||
try:
|
||||
run_once_job(
|
||||
f"move_device_{agent.hostame}_{batchid}",
|
||||
"move_device",
|
||||
run_timestamp,
|
||||
[api, agent, policy_relationship_map],
|
||||
"enforcement",
|
||||
)
|
||||
|
||||
print(
|
||||
f"Scheduled move_device for device {agent.hostname} in batch {batchid} at {run_time}"
|
||||
)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"Failed to schedule move_device for device {agent.hostname} in batch {batchid}: {e}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to process batch {batchid}: {e}")
|
||||
|
||||
|
||||
def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant
|
||||
):
|
||||
"""
|
||||
# Get unique policy names from device list
|
||||
policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist())
|
||||
|
||||
# Create inverse map to go from Audit to Enforcement
|
||||
inverse_map = {v: k for k, v in policy_relationship_map.items()}
|
||||
|
||||
# Fetch all policies
|
||||
all_policies = [Policy(row['groupid'], row['hidden'], row['name'], row['parent']) for _, row in api.policy_find_all().iterrows()]
|
||||
|
||||
# Define policy types
|
||||
policy_types = [1, 2, 6, 7]
|
||||
|
||||
#TODO finish logic for adding hashes
|
||||
"""
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")
|
||||
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
|
||||
pups = load_env_json("PUPS", "[]")
|
||||
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE")
|
||||
print(f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}")
|
||||
|
||||
def moveToLocalApproval(api: AirlockAPIWrapper):
|
||||
possible_durations = [15, 60, 360, 1440, 10080]
|
||||
duration_selected = None
|
||||
|
||||
print(colorText("Please select a duration:", "white"))
|
||||
for i, option in enumerate(possible_durations, start=1):
|
||||
print(f"{i}. {option}")
|
||||
|
||||
try:
|
||||
|
||||
choice = int(get_sanitized_input("Enter the number of your choice:"))
|
||||
if 1 <= choice <= len(possible_durations):
|
||||
duration_selected = possible_durations[choice - 1]
|
||||
print(colorText(f"You selected: {duration_selected}", "yellow"))
|
||||
logger.debug(f"You selected: {duration_selected}")
|
||||
else:
|
||||
print(colorText("❌ Invalid choice.", "red"))
|
||||
logger.debug("Invalid Input")
|
||||
return
|
||||
except ValueError:
|
||||
print(colorText("❌ Invalid input. Please enter a number.", "red"))
|
||||
logger.debug("Invalid Input")
|
||||
return
|
||||
|
||||
agents = selectAgents(api)
|
||||
batch = int(time.time())
|
||||
|
||||
if not agents:
|
||||
print(colorText("❌ No agents found or error retrieving agents.", "red"))
|
||||
logger.debug("No agents found or error retrieving agents")
|
||||
return
|
||||
|
||||
for agent in agents:
|
||||
try:
|
||||
addLocalApproval(api, batch, duration_selected, agent.agentid)
|
||||
moveAgentToRelatedPolicy(api, agent, "audit")
|
||||
except Exception as e:
|
||||
print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red"))
|
||||
|
||||
|
||||
def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid):
|
||||
|
||||
purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
|
||||
api.otp_generate(agentid, duration_selected, purpose)
|
||||
|
||||
|
||||
|
||||
def monitorAuditStatus(api: AirlockAPIWrapper):
|
||||
current_agents = findAllAgents(api)
|
||||
last_agents = []
|
||||
if not last_agents:
|
||||
last_agents = current_agents
|
||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")
|
||||
|
||||
# Reverse map for audit → enforcement
|
||||
reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
|
||||
known_transitions = set(policy_relationship_map.items()) | set(reverse_policy_map.items())
|
||||
|
||||
# Index last_agents by hostname for quick lookup
|
||||
last_agent_map = {agent.hostname: agent for agent in last_agents}
|
||||
|
||||
# Result buckets
|
||||
newly_added = []
|
||||
same_policy = []
|
||||
moved_to_audit = []
|
||||
moved_to_enforcement = []
|
||||
unusual_move = []
|
||||
|
||||
for current in current_agents:
|
||||
previous = last_agent_map.get(current.hostname)
|
||||
|
||||
if not previous:
|
||||
newly_added.append(current)
|
||||
continue
|
||||
|
||||
if current.groupid == previous.groupid:
|
||||
same_policy.append(current)
|
||||
elif (previous.groupid, current.groupid) in known_transitions:
|
||||
moved_to_audit.append(current)
|
||||
elif (current.groupid, previous.groupid) in known_transitions:
|
||||
moved_to_enforcement.append(current)
|
||||
else:
|
||||
unusual_move.append(current)
|
||||
|
||||
# Return all five DataFrames
|
||||
return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move
|
||||
|
||||
|
||||
def getNewLocalApprovals(api: AirlockAPIWrapper):
|
||||
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
current_la = getLocalApprovals(api)
|
||||
|
||||
# Load old approval list
|
||||
old_la_path = f"{working_dir}\\Scheduling\\last_local_approval.parquet"
|
||||
if os.path.exists(old_la_path):
|
||||
old_la = pd.read_parquet(old_la_path)
|
||||
else:
|
||||
old_la = pd.DataFrame(columns=current_la.columns)
|
||||
|
||||
# Create composite keys
|
||||
current_la["key"] = current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str)
|
||||
old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str)
|
||||
|
||||
# Find new entries
|
||||
new_entries = current_la[~current_la["key"].isin(old_la["key"])]
|
||||
|
||||
# Convert 'granted' to datetime and filter by last 10 minutes
|
||||
new_entries["granted"] = pd.to_datetime(new_entries["granted"], utc=True, errors="coerce")
|
||||
ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10)
|
||||
recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago]
|
||||
|
||||
# Save current approvals for next run
|
||||
current_la.drop(columns=["key"], inplace=True)
|
||||
current_la.to_parquet(old_la_path, index=False)
|
||||
|
||||
return recent_entries
|
||||
@@ -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,177 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Callable, Dict, List
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Registry of functions that can be scheduled
|
||||
FUNCTION_MAP: Dict[str, Callable] = {}
|
||||
|
||||
# Dictionary to manually track scheduled jobs by ID
|
||||
scheduled_jobs: Dict[str, asyncio.TimerHandle] = {}
|
||||
|
||||
# Path to the JSON file for job persistence TODO - pin this to the correct place
|
||||
JOBS_FILE = os.path.join(os.getcwd(), "jobs.json")
|
||||
|
||||
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
|
||||
|
||||
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 save_jobs(jobs: List[Dict[str, Any]]):
|
||||
"""
|
||||
Save jobs to the JSON file (overwrite).
|
||||
"""
|
||||
with open(JOBS_FILE, "w") as f:
|
||||
json.dump(jobs, f, indent=4)
|
||||
|
||||
def cancel_job(job_id: str):
|
||||
"""
|
||||
Cancel a scheduled job by ID and remove it from the registry and persistence.
|
||||
"""
|
||||
handle = scheduled_jobs.pop(job_id, None)
|
||||
if handle:
|
||||
handle.cancel()
|
||||
logger.info(f"Cancelled job '{job_id}'")
|
||||
|
||||
jobs = [j for j in load_jobs() if j.get("id") != job_id]
|
||||
save_jobs(jobs)
|
||||
|
||||
def run_once_job(job_id: str, func_name: str, delay_seconds: float, args=None, kwargs=None, persist=True):
|
||||
"""
|
||||
Schedule a job to run once after a delay (in seconds).
|
||||
"""
|
||||
args = args or []
|
||||
kwargs = kwargs or {}
|
||||
|
||||
def job_wrapper():
|
||||
func = FUNCTION_MAP.get(func_name)
|
||||
if func is None:
|
||||
logger.error(f"Function '{func_name}' is not registered.")
|
||||
return
|
||||
func(*args, **kwargs)
|
||||
cancel_job(job_id)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
handle = loop.call_later(delay_seconds, job_wrapper)
|
||||
scheduled_jobs[job_id] = handle
|
||||
|
||||
if persist:
|
||||
jobs = [j for j in load_jobs() if j.get("id") != job_id]
|
||||
jobs.append({
|
||||
"id": job_id,
|
||||
"type": "once",
|
||||
"delay": delay_seconds,
|
||||
"function": func_name,
|
||||
"args": args,
|
||||
"kwargs": kwargs
|
||||
})
|
||||
save_jobs(jobs)
|
||||
logger.info(f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds.")
|
||||
|
||||
def recurring_job(job_id: str, func_name: str, interval: float, args=None, kwargs=None, persist=True):
|
||||
"""
|
||||
Schedule a recurring job.
|
||||
"""
|
||||
args = args or []
|
||||
kwargs = kwargs or {}
|
||||
|
||||
def job_wrapper():
|
||||
func = FUNCTION_MAP.get(func_name)
|
||||
if func is None:
|
||||
logger.error(f"Function '{func_name}' is not registered.")
|
||||
return
|
||||
func(*args, **kwargs)
|
||||
# Reschedule the job
|
||||
handle = asyncio.get_event_loop().call_later(interval, job_wrapper)
|
||||
scheduled_jobs[job_id] = handle
|
||||
|
||||
cancel_job(job_id)
|
||||
handle = asyncio.get_event_loop().call_later(interval, job_wrapper)
|
||||
scheduled_jobs[job_id] = handle
|
||||
|
||||
if persist:
|
||||
jobs = [j for j in load_jobs() if j.get("id") != job_id]
|
||||
jobs.append({
|
||||
"id": job_id,
|
||||
"type": "recurring",
|
||||
"interval": interval,
|
||||
"function": func_name,
|
||||
"args": args,
|
||||
"kwargs": kwargs
|
||||
})
|
||||
save_jobs(jobs)
|
||||
logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.")
|
||||
|
||||
def reload_jobs():
|
||||
"""
|
||||
Reload jobs from JSON and reschedule them.
|
||||
"""
|
||||
jobs = load_jobs()
|
||||
for job in jobs:
|
||||
if job["type"] == "once":
|
||||
run_once_job(
|
||||
job["id"],
|
||||
job["function"],
|
||||
job["delay"],
|
||||
job.get("args"),
|
||||
job.get("kwargs"),
|
||||
persist=False
|
||||
)
|
||||
elif job["type"] == "recurring":
|
||||
recurring_job(
|
||||
job["id"],
|
||||
job["function"],
|
||||
job["interval"],
|
||||
job.get("args"),
|
||||
job.get("kwargs"),
|
||||
persist=False
|
||||
)
|
||||
|
||||
async def start_scheduler():
|
||||
"""
|
||||
Start the asynchronous scheduler loop.
|
||||
|
||||
This function is a placeholder to keep the event loop alive.
|
||||
Jobs are scheduled using asyncio.call_later and do not require polling.
|
||||
"""
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
logger.critical("Scheduler stopped.")
|
||||
|
||||
"""
|
||||
Start the asynchronous scheduler loop.
|
||||
|
||||
This function is a placeholder for compatibility. Since we use asyncio.call_later,
|
||||
jobs are scheduled directly on the event loop and no polling is required.
|
||||
|
||||
Usage:
|
||||
# In an async app (e.g., Textual)
|
||||
asyncio.create_task(start_scheduler())
|
||||
|
||||
# Or in a standalone script
|
||||
async def main():
|
||||
await start_scheduler()
|
||||
|
||||
asyncio.run(main())
|
||||
"""
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(3600) # Sleep indefinitely; jobs run via call_later
|
||||
except asyncio.CancelledError:
|
||||
logger.critical("Scheduler stopped.")
|
||||
@@ -0,0 +1,95 @@
|
||||
|
||||
-- Current state of agents
|
||||
CREATE TABLE IF NOT EXISTS agents (
|
||||
agentid TEXT PRIMARY KEY,
|
||||
clientversion TEXT,
|
||||
domain TEXT,
|
||||
freespace INTEGER,
|
||||
groupid TEXT,
|
||||
hostname TEXT,
|
||||
ip TEXT,
|
||||
localip TEXT,
|
||||
lastcheckin TEXT,
|
||||
os TEXT,
|
||||
policyversion TEXT,
|
||||
status INTEGER,
|
||||
username TEXT,
|
||||
groupname TEXT,
|
||||
status_text TEXT,
|
||||
firstseen TEXT
|
||||
);
|
||||
|
||||
-- History of clientversion changes
|
||||
CREATE TABLE IF NOT EXISTS agent_clientversion_history (
|
||||
agentid TEXT,
|
||||
clientversion TEXT,
|
||||
timestamp TEXT,
|
||||
PRIMARY KEY (agentid, clientversion, timestamp)
|
||||
);
|
||||
|
||||
-- History of IP changes
|
||||
CREATE TABLE IF NOT EXISTS agent_ip_history (
|
||||
agentid TEXT,
|
||||
ip TEXT,
|
||||
timestamp TEXT,
|
||||
PRIMARY KEY (agentid, ip, timestamp)
|
||||
);
|
||||
|
||||
-- History of local IP changes
|
||||
CREATE TABLE IF NOT EXISTS agent_localip_history (
|
||||
agentid TEXT,
|
||||
localip TEXT,
|
||||
timestamp TEXT,
|
||||
PRIMARY KEY (agentid, localip, timestamp)
|
||||
);
|
||||
|
||||
-- History of policyversion changes
|
||||
CREATE TABLE IF NOT EXISTS agent_policyversion_history (
|
||||
agentid TEXT,
|
||||
policyversion TEXT,
|
||||
timestamp TEXT,
|
||||
PRIMARY KEY (agentid, policyversion, timestamp)
|
||||
);
|
||||
|
||||
|
||||
-- Current state of policies
|
||||
CREATE TABLE IF NOT EXISTS policies (
|
||||
groupid TEXT PRIMARY KEY,
|
||||
hidden BOOLEAN,
|
||||
name TEXT,
|
||||
parent TEXT,
|
||||
firstseen TEXT
|
||||
);
|
||||
|
||||
-- History of policy name changes
|
||||
CREATE TABLE IF NOT EXISTS policy_name_history (
|
||||
groupid TEXT,
|
||||
name TEXT,
|
||||
timestamp TEXT,
|
||||
PRIMARY KEY (groupid, name, timestamp)
|
||||
);
|
||||
|
||||
|
||||
-- Current state of allowlists
|
||||
CREATE TABLE IF NOT EXISTS allowlists (
|
||||
applicationid TEXT PRIMARY KEY,
|
||||
name TEXT,
|
||||
version TEXT,
|
||||
firstseen TEXT
|
||||
);
|
||||
|
||||
-- History of allowlist name changes
|
||||
CREATE TABLE IF NOT EXISTS allowlist_name_history (
|
||||
applicationid TEXT,
|
||||
name TEXT,
|
||||
timestamp TEXT,
|
||||
PRIMARY KEY (applicationid, name, timestamp)
|
||||
);
|
||||
|
||||
-- History of allowlist version changes
|
||||
CREATE TABLE IF NOT EXISTS allowlist_version_history (
|
||||
applicationid TEXT,
|
||||
version TEXT,
|
||||
timestamp TEXT,
|
||||
PRIMARY KEY (applicationid, version, timestamp)
|
||||
);
|
||||
@@ -0,0 +1,232 @@
|
||||
import sqlite3
|
||||
import logging
|
||||
from models.agent import Agent
|
||||
from models.policy import Policy, Allowlist
|
||||
from datetime import datetime
|
||||
from typing import List
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ------------------ Database Initialization ------------------
|
||||
|
||||
def initialize_db(db_path: str, schema_path: str, schema_version: str = "1.0"):
|
||||
try:
|
||||
with open(schema_path, 'r') as f:
|
||||
schema_sql = f.read()
|
||||
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.executescript(schema_sql)
|
||||
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS schema_version (
|
||||
version TEXT,
|
||||
applied_on TEXT
|
||||
)
|
||||
""")
|
||||
cursor.execute("""
|
||||
INSERT INTO schema_version (version, applied_on)
|
||||
VALUES (?, ?)
|
||||
""", (schema_version, datetime.now().isoformat()))
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
logger.info(f"Tracking database initialized at {db_path} with schema version {schema_version}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error initializing database: {e}")
|
||||
|
||||
# ------------------ Change Tracking ------------------
|
||||
|
||||
def compare_and_track_changes(
|
||||
conn,
|
||||
current_agents: List[Agent],
|
||||
current_policies: List[Policy],
|
||||
current_allowlists: List[Allowlist]
|
||||
):
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# --- Agents ---
|
||||
cursor.execute("SELECT * FROM agents")
|
||||
agent_rows = cursor.fetchall()
|
||||
agent_columns = [desc[0] for desc in cursor.description]
|
||||
previous_agents = {
|
||||
row[agent_columns.index("agentid")]: Agent(**dict(zip(agent_columns, row)))
|
||||
for row in agent_rows
|
||||
}
|
||||
|
||||
for agent in current_agents:
|
||||
if agent.agentid not in previous_agents:
|
||||
insert_new_agent(conn, agent)
|
||||
else:
|
||||
log_agent_changes(conn, previous_agents[agent.agentid], agent)
|
||||
|
||||
# --- Policies ---
|
||||
cursor.execute("SELECT * FROM policies")
|
||||
policy_rows = cursor.fetchall()
|
||||
policy_columns = [desc[0] for desc in cursor.description]
|
||||
previous_policies = {
|
||||
row[policy_columns.index("groupid")]: Policy(**dict(zip(policy_columns, row)))
|
||||
for row in policy_rows
|
||||
}
|
||||
|
||||
for policy in current_policies:
|
||||
if policy.groupid not in previous_policies:
|
||||
insert_new_policy(conn, policy)
|
||||
else:
|
||||
log_policy_changes(conn, previous_policies[policy.groupid], policy)
|
||||
|
||||
# --- Allowlists ---
|
||||
cursor.execute("SELECT * FROM allowlists")
|
||||
allowlist_rows = cursor.fetchall()
|
||||
allowlist_columns = [desc[0] for desc in cursor.description]
|
||||
previous_allowlists = {
|
||||
row[allowlist_columns.index("applicationid")]: Allowlist(**dict(zip(allowlist_columns, row)))
|
||||
for row in allowlist_rows
|
||||
}
|
||||
|
||||
for allowlist in current_allowlists:
|
||||
if allowlist.applicationid not in previous_allowlists:
|
||||
insert_new_allowlist(conn, allowlist)
|
||||
else:
|
||||
log_allowlist_changes(conn, previous_allowlists[allowlist.applicationid], allowlist)
|
||||
|
||||
conn.commit()
|
||||
logger.info("Change tracking completed successfully.")
|
||||
except Exception as e:
|
||||
logger.error(f"Error during change tracking: {e}")
|
||||
|
||||
# ------------------ Insert and Log Functions ------------------
|
||||
|
||||
def insert_new_agent(conn, agent: Agent):
|
||||
try:
|
||||
timestamp = datetime.now().isoformat()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
INSERT INTO agents (
|
||||
agentid, clientversion, domain, freespace, groupid, hostname, ip, localip,
|
||||
lastcheckin, os, policyversion, status, username, groupname, status_text, firstseen
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(agentid) DO NOTHING
|
||||
""", (
|
||||
agent.agentid, agent.clientversion, agent.domain, agent.freespace, agent.groupid,
|
||||
agent.hostname, agent.ip, agent.localip, agent.lastcheckin, agent.os,
|
||||
agent.policyversion, agent.status, agent.username, agent.groupname, agent.status_text,
|
||||
timestamp
|
||||
))
|
||||
|
||||
for field, value in [
|
||||
("clientversion", agent.clientversion),
|
||||
("ip", agent.ip),
|
||||
("localip", agent.localip),
|
||||
("policyversion", agent.policyversion),
|
||||
("hostname", agent.hostname),
|
||||
("status", agent.status)
|
||||
]:
|
||||
cursor.execute(f"""
|
||||
INSERT INTO agent_{field}_history (agentid, {field}, timestamp)
|
||||
VALUES (?, ?, ?)
|
||||
""", (agent.agentid, value, timestamp))
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Inserted new agent: {agent.agentid}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error inserting new agent {agent.agentid}: {e}")
|
||||
|
||||
def log_agent_changes(conn, previous: Agent, current: Agent):
|
||||
try:
|
||||
timestamp = datetime.now().isoformat()
|
||||
cursor = conn.cursor()
|
||||
|
||||
for field in ["clientversion", "ip", "localip", "policyversion", "hostname", "status"]:
|
||||
if getattr(previous, field) != getattr(current, field):
|
||||
cursor.execute(f"""
|
||||
INSERT INTO agent_{field}_history (agentid, {field}, timestamp)
|
||||
VALUES (?, ?, ?)
|
||||
""", (current.agentid, getattr(current, field), timestamp))
|
||||
|
||||
cursor.execute("""
|
||||
UPDATE agents SET lastcheckin = ? WHERE agentid = ?
|
||||
""", (current.lastcheckin, current.agentid))
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Logged changes for agent: {current.agentid}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error logger changes for agent {current.agentid}: {e}")
|
||||
|
||||
def insert_new_policy(conn, policy: Policy):
|
||||
try:
|
||||
timestamp = datetime.now().isoformat()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
INSERT INTO policies (groupid, hidden, name, parent, firstseen)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
ON CONFLICT(groupid) DO NOTHING
|
||||
""", (policy.groupid, policy.hidden, policy.name, policy.parent, timestamp))
|
||||
|
||||
cursor.execute("""
|
||||
INSERT INTO policy_name_history (groupid, name, timestamp)
|
||||
VALUES (?, ?, ?)
|
||||
""", (policy.groupid, policy.name, timestamp))
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Inserted new policy: {policy.groupid}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error inserting new policy {policy.groupid}: {e}")
|
||||
|
||||
def log_policy_changes(conn, previous: Policy, current: Policy):
|
||||
try:
|
||||
timestamp = datetime.now().isoformat()
|
||||
cursor = conn.cursor()
|
||||
|
||||
if previous.name != current.name:
|
||||
cursor.execute("""
|
||||
INSERT INTO policy_name_history (groupid, name, timestamp)
|
||||
VALUES (?, ?, ?)
|
||||
""", (current.groupid, current.name, timestamp))
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Logged changes for policy: {current.groupid}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error logger changes for policy {current.groupid}: {e}")
|
||||
|
||||
def insert_new_allowlist(conn, allowlist: Allowlist):
|
||||
try:
|
||||
timestamp = datetime.now().isoformat()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
INSERT INTO allowlists (applicationid, name, version, firstseen)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT(applicationid) DO NOTHING
|
||||
""", (allowlist.applicationid, allowlist.name, allowlist.version, timestamp))
|
||||
|
||||
for field, value in [("name", allowlist.name), ("version", allowlist.version)]:
|
||||
cursor.execute(f"""
|
||||
INSERT INTO allowlist_{field}_history (applicationid, {field}, timestamp)
|
||||
VALUES (?, ?, ?)
|
||||
""", (allowlist.applicationid, value, timestamp))
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Inserted new allowlist: {allowlist.applicationid}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error inserting new allowlist {allowlist.applicationid}: {e}")
|
||||
|
||||
def log_allowlist_changes(conn, previous: Allowlist, current: Allowlist):
|
||||
try:
|
||||
timestamp = datetime.now().isoformat()
|
||||
cursor = conn.cursor()
|
||||
|
||||
for field in ["name", "version"]:
|
||||
if getattr(previous, field) != getattr(current, field):
|
||||
cursor.execute(f"""
|
||||
INSERT INTO allowlist_{field}_history (applicationid, {field}, timestamp)
|
||||
VALUES (?, ?, ?)
|
||||
""", (current.applicationid, getattr(current, field), timestamp))
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Logged changes for allowlist: {current.applicationid}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error logger changes for allowlist {current.applicationid}: {e}")
|
||||
@@ -0,0 +1,307 @@
|
||||
# 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_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 = {
|
||||
"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"])
|
||||
|
||||
def otp_revoke(self, otpid: str) -> dict:
|
||||
"""
|
||||
Revoke an active OTP.
|
||||
Parameters:
|
||||
- otpid (str): The ID of the OTP to revoke.
|
||||
Returns:
|
||||
- dict: JSON response from the API.
|
||||
"""
|
||||
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)
|
||||
|
||||
|
||||
# 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_list_allowlists(self, groupid: str) -> pd.DataFrame:
|
||||
"""List allowlists assigned to a specific policy group."""
|
||||
payload = {"groupid": groupid}
|
||||
result = self._post("/v1/group/policies", payload)
|
||||
return pd.DataFrame(result["response"]["applications"])
|
||||
|
||||
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)
|
||||
|
||||
def policy_set_script_custom(
|
||||
self,
|
||||
groupid: str,
|
||||
script_custom: int,
|
||||
scripts_audit: Optional[List[str]] = None,
|
||||
scripts_disabled: Optional[List[str]] = None,
|
||||
scripts_respect: Optional[List[str]] = None
|
||||
) -> dict:
|
||||
"""
|
||||
Set the script control behavior for a policy group.
|
||||
|
||||
Parameters:
|
||||
- groupid (str): Target Group ID.
|
||||
- script_custom (int): 0 = Disabled, 1 = Enabled.
|
||||
- scripts_audit (List[str], optional): Script types to audit.
|
||||
- scripts_disabled (List[str], optional): Script types to disable.
|
||||
- scripts_respect (List[str], optional): Script types to respect current policy.
|
||||
|
||||
Returns:
|
||||
- dict: JSON response from the API.
|
||||
|
||||
Available script types are "batch","powershell","command","vbscript","javascript","windowsinstaller","htmlapplication","javaapplication","windowsscriptcomponent","compiledhtml","shellscript","dylib","python"
|
||||
|
||||
"""
|
||||
payload = {
|
||||
"groupid": groupid,
|
||||
"script_custom": script_custom,
|
||||
"scripts_audit": scripts_audit or [],
|
||||
"scripts_disabled": scripts_disabled or [],
|
||||
"scripts_respect": scripts_respect or []
|
||||
}
|
||||
return self._post("/v1/group/settings/script_custom", 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,75 @@
|
||||
# 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/>.
|
||||
|
||||
|
||||
#TODO Continue implementing logger
|
||||
#TODO Add input sanitation and CSV injection prevention
|
||||
#TODO Continue OTP and Local approval rewrites
|
||||
#TODO Explore pywin32
|
||||
#TODO Fix Requirements.txt
|
||||
#TODO Create Generic system_config.json for gitea
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
import dotenv
|
||||
import urllib3
|
||||
|
||||
import utils.menus as menus
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.security import getAPI
|
||||
from utils.setup import setup
|
||||
|
||||
urllib3.disable_warnings(
|
||||
urllib3.exceptions.InsecureRequestWarning
|
||||
)
|
||||
|
||||
def main():
|
||||
|
||||
#Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
|
||||
|
||||
working_dir = setup()
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.debug("🔍 Logging test: this should appear in both console and file.")
|
||||
dotenv.load_dotenv(dotenv_path=working_dir / ".env")
|
||||
|
||||
try:
|
||||
url = os.getenv("URL")
|
||||
username = os.getenv("USERNAME")
|
||||
|
||||
if not url:
|
||||
raise ValueError("Missing URL in environment variables.")
|
||||
if not username:
|
||||
raise ValueError("Missing USERNAME in environment variables.")
|
||||
|
||||
logger.debug(f"Retrieved URL: {url}")
|
||||
logger.debug(f"Retrieved Username: {username}")
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"Configuration error: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
if username:
|
||||
api = AirlockAPIWrapper(
|
||||
base_url=str(os.getenv("URL")),
|
||||
api_key = getAPI(username, "AirlockTools"), # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
|
||||
menus.menu_main(api)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,250 @@
|
||||
# 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 utils.configmanager import get_protected_json, load_env
|
||||
from utils.selector import Selector
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
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 not agents:
|
||||
logging.warning("No agents or policies found.")
|
||||
print("No agents matched the criteria.")
|
||||
return
|
||||
|
||||
# Convert enriched agents to DataFrame
|
||||
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
|
||||
|
||||
# Otherwise, print and optionally export
|
||||
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':
|
||||
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.")
|
||||
|
||||
|
||||
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"))
|
||||
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||
|
||||
device_input_lines = []
|
||||
empty_line_count = 0
|
||||
|
||||
# Regex to validate each line
|
||||
valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')
|
||||
|
||||
while True:
|
||||
line = get_sanitized_input("")
|
||||
stripped_line = line.strip()
|
||||
|
||||
if stripped_line == "":
|
||||
empty_line_count += 1
|
||||
if empty_line_count == 2:
|
||||
break
|
||||
continue # Don't validate empty lines
|
||||
else:
|
||||
empty_line_count = 0
|
||||
|
||||
# Validate only non-empty lines
|
||||
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"))
|
||||
|
||||
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 to match hostnames
|
||||
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"))
|
||||
|
||||
# Enrich each agent using its class method
|
||||
for agent in matched_agents:
|
||||
agent.enrich_with_policies(policies)
|
||||
|
||||
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 = get_protected_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,115 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PROTECTED_KEYS = [
|
||||
"APPNAME",
|
||||
"LOG_LEVEL",
|
||||
"PATH_EXCLUSION_CONST",
|
||||
"MIN_FILES_FOR_PATH",
|
||||
"VT_THREAT_TOLERANCE",
|
||||
"POLICY_MAP_ENF_AUD"
|
||||
]
|
||||
|
||||
_protected_config = {}
|
||||
|
||||
def get_system_config_path() -> Path:
|
||||
# Check inside bundled EXE directory first
|
||||
bundled_dir = Path(getattr(sys, '_MEIPASS', ''))
|
||||
bundled_path = bundled_dir / "system_config.json"
|
||||
if bundled_path.exists():
|
||||
return bundled_path
|
||||
|
||||
# Fallback to external location
|
||||
return Path(__file__).parent.parent / "system_config.json"
|
||||
|
||||
def load_protected_config() -> dict:
|
||||
global _protected_config
|
||||
try:
|
||||
with open(get_system_config_path(), "r") as f:
|
||||
system_config = json.load(f)
|
||||
except FileNotFoundError:
|
||||
logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
|
||||
system_config = {
|
||||
"APPNAME": "AirlockTools",
|
||||
"PATH_EXCLUSION_CONST": 4,
|
||||
"MIN_FILES_FOR_PATH": 4,
|
||||
"VT_THREAT_TOLERANCE": 4,
|
||||
"POLICY_MAP_ENF_AUD": {
|
||||
"enforced_id": "audit_id"
|
||||
}
|
||||
}
|
||||
|
||||
_protected_config = {key: system_config[key] for key in PROTECTED_KEYS}
|
||||
return _protected_config
|
||||
|
||||
def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
|
||||
value = _protected_config.get(key)
|
||||
if value is None:
|
||||
logging.warning(f"Protected config key '{key}' not found.")
|
||||
return default
|
||||
try:
|
||||
if isinstance(value, str):
|
||||
value = value.strip("'\"")
|
||||
return cast_type(value)
|
||||
except (ValueError, TypeError):
|
||||
logging.warning(f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}.")
|
||||
return default
|
||||
|
||||
def get_protected_json(key: str, default: str = "{}") -> dict:
|
||||
raw = _protected_config.get(key, default)
|
||||
if isinstance(raw, dict):
|
||||
return raw
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
escaped = raw.encode('unicode_escape').decode('utf-8')
|
||||
return json.loads(escaped)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to parse protected JSON key '{key}': {e}")
|
||||
return json.loads(default)
|
||||
|
||||
|
||||
|
||||
|
||||
def load_env_json(key: str, default: str):
|
||||
raw = os.getenv(key, default)
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
escaped = raw.encode('unicode_escape').decode('utf-8')
|
||||
return json.loads(escaped)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to parse {key}: {e}")
|
||||
return json.loads(default)
|
||||
|
||||
def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
|
||||
"""
|
||||
Safely retrieves an environment variable and casts it to the desired type.
|
||||
|
||||
Parameters:
|
||||
key (str): The name of the environment variable.
|
||||
cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str.
|
||||
default (Optional[T], optional): Default value if the variable is not set or invalid.
|
||||
|
||||
Returns:
|
||||
Optional[T]: The casted value or the default.
|
||||
"""
|
||||
value = os.getenv(key)
|
||||
if value is None:
|
||||
logger.warning(f"Environment variable '{key}' not set.")
|
||||
return default
|
||||
try:
|
||||
value = value.strip("'\"") # Strip surrounding quotes
|
||||
return cast_type(value)
|
||||
except (ValueError, TypeError):
|
||||
logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.")
|
||||
return default
|
||||
@@ -0,0 +1,398 @@
|
||||
# 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 inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import List, Optional
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
from services.policyhandler import pullPolicyExechistories
|
||||
from utils.configmanager import get_protected_value, load_env_json
|
||||
from utils.utils import colorText, regulator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
class Hash:
|
||||
"""
|
||||
Hash model representing Hash data
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sha256,
|
||||
applications=None,
|
||||
baselines=None,
|
||||
blocklists=None,
|
||||
createtime=None,
|
||||
datetime=None,
|
||||
description=None,
|
||||
filename=None,
|
||||
filepath=None,
|
||||
filesize=None,
|
||||
md5=None,
|
||||
modtime=None,
|
||||
origname=None,
|
||||
productname=None,
|
||||
productversion=None,
|
||||
publisher=None,
|
||||
reputation=None,
|
||||
sha128=None,
|
||||
sha384=None,
|
||||
sha512=None,
|
||||
):
|
||||
self.sha256 = sha256
|
||||
self.applications = applications
|
||||
self.baselines = baselines
|
||||
self.blocklists = blocklists
|
||||
self.createtime = createtime
|
||||
self.datetime = datetime
|
||||
self.description = description
|
||||
self.filename = filename
|
||||
self.filepath = filepath
|
||||
self.filesize = filesize
|
||||
self.md5 = md5
|
||||
self.modtime = modtime
|
||||
self.origname = origname
|
||||
self.productname = productname
|
||||
self.productversion = productversion
|
||||
self.publisher = publisher
|
||||
self.reputation = reputation
|
||||
self.sha128 = sha128
|
||||
self.sha384 = sha384
|
||||
self.sha512 = sha512
|
||||
|
||||
def __repr__(self):
|
||||
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
|
||||
return f"<Hash({attrs})>"
|
||||
|
||||
def __eq__(self, other):
|
||||
if isinstance(other, Hash):
|
||||
return self.sha256 == other.sha256
|
||||
return False
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self.sha256)
|
||||
|
||||
def to_dict(self):
|
||||
"""Returns a dictionary representation of the hash."""
|
||||
return self.__dict__
|
||||
|
||||
@staticmethod
|
||||
def safe_int(value, default=0):
|
||||
"""Safely convert a value to int, returning default on failure."""
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
@classmethod
|
||||
def deduplicate(cls, hash_list):
|
||||
"""
|
||||
Deduplicates a list of Hash objects based on sha256.
|
||||
Args:
|
||||
hash_list (list): List of Hash instances.
|
||||
Returns:
|
||||
list: Deduplicated list of Hash instances.
|
||||
"""
|
||||
seen = set()
|
||||
deduped = []
|
||||
for h in hash_list:
|
||||
if h.sha256 not in seen:
|
||||
seen.add(h.sha256)
|
||||
deduped.append(h)
|
||||
return deduped
|
||||
|
||||
@classmethod
|
||||
def categorize_hashes(cls, hashes):
|
||||
threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int)
|
||||
bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
|
||||
pups_pattern = regulator(load_env_json("PUPS", "[]"))
|
||||
|
||||
needs_review = []
|
||||
approved = []
|
||||
unapproved = []
|
||||
|
||||
for hash_obj in hashes:
|
||||
publisher = hash_obj.publisher or ""
|
||||
description = hash_obj.description or ""
|
||||
reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
|
||||
scannermatch = reputation.get("scannermatch")
|
||||
|
||||
logger.debug(f"Evaluating hash: {hash_obj}")
|
||||
logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}")
|
||||
|
||||
# 1. Unapproved: bad publisher or PUP
|
||||
if re.search(bad_publishers_pattern, publisher, re.IGNORECASE):
|
||||
logger.debug("Unapproved: Publisher matches bad publisher pattern.")
|
||||
unapproved.append(hash_obj)
|
||||
continue
|
||||
|
||||
if re.search(pups_pattern, description, re.IGNORECASE):
|
||||
logger.debug("Unapproved: Description matches PUP pattern.")
|
||||
unapproved.append(hash_obj)
|
||||
continue
|
||||
|
||||
# 2. Approved: signed
|
||||
if publisher != "Not Signed":
|
||||
logger.debug("Approved: File is signed and not flagged.")
|
||||
approved.append(hash_obj)
|
||||
continue
|
||||
|
||||
# 3. Approved or Unapproved based on threat level
|
||||
try:
|
||||
score = int(scannermatch) # pyright: ignore[reportArgumentType]
|
||||
logger.debug(f"Parsed scannermatch score: {score}")
|
||||
if score > threat_tolerance: # pyright: ignore[reportOperatorIssue]
|
||||
logger.debug("Unapproved: Unsigned file with high threat score.")
|
||||
unapproved.append(hash_obj)
|
||||
else:
|
||||
logger.debug("Approved: Unsigned file with low threat score.")
|
||||
approved.append(hash_obj)
|
||||
except (ValueError, TypeError):
|
||||
logger.debug("Needs Review: Scannermatch score is missing or invalid.")
|
||||
needs_review.append(hash_obj)
|
||||
|
||||
|
||||
logger.debug(f"Final counts — Needs Review: {len(needs_review)}, Approved: {len(approved)}, Unapproved: {len(unapproved)}")
|
||||
return needs_review, approved, unapproved
|
||||
|
||||
|
||||
@classmethod
|
||||
def export_to_csv(cls, hash_list, directory_path):
|
||||
"""
|
||||
Exports a list of Hash objects to a CSV file in the specified directory.
|
||||
The filename is derived from the variable name of the list if possible,
|
||||
and includes a timestamp to ensure uniqueness.
|
||||
"""
|
||||
filename = "hashes_export.csv"
|
||||
frame = inspect.currentframe()
|
||||
if frame is not None and frame.f_back is not None:
|
||||
callers_local_vars = frame.f_back.f_locals.items()
|
||||
for var_name, var_val in callers_local_vars:
|
||||
if var_val is hash_list:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"{var_name}_{timestamp}.csv"
|
||||
break
|
||||
else:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"hashes_export_{timestamp}.csv"
|
||||
|
||||
os.makedirs(directory_path, exist_ok=True)
|
||||
file_path = os.path.join(directory_path, filename)
|
||||
|
||||
df = pd.DataFrame([h.to_dict() for h in hash_list])
|
||||
df.to_csv(file_path, index=False)
|
||||
|
||||
logger.info(f"CSV file saved to: {file_path}")
|
||||
|
||||
|
||||
"""
|
||||
#Example - Convert Dataframe returned by hash query into hash objects
|
||||
hash_objects = []
|
||||
for _, row in df.iterrows():
|
||||
try:
|
||||
parsed_data = ast.literal_eval(row['data'])
|
||||
hash_obj = Hash(sha256=row['sha256'], **parsed_data)
|
||||
hash_objects.append(hash_obj)
|
||||
except Exception as e:
|
||||
print(f"Error parsing row: {e}")
|
||||
|
||||
# Display the created Hash objects
|
||||
for obj in hash_objects:
|
||||
print(obj)
|
||||
|
||||
|
||||
# Categorize hashes
|
||||
needs_review, approved, unapproved = Hash.categorize_hashes(
|
||||
hashes=hash_objects,
|
||||
threat_tolerance=3,
|
||||
untrusted_pattern=untrusted_pattern,
|
||||
pups_pattern=pups_pattern
|
||||
)
|
||||
|
||||
# Deduplicate
|
||||
deduped_hashes = Hash.deduplicate(hash_list)
|
||||
|
||||
# Specify the directory where you want to save the CSV
|
||||
output_directory = "C:/Users/Brandon/Documents/HashExports"
|
||||
|
||||
# Call the export method
|
||||
Hash.export_to_csv(hashes_for_export, output_directory)
|
||||
|
||||
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExecutionHistoryRecord:
|
||||
# Mandatory fields
|
||||
username: str
|
||||
hostname: str
|
||||
netdomain: str
|
||||
filename: str
|
||||
ppolicy: str
|
||||
policyname: str
|
||||
policyver: str
|
||||
commandline: str
|
||||
publisher: str
|
||||
sha256: str
|
||||
datetime: str
|
||||
|
||||
# Optional fields
|
||||
type: Optional[int] = None
|
||||
pprocess: Optional[str] = None
|
||||
gprocess: Optional[str] = None
|
||||
md5: Optional[str] = None
|
||||
sha128: Optional[str] = None
|
||||
sha384: Optional[str] = None
|
||||
sha512: Optional[str] = None
|
||||
ip: Optional[str] = None
|
||||
localip: Optional[str] = None
|
||||
extid: Optional[str] = None
|
||||
extname: Optional[str] = None
|
||||
exttype: Optional[int] = None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension
|
||||
extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge
|
||||
|
||||
@staticmethod
|
||||
def enrich_with_hashes_and_export(
|
||||
executions: list, hashes: list, directory_path: str, label: str = "enriched"
|
||||
):
|
||||
exec_df = pd.DataFrame([e.__dict__ for e in executions])
|
||||
hash_df = pd.DataFrame([h.to_dict() for h in hashes])
|
||||
|
||||
logger.debug(f"Execution DataFrame columns: {exec_df.columns}")
|
||||
logger.debug(f"Hash DataFrame columns: {hash_df.columns}")
|
||||
|
||||
if hash_df.empty:
|
||||
logger.warning(f"hash_df is empty for label: {label}. Skipping merge.")
|
||||
merged_df = exec_df.copy()
|
||||
else:
|
||||
merged_df = pd.merge(
|
||||
exec_df,
|
||||
hash_df,
|
||||
on="sha256",
|
||||
how="left", # Preserve all executions, enrich where possible
|
||||
suffixes=("_exec", "_hash")
|
||||
)
|
||||
merged_df.sort_values(by="filename_exec", inplace=True)
|
||||
logger.info(f"Merged {len(merged_df)} rows. Non-null hash matches: {merged_df['sha256'].notna().sum()}")
|
||||
|
||||
filename = f"{label}_executions.csv"
|
||||
os.makedirs(directory_path, exist_ok=True)
|
||||
file_path = os.path.join(directory_path, filename)
|
||||
merged_df.to_csv(file_path, index=False)
|
||||
|
||||
logger.info(f"CSV file saved to: {file_path}")
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict):
|
||||
mandatory_fields = [
|
||||
"username",
|
||||
"hostname",
|
||||
"netdomain",
|
||||
"filename",
|
||||
"ppolicy",
|
||||
"policyname",
|
||||
"policyver",
|
||||
"commandline",
|
||||
"publisher",
|
||||
"sha256",
|
||||
"datetime",
|
||||
]
|
||||
missing_fields = [
|
||||
field for field in mandatory_fields if field not in data or data[field] is None
|
||||
]
|
||||
if missing_fields:
|
||||
raise ValueError(f"Missing mandatory fields: {missing_fields}")
|
||||
|
||||
return cls(
|
||||
username=data["username"],
|
||||
hostname=data["hostname"],
|
||||
netdomain=data["netdomain"],
|
||||
filename=data["filename"],
|
||||
ppolicy=data["ppolicy"],
|
||||
policyname=data["policyname"],
|
||||
policyver=data["policyver"],
|
||||
commandline=data["commandline"],
|
||||
publisher=data["publisher"],
|
||||
sha256=data["sha256"],
|
||||
datetime=data["datetime"],
|
||||
type=data.get("type"),
|
||||
pprocess=data.get("pprocess"),
|
||||
gprocess=data.get("gprocess"),
|
||||
md5=data.get("md5"),
|
||||
sha128=data.get("sha128"),
|
||||
sha384=data.get("sha384"),
|
||||
sha512=data.get("sha512"),
|
||||
ip=data.get("ip"),
|
||||
localip=data.get("localip"),
|
||||
extid=data.get("extid"),
|
||||
extname=data.get("extname"),
|
||||
exttype=data.get("exttype"),
|
||||
extbrowser=data.get("extbrowser"),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_policies(
|
||||
cls, api, selected_policies, type_: list, history_days: int
|
||||
) -> List["ExecutionHistoryRecord"]:
|
||||
executions = []
|
||||
for policy in selected_policies:
|
||||
execs = pullPolicyExechistories(
|
||||
api, policy, type_, history_days, True
|
||||
)
|
||||
if execs:
|
||||
data = json.loads(execs)
|
||||
exechistories = data.get("response", {}).get("exechistories", [])
|
||||
if not exechistories:
|
||||
continue
|
||||
|
||||
df = pd.DataFrame(exechistories)
|
||||
df = df.drop_duplicates(subset=["sha256", "filename", "hostname"])
|
||||
df = df.sort_values(by=["sha256", "filename"])
|
||||
|
||||
executions.extend([cls.from_dict(row.to_dict()) for _, row in df.iterrows()])
|
||||
logger.debug(f"Staging of Execution history for policy: {policy.name} is complete")
|
||||
print(
|
||||
colorText(
|
||||
f"Staging of Execution history for policy: {policy.name} is complete",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
|
||||
return executions
|
||||
|
||||
def __repr__(self):
|
||||
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
|
||||
return f"<Execution({attrs})>"
|
||||
|
||||
|
||||
"""
|
||||
executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30)
|
||||
|
||||
ExecutionHistoryRecord.enrich_with_hashes_and_export(executions, hash_objects, "C:/Users/Brandon/Documents/EnrichedExports")
|
||||
"""
|
||||
@@ -0,0 +1,316 @@
|
||||
# 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
|
||||
import os
|
||||
import re
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
import services.policyhandler as policyh
|
||||
from flows.otp import generate, otp_activities_by_agent, revoke
|
||||
from flows.prepPolicy import (
|
||||
buildPathsandPublishers,
|
||||
buildPreflights,
|
||||
selectAllowlists,
|
||||
selectPolicies,
|
||||
sortHashes,
|
||||
)
|
||||
from flows.quietAgent import findQuietAgents
|
||||
from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import load_env
|
||||
from utils.selector import Selector
|
||||
from utils.utils import (
|
||||
areYouSure,
|
||||
colorText,
|
||||
displayIntro,
|
||||
get_sanitized_input,
|
||||
open_directory,
|
||||
printEnforceChecklist,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
def menu_main(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
extras = load_env("EXTRAS")
|
||||
while True:
|
||||
displayIntro()
|
||||
# Add Settings, and give option to change working dir
|
||||
print(colorText("1. ✅ - Move Device(s) to local approval", "yellow"))
|
||||
print(colorText("2. 🎫 - OTP", "yellow"))
|
||||
print(colorText("3. 🔄 - Move to Audit/Enforcement", "yellow"))
|
||||
print(colorText("4. 🔍 - Device Search", "yellow"))
|
||||
print(colorText("5. 🔇 - Find Quiet Hosts", "yellow"))
|
||||
if extras == "POLICYPREP" : print(colorText("6. 🛡️ - Policy Enforcement Tools", "yellow"))
|
||||
print(colorText("F. 📂 - Open Working Directory", "yellow"))
|
||||
print(colorText("S. 🛠️ - Settings", "yellow"))
|
||||
print(colorText("Q. 🔚 - Quit", "yellow"))
|
||||
|
||||
choice = get_sanitized_input("\nEnter Menu Item: ")
|
||||
if choice == "1":
|
||||
print("This Feature is still in development")
|
||||
get_sanitized_input("Press enter to continue")
|
||||
elif choice == "2":
|
||||
menu_otp(api)
|
||||
elif choice == "3":
|
||||
choices = ["audit", "enforcement"]
|
||||
print(colorText("Move devices to which state?:", "yellow"))
|
||||
direction = Selector.select_string(choices, False, False)
|
||||
devices = selectAgents(api)
|
||||
print(colorText("Would you like to continue with these devices?","white"))
|
||||
for device in devices:
|
||||
print(device.hostname)
|
||||
confirm = Selector.confirm()
|
||||
if direction and devices and confirm:
|
||||
for device in devices:
|
||||
moveAgentToRelatedPolicy(api,device, direction[0])
|
||||
elif choice == "4":
|
||||
findAgents(api,False)
|
||||
elif choice == "5":
|
||||
findQuietAgents(api)
|
||||
elif choice == "6":
|
||||
if extras == "POLICYPREP": menu_policymanagment(api)
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
elif choice.upper() == "S":
|
||||
menu_settings()
|
||||
elif choice.upper() == "Q":
|
||||
break
|
||||
else:
|
||||
print(colorText("Invalid choice. Please try again.", "red"))
|
||||
|
||||
|
||||
|
||||
def menu_policy_enforce(api: AirlockAPIWrapper):
|
||||
selected_policies = []
|
||||
destination_policy = []
|
||||
destination_allowlist = []
|
||||
processed_paths = []
|
||||
processed_hashes = []
|
||||
processed_publishers = []
|
||||
tested = False
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
while True:
|
||||
printEnforceChecklist(selected_policies, destination_policy, destination_allowlist)
|
||||
choice = get_sanitized_input("\nEnter your choice: ")
|
||||
|
||||
if choice == "1":
|
||||
selected_policies = selectPolicies(api,True)
|
||||
|
||||
elif choice == "2":
|
||||
print(colorText("Please choose destination_name Policy for Path Exclusions", "white"))
|
||||
|
||||
destination_policy = selectPolicies(api, False)
|
||||
|
||||
print(colorText("Please choose Allowlist for Hashes", "white"))
|
||||
|
||||
destination_allowlist = selectAllowlists(api, destination_policy, False)
|
||||
|
||||
elif choice == "3":
|
||||
sortHashes(
|
||||
api,
|
||||
selected_policies,
|
||||
type=[1, 2, 6, 7],
|
||||
)
|
||||
|
||||
elif choice == "4":
|
||||
if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
|
||||
buildPathsandPublishers(False)
|
||||
else:
|
||||
print("File not found. Please make sure it's saved correctly and try again.")
|
||||
|
||||
elif choice == "5":
|
||||
if os.path.exists(f"{working_dir}\\Approved\\hashes_to_add.csv") and os.path.exists(
|
||||
f"{working_dir}\\Approved\\primary_Paths.csv"
|
||||
):
|
||||
buildPreflights()
|
||||
else:
|
||||
print("File not found. Please make sure it's saved correctly and try again.")
|
||||
|
||||
elif choice == "6":
|
||||
if (
|
||||
os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv")
|
||||
and os.path.exists(f"{working_dir}\\Preflight\\approved_hashes.csv")
|
||||
and destination_policy
|
||||
and destination_allowlist
|
||||
):
|
||||
print(colorText("These path exclusions would be added to:", "yellow"))
|
||||
print(destination_policy)
|
||||
|
||||
pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\approved_paths.csv")
|
||||
hashes = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv")
|
||||
|
||||
unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates()
|
||||
|
||||
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
|
||||
processed_paths = [
|
||||
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
|
||||
for path, ext in unique_combinations.itertuples(index=False, name=None)
|
||||
]
|
||||
|
||||
print(processed_paths)
|
||||
print(colorText("These publishers would added", "yellow"))
|
||||
|
||||
if os.path.exists(f"{working_dir}\\Preflight\\approved_publishers.csv"):
|
||||
publishers = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv")
|
||||
if publishers.empty:
|
||||
print(colorText("The publishers list is empty.", "red"))
|
||||
else:
|
||||
processed_publishers = (
|
||||
publishers[publishers["publisher_hash"] != "Not Signed"]
|
||||
["publisher_hash"]
|
||||
.drop_duplicates()
|
||||
.tolist()
|
||||
)
|
||||
print(processed_publishers)
|
||||
|
||||
print(colorText("These hashes would be added to:", "yellow"))
|
||||
print(destination_allowlist)
|
||||
|
||||
processed_hashes = hashes["sha256"].unique().tolist()
|
||||
print(processed_hashes)
|
||||
|
||||
if processed_paths and processed_hashes:
|
||||
tested = True
|
||||
else:
|
||||
# Log which condition(s) failed
|
||||
missing_items = []
|
||||
if not os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv"):
|
||||
missing_items.append("approved_paths.csv not found")
|
||||
if not os.path.exists(f"{working_dir}\\Preflight\\approved_hashes.csv"):
|
||||
missing_items.append("approved_hashes.csv not found")
|
||||
if not destination_policy:
|
||||
missing_items.append("destination_policy is empty or None")
|
||||
if not destination_allowlist:
|
||||
missing_items.append("destination_allowlist is empty or None")
|
||||
|
||||
logger.error("Preflight check failed due to the following:")
|
||||
for item in missing_items:
|
||||
logger.error(f" - {item}")
|
||||
|
||||
|
||||
elif choice == "7":
|
||||
areYouSure()
|
||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
if (
|
||||
tested
|
||||
and destination_policy
|
||||
and destination_allowlist
|
||||
and confirmation.strip() == "I AGREE"
|
||||
):
|
||||
print(colorText("Proceeding with the code...", "yellow"))
|
||||
api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
|
||||
api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths)
|
||||
if processed_publishers:
|
||||
api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)
|
||||
else:
|
||||
logger.error("Confirmation block failed. Reasons:")
|
||||
if not tested:
|
||||
logger.error(" - Preflight checks were not completed successfully (`tested` is False).")
|
||||
if not destination_policy:
|
||||
logger.error(" - `destination_policy` is missing or invalid.")
|
||||
if not destination_allowlist:
|
||||
logger.error(" - `destination_allowlist` is missing or invalid.")
|
||||
if confirmation.strip() != "I AGREE":
|
||||
logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip())
|
||||
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
elif choice.upper() == "S":
|
||||
menu_settings()
|
||||
elif choice.upper == "B":
|
||||
break
|
||||
|
||||
|
||||
else:
|
||||
print(colorText("Invalid choice. Please try again.", "red"))
|
||||
|
||||
|
||||
def menu_otp(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
while True:
|
||||
|
||||
print(colorText("\n--- 🎫 OTP Submenu 🎫 ---", "cyan"))
|
||||
print(colorText("1. 🔐 -Generate OTPs", "cyan"))
|
||||
print(colorText("2. 📊 -OTP Activities By Agent", "cyan"))
|
||||
print(colorText("3. ❌ -Revoke OTPs", "cyan"))
|
||||
print(colorText("F. 📂 - Open Working Directory", "yellow"))
|
||||
print(colorText("S. 🛠️ - Settings", "yellow"))
|
||||
print(colorText("B. 🔙 - Back", "yellow"))
|
||||
|
||||
choice = get_sanitized_input("Enter your choice: ")
|
||||
|
||||
if choice == "1":
|
||||
otp_list = generate(api)
|
||||
print(colorText(otp_list,"green"))
|
||||
elif choice == "2":
|
||||
otp_activities_by_agent(api)
|
||||
elif choice == "3":
|
||||
revoke(api)
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
elif choice.upper() == "S":
|
||||
menu_settings()
|
||||
elif choice.upper() == "B":
|
||||
break
|
||||
|
||||
def menu_policymanagment(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
while True:
|
||||
print(colorText("1. 🔒 - Prepare Policy For Enforcement", "yellow"))
|
||||
print(colorText("2. 🔄 - Update Audit Policies from Enforcement Policies", "yellow"))
|
||||
print(colorText("F. 📂 - Open Working Directory", "yellow"))
|
||||
print(colorText("S. 🛠️ - Settings", "yellow"))
|
||||
print(colorText("B. 🔙 - Back", "yellow"))
|
||||
choice = get_sanitized_input("\n Enter Menu Item: ")
|
||||
|
||||
if choice == "1":
|
||||
menu_policy_enforce(api)
|
||||
elif choice == "2":
|
||||
areYouSure()
|
||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
if confirmation.strip() == "I AGREE":
|
||||
policyh.updateAuditPoliciesFromEnforcementPolices(api)
|
||||
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
elif choice.upper() == "S":
|
||||
menu_settings()
|
||||
elif choice.upper() == "B":
|
||||
break
|
||||
else:
|
||||
print(colorText("Invalid choice. Please try again.", "red"))
|
||||
def menu_settings():
|
||||
while True:
|
||||
print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
|
||||
print(colorText("This Feature is still in development", "cyan"))
|
||||
# print(colorText("2. Sub-option B","cyan"))
|
||||
print(colorText("B. 🔙 - Back", "yellow"))
|
||||
choice = get_sanitized_input("Enter your choice: ")
|
||||
|
||||
if choice == "1":
|
||||
pass #TODO ADD CHANGE WORKDIR CODE
|
||||
|
||||
elif choice.upper() == "B":
|
||||
print("Returning to Main Menu...")
|
||||
break
|
||||
else:
|
||||
print("Invalid choice. Please try again.")
|
||||
@@ -0,0 +1,69 @@
|
||||
# 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 services.agenthandler import selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.selector import Selector
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
||||
|
||||
def generate(api: AirlockAPIWrapper):
|
||||
otp_dict = {}
|
||||
agents = selectAgents(api)
|
||||
print(colorText("Would you like to continue with these devices?","white"))
|
||||
for agent in agents:
|
||||
print(agent.hostname)
|
||||
confirm = Selector.confirm()
|
||||
if agents and confirm:
|
||||
requester = get_sanitized_input("Who is requesting the OTP: ")
|
||||
because = get_sanitized_input("Why/What work are they doing?: ")
|
||||
|
||||
purpose = f"Requester: {requester} - for : {because}"
|
||||
possible_durations = [15, 60, 360, 1440, 10080]
|
||||
|
||||
print(colorText("Please select a duration in minutes: ", "white"))
|
||||
print(colorText("15 mins, 60 mins, 360 mins(6 Hours), 1440 mins (24 Hours), 10080 mins (7 Days):", "white"))
|
||||
duration_selected = Selector.select_int(possible_durations)
|
||||
if duration_selected:
|
||||
for agent in agents:
|
||||
otp_code = api.otp_generate(agent.agentid, duration_selected, purpose)
|
||||
logger.info(f"Generated OTP for {agent.hostname}: {otp_code}")
|
||||
otp_dict[agent.hostname] = otp_code
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,237 @@
|
||||
# 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.configmanager import get_protected_json
|
||||
from utils.setup import get_base_directory
|
||||
from utils.utils import colorText
|
||||
|
||||
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 = get_protected_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,365 @@
|
||||
# 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
|
||||
import os
|
||||
import os.path
|
||||
from typing import List, Optional
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
from models.execution import ExecutionHistoryRecord, Hash
|
||||
from models.policy import Allowlist, Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_value, load_env, load_env_json
|
||||
from utils.selector import Selector
|
||||
from utils.utils import (
|
||||
colorText,
|
||||
formatHTML,
|
||||
import_to_dataframe,
|
||||
regulator,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
|
||||
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
|
||||
|
||||
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||
logger.debug("Prompting for Policies")
|
||||
print(colorText("Please select policy/policies", "white"))
|
||||
selected = Selector.select_objects(policies, allow_multiple, prompt_each=True)
|
||||
|
||||
if selected is None:
|
||||
return []
|
||||
|
||||
# Normalize to always return a list
|
||||
logger.debug("Returning {selected.dict}")
|
||||
return selected if isinstance(selected, list) else [selected]
|
||||
|
||||
|
||||
def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) -> List[Allowlist]:
|
||||
if policy == "all": allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()]
|
||||
else: allowlists = [Allowlist(**row.to_dict()) for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()]
|
||||
logger.debug("Prompting for Allowlist(s)")
|
||||
print(colorText("Please select allowlist(s)", "white"))
|
||||
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
|
||||
|
||||
if selected is None:
|
||||
return []
|
||||
|
||||
# Normalize to always return a list
|
||||
logger.debug(f"Returning {selected}")
|
||||
return selected if isinstance(selected, list) else [selected]
|
||||
|
||||
|
||||
|
||||
def sortHashes(
|
||||
api: AirlockAPIWrapper,
|
||||
selected_policies: List[Policy],
|
||||
type=[1, 2, 6, 7],
|
||||
history_days: Optional[int] = None
|
||||
):
|
||||
|
||||
|
||||
if history_days is None:
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
|
||||
|
||||
logger.debug(f"{history_days} day selected for history")
|
||||
|
||||
if history_days is None:
|
||||
logging.warning("No history range selected. Aborting.")
|
||||
return
|
||||
|
||||
executions = []
|
||||
hashes = []
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
# Pull execution histories for each policy
|
||||
|
||||
policy_executions = ExecutionHistoryRecord.from_policies(
|
||||
api, selected_policies, type_=type, history_days=history_days
|
||||
)
|
||||
|
||||
logger.debug(f"Policy_executions is {policy_executions}")
|
||||
|
||||
executions.extend(policy_executions)
|
||||
logger.debug(f"Executions contains {executions}")
|
||||
if executions:
|
||||
hashes = [Hash(sha256=row["sha256"], **row["data"]) for _, row in api.hash_query([record.sha256 for record in executions]).iterrows()
|
||||
]
|
||||
|
||||
if hashes:
|
||||
unique_hashes = Hash.deduplicate(hashes)
|
||||
|
||||
needs_review, approved, unapproved = Hash.categorize_hashes(
|
||||
hashes=unique_hashes
|
||||
)
|
||||
|
||||
categories = {
|
||||
"needs_review": needs_review,
|
||||
"approved": approved,
|
||||
"unapproved": unapproved,
|
||||
}
|
||||
|
||||
for label, category in categories.items():
|
||||
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
|
||||
html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html"
|
||||
|
||||
ExecutionHistoryRecord.enrich_with_hashes_and_export(
|
||||
executions, category, f"{working_dir}\\Needs_Review\\Review_First", label=label
|
||||
)
|
||||
df = import_to_dataframe(csv_path)
|
||||
formatHTML(df, html_path)
|
||||
|
||||
def buildPathsandPublishers(split):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
df1 = pd.DataFrame()
|
||||
df2 = pd.DataFrame()
|
||||
all_approved_hashes = pd.DataFrame()
|
||||
path1 = f"{working_dir}\\Approved\\approved_executions.csv"
|
||||
path2 = f"{working_dir}\\Approved\\needs_review_executions.csv"
|
||||
|
||||
if os.path.exists(path1):
|
||||
df1 = pd.read_csv(path1)
|
||||
else:
|
||||
logger.warning(f"File not found: {path1}")
|
||||
|
||||
if os.path.exists(path2):
|
||||
df2 = pd.read_csv(path2)
|
||||
else:
|
||||
logger.warning(f"File not found: {path2}")
|
||||
|
||||
if df1.empty and df2.empty:
|
||||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||
all_approved_hashes = pd.DataFrame()
|
||||
logger.debug(all_approved_hashes.head)
|
||||
else:
|
||||
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
|
||||
if "filename_exec" in all_approved_hashes.columns:
|
||||
all_approved_hashes = all_approved_hashes.sort_values(by="filename_exec")
|
||||
else:
|
||||
logger.warning("Warning: 'filename_exec' column not found in concatenated DataFrame.")
|
||||
|
||||
if not all_approved_hashes.empty:
|
||||
primary_path_exclusions = calculatePath(
|
||||
all_approved_hashes,
|
||||
split,
|
||||
)
|
||||
remaining_hashes = all_approved_hashes[
|
||||
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
|
||||
]
|
||||
secondary_path_exclusions = calculatePath(
|
||||
remaining_hashes, split
|
||||
)
|
||||
remaining_hashes = remaining_hashes[
|
||||
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
|
||||
]
|
||||
dataframes = {
|
||||
"primary_Paths": primary_path_exclusions,
|
||||
"secondary_Paths": secondary_path_exclusions,
|
||||
"hashes_to_add": remaining_hashes,
|
||||
}
|
||||
logger.debug("Preparing to sort dataframes")
|
||||
for name, df in dataframes.items():
|
||||
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||||
if name == "hashes_to_add": df.sort_values(by="filename_exec", inplace=True)
|
||||
else: df.sort_values(by="longestcfp", inplace=True)
|
||||
|
||||
df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{name}.csv", index=False)
|
||||
formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{name}.html")
|
||||
|
||||
if not all_approved_hashes.empty:
|
||||
# Drop all not signed, only keep unique values
|
||||
publist = all_approved_hashes[
|
||||
all_approved_hashes["publisher_hash"] != "Not Signed"
|
||||
].drop_duplicates(subset=["publisher_hash"])
|
||||
# Remove Bad publisher if somehow they made it this far
|
||||
pattern = regulator(load_env_json("BAD_PUBLISHERS","[]"))
|
||||
publist = publist[~publist["publisher_hash"].str.contains(pattern, na=False)]
|
||||
publist = publist[["publisher_hash"]]
|
||||
publist.sort_values(by="publisher_hash", inplace=True)
|
||||
publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\publishers.csv", index=False)
|
||||
|
||||
def buildPreflights():
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
df1 = pd.DataFrame()
|
||||
df2 = pd.DataFrame()
|
||||
approved_hashes = pd.DataFrame()
|
||||
approved_publishers = pd.DataFrame()
|
||||
|
||||
hash = f"{working_dir}\\Approved\\hashes_to_add.csv"
|
||||
path1 = f"{working_dir}\\Approved\\primary_Paths.csv"
|
||||
path2 = f"{working_dir}\\Approved\\secondary_Paths.csv"
|
||||
publishers = f"{working_dir}\\Approved\\publishers.csv"
|
||||
|
||||
if os.path.exists(hash):
|
||||
approved_hashes = pd.read_csv(hash)
|
||||
|
||||
else:
|
||||
logger.warning(f"File not found: {hash}")
|
||||
|
||||
if os.path.exists(path1):
|
||||
df1 = pd.read_csv(path1)
|
||||
else:
|
||||
logger.warning(f"File not found: {path1}")
|
||||
|
||||
if os.path.exists(path2):
|
||||
df2 = pd.read_csv(path2)
|
||||
else:
|
||||
logger.warning(f"File not found: {path2}")
|
||||
|
||||
if df1.empty and df2.empty:
|
||||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||
approved_paths = pd.DataFrame()
|
||||
else:
|
||||
approved_paths = pd.concat([df1, df2], ignore_index=True)
|
||||
|
||||
if os.path.exists(publishers):
|
||||
approved_publishers = pd.read_csv(publishers)
|
||||
|
||||
else:
|
||||
logger.warning(f"File not found: {publishers}")
|
||||
|
||||
dataframes = {"approved_paths": approved_paths, "approved_hashes": approved_hashes, "approved_publishers": approved_publishers}
|
||||
|
||||
for name, df in dataframes.items():
|
||||
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||||
if name == "approved_paths":df.sort_values(by="longestcfp", inplace=True)
|
||||
elif name == "approved_hashes":df.sort_values(by="filename_exec", inplace=True)
|
||||
elif name == "approved_publishers" : df.sort_values(by="publisher_hash", inplace=True)
|
||||
|
||||
df.to_csv(f"{working_dir}\\Preflight\\{name}.csv", index=False)
|
||||
formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{name}.html")
|
||||
|
||||
def splitFilepathsGrouped(df, col="filename"):
|
||||
path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type= int)
|
||||
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type= int)
|
||||
|
||||
def clean_split(path):
|
||||
if not isinstance(path, (str, bytes, os.PathLike)):
|
||||
return []
|
||||
parts = os.path.normpath(path).split(os.sep)
|
||||
parts = [p for p in parts if p] # Remove empty strings
|
||||
return parts
|
||||
|
||||
# Diagnostic: log any non-string entries
|
||||
non_string_entries = df[~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))]
|
||||
if not non_string_entries.empty:
|
||||
print(f"[WARNING] Non-string entries found in column '{col}':")
|
||||
print(non_string_entries)
|
||||
|
||||
df = df.copy()
|
||||
split_paths = df[col].apply(clean_split)
|
||||
|
||||
# Filter out paths with fewer than `min_files_for_path` components
|
||||
df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy()
|
||||
split_paths = split_paths[df.index]
|
||||
|
||||
df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant]))
|
||||
grouped = df.groupby("group_key")
|
||||
new_rows = []
|
||||
|
||||
for _, group_df in grouped:
|
||||
paths = group_df[col].tolist()
|
||||
split_parts = [clean_split(p) for p in paths]
|
||||
|
||||
def longest_common_prefix(paths):
|
||||
if not paths:
|
||||
return []
|
||||
prefix = paths[0]
|
||||
for path in paths[1:]:
|
||||
prefix = [a for a, b in zip(prefix, path) if a == b]
|
||||
if not prefix:
|
||||
break
|
||||
return prefix
|
||||
|
||||
common_prefix = longest_common_prefix(split_parts)
|
||||
prefix_str = os.sep.join(common_prefix)
|
||||
|
||||
for i, parts in enumerate(split_parts):
|
||||
filename = parts[-1]
|
||||
middle = (
|
||||
os.sep.join(parts[len(common_prefix):-1])
|
||||
if len(parts) > len(common_prefix) + 1
|
||||
else ""
|
||||
)
|
||||
row = group_df.iloc[i].copy()
|
||||
row["longestcfp"] = prefix_str
|
||||
row["middle"] = middle
|
||||
row["filename_only"] = filename
|
||||
row["file_extension"] = os.path.splitext(filename)[1].lower()
|
||||
new_rows.append(row)
|
||||
|
||||
return pd.DataFrame(new_rows).drop(columns=["group_key"])
|
||||
|
||||
def calculatePath(approved_hashes, split):
|
||||
if split:
|
||||
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
|
||||
else:
|
||||
dfs_by_policy = [approved_hashes]
|
||||
|
||||
badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
|
||||
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type = int)
|
||||
|
||||
processed_dfs = []
|
||||
|
||||
for df in dfs_by_policy:
|
||||
haslcp = splitFilepathsGrouped(df, "filename_exec")
|
||||
haslcp = haslcp.drop_duplicates()
|
||||
|
||||
forbidden = regulator(badpathparts, True)
|
||||
forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
|
||||
|
||||
logger.debug("Removing forbidden filepaths for path exceptions")
|
||||
print(colorText("Removing forbidden filepaths for path exceptions", "green"))
|
||||
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
|
||||
|
||||
lcp_not_forbidden_review = lcp_not_forbidden[
|
||||
[
|
||||
"policyname",
|
||||
"longestcfp",
|
||||
"middle",
|
||||
"filename_only",
|
||||
"file_extension",
|
||||
"sha256",
|
||||
]
|
||||
]
|
||||
|
||||
unique_sha_counts = (
|
||||
lcp_not_forbidden_review.groupby("longestcfp")["sha256"].nunique().reset_index()
|
||||
)
|
||||
unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
|
||||
|
||||
lcp_not_forbidden_review = lcp_not_forbidden_review.merge(
|
||||
unique_sha_counts, on="longestcfp", how="left"
|
||||
)
|
||||
lcp_not_forbidden_review = lcp_not_forbidden_review[
|
||||
lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
|
||||
]
|
||||
processed_dfs.append(lcp_not_forbidden_review)
|
||||
|
||||
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
|
||||
|
||||
return pathExclusions
|
||||
@@ -0,0 +1,105 @@
|
||||
# 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 logging
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
from flows.prepPolicy import selectPolicies
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import getPolicyInfo
|
||||
from utils.selector import Selector
|
||||
from utils.utils import colorText, load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
def findQuietAgents(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
selected_policy = selectPolicies(api, False)
|
||||
if selected_policy:
|
||||
agents = api.agents_find_by_group(selected_policy[0].groupid)
|
||||
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
required_quiet = Selector.select_value(
|
||||
prompt="Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
|
||||
policy_exec_history = getPolicyInfo(
|
||||
api, selected_policy[0], [1, 2, 6, 7], history_days
|
||||
)
|
||||
|
||||
if policy_exec_history.empty:
|
||||
logging.info("No execution history found for the selected policy and time range.")
|
||||
return
|
||||
|
||||
policy_exec_history.loc[:, "datetime"] = pd.to_datetime(
|
||||
policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
|
||||
)
|
||||
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
|
||||
policy_exec_history.loc[:, "days_ago"] = policy_exec_history["datetime"].apply(
|
||||
lambda dt: (now - dt).days
|
||||
)
|
||||
|
||||
hostname_counts = policy_exec_history["hostname"].value_counts()
|
||||
|
||||
agents.loc[:, "execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int)
|
||||
|
||||
most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates(
|
||||
subset="hostname", keep="first"
|
||||
)
|
||||
|
||||
agents.loc[:, "days_since"] = agents["hostname"].map(
|
||||
most_recent_exec.set_index("hostname")["days_ago"]
|
||||
)
|
||||
|
||||
agents.loc[:, "required_quiet"] = required_quiet
|
||||
agents.loc[:, "enforce_ready"] = agents["days_since"].apply(
|
||||
lambda x: True if pd.isna(x) or x > required_quiet else False
|
||||
)
|
||||
|
||||
agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True])
|
||||
|
||||
filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
|
||||
logging.debug(f"Saving CSV to {filename}")
|
||||
print(colorText(f"Saving CSV to {filename}", "green"))
|
||||
agents.to_csv(filename, index=False)
|
||||
|
||||
total_agents = len(agents)
|
||||
ready_agents = agents["enforce_ready"].sum()
|
||||
not_ready_agents = total_agents - ready_agents
|
||||
ready_percentage = (ready_agents / total_agents) * 100
|
||||
|
||||
message = (
|
||||
f"Total agents: {total_agents}\n"
|
||||
f"Agents marked as 'enforce_ready': {ready_agents}\n"
|
||||
f"Agents not ready: {not_ready_agents}\n"
|
||||
f"Percentage ready for enforcement: {ready_percentage:.2f}%"
|
||||
)
|
||||
logger.info(message)
|
||||
colorText(message, "green")
|
||||
@@ -0,0 +1,159 @@
|
||||
# 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
|
||||
from utils.utils import colorText
|
||||
from sys import exit
|
||||
|
||||
# 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.")
|
||||
print(colorText("❌ Authentication failed. Exiting.","red"))
|
||||
exit(1) # Exit cleanly without traceback
|
||||
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")
|
||||
|
||||
while True:
|
||||
password = getpass("🔓 Create a password to encrypt your API key: ")
|
||||
confirm_password = getpass("🔒 Confirm your password: ")
|
||||
|
||||
if password != confirm_password:
|
||||
logging.warning("❌ Passwords do not match. Try again.")
|
||||
continue
|
||||
|
||||
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:
|
||||
logging.warning("❌ Password does not meet complexity requirements. Try again.")
|
||||
|
||||
|
||||
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,171 @@
|
||||
# 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, Callable, List, Optional, Union
|
||||
|
||||
from utils.utils import get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Selector:
|
||||
@staticmethod
|
||||
def _display_choices(
|
||||
items: List[Any],
|
||||
label_func: Callable[[Any], str],
|
||||
num_columns: int = 3,
|
||||
header: str = "Available Choices:"
|
||||
) -> None:
|
||||
sorted_items = sorted(items, key=lambda item: label_func(item).lower())
|
||||
rows = (len(sorted_items) + num_columns - 1) // num_columns
|
||||
print(f"\n{header}")
|
||||
for row in range(rows):
|
||||
line = ""
|
||||
for col in range(num_columns):
|
||||
idx = row + col * rows
|
||||
if idx < len(sorted_items):
|
||||
label = label_func(sorted_items[idx])
|
||||
line += f"{idx + 1}: {label:<30}"
|
||||
print(line)
|
||||
|
||||
@staticmethod
|
||||
def _select_from_list(
|
||||
items: List[Any],
|
||||
label_func: Callable[[Any], str],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False,
|
||||
header: str = "Available Choices:"
|
||||
) -> Union[Optional[Any], List[Any]]:
|
||||
if not items:
|
||||
logger.warning("No items available for selection.")
|
||||
return None
|
||||
|
||||
Selector._display_choices(items, label_func, header=header)
|
||||
sorted_items = sorted(items, key=lambda item: label_func(item).lower())
|
||||
selected = []
|
||||
|
||||
if allow_multiple:
|
||||
while True:
|
||||
choice = get_sanitized_input("Select an item by number (or Q to finish): ").strip().lower()
|
||||
if choice == "q":
|
||||
break
|
||||
try:
|
||||
index = int(choice)
|
||||
if 1 <= index <= len(sorted_items):
|
||||
item = sorted_items[index - 1]
|
||||
if item not in selected:
|
||||
selected.append(item)
|
||||
if prompt_each:
|
||||
logger.info(f"Selected: {label_func(item)}")
|
||||
else:
|
||||
logger.warning("Item 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(get_sanitized_input("Select one item by number: "))
|
||||
if 1 <= choice <= len(sorted_items):
|
||||
selected_item = sorted_items[choice - 1]
|
||||
logger.info(f"Selected: {label_func(selected_item)}")
|
||||
return selected_item
|
||||
else:
|
||||
logger.warning("Selection out of range.")
|
||||
except ValueError:
|
||||
logger.warning("Invalid input.")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def select_objects(
|
||||
objects: List[Any],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False
|
||||
) -> Union[Optional[Any], List[Any]]:
|
||||
return Selector._select_from_list(
|
||||
objects,
|
||||
label_func=lambda obj: getattr(obj, "name", str(obj)),
|
||||
allow_multiple=allow_multiple,
|
||||
prompt_each=prompt_each,
|
||||
header="Available Objects:"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def select_string(
|
||||
options: List[str],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False
|
||||
) -> Union[Optional[str], List[str]]:
|
||||
return Selector._select_from_list(
|
||||
options,
|
||||
label_func=str,
|
||||
allow_multiple=allow_multiple,
|
||||
prompt_each=prompt_each,
|
||||
header="Available Options:"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def select_int(
|
||||
options: List[int],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False
|
||||
) -> Union[Optional[int], List[int]]:
|
||||
return Selector._select_from_list(
|
||||
options,
|
||||
label_func=lambda x: str(x),
|
||||
allow_multiple=allow_multiple,
|
||||
prompt_each=prompt_each,
|
||||
header="Available Integers:"
|
||||
)
|
||||
|
||||
@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 = get_sanitized_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 = get_sanitized_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,175 @@
|
||||
# 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
|
||||
|
||||
from utils.configmanager import PROTECTED_KEYS, load_protected_config
|
||||
|
||||
|
||||
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))
|
||||
|
||||
# 🔧 Clear existing handlers
|
||||
for handler in logger.handlers[:]:
|
||||
logger.removeHandler(handler)
|
||||
|
||||
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():
|
||||
if key in PROTECTED_KEYS:
|
||||
continue # Skip protected keys
|
||||
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": []
|
||||
}
|
||||
|
||||
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}
|
||||
|
||||
protected_config = load_protected_config()
|
||||
merged_config.update(protected_config)
|
||||
|
||||
# ✅ 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
|
||||
@@ -0,0 +1,708 @@
|
||||
# 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
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
import tkinter as tk
|
||||
from tkinter import filedialog
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from utils.configmanager import load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
||||
|
||||
def import_to_dataframe(file_path: str) -> pd.DataFrame:
|
||||
df = pd.DataFrame()
|
||||
|
||||
try:
|
||||
if not os.path.exists(file_path):
|
||||
print(colorText(f"Error: File '{file_path}' does not exist.", "red"))
|
||||
return df
|
||||
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
|
||||
if ext == ".csv":
|
||||
df = pd.read_csv(file_path)
|
||||
elif ext == ".parquet":
|
||||
df = pd.read_parquet(file_path)
|
||||
else:
|
||||
print(colorText(f"Error: Unsupported file extension '{ext}'.", "red"))
|
||||
return df
|
||||
|
||||
if df.empty:
|
||||
print(colorText("Error: File has headers but no data rows.", "red"))
|
||||
else:
|
||||
print(colorText(f"Data loaded successfully from {file_path}", "green"))
|
||||
|
||||
return df
|
||||
|
||||
except pd.errors.EmptyDataError:
|
||||
print(
|
||||
colorText(
|
||||
"Notice: CSV file is completely empty, falling back to empty frame",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
return pd.DataFrame()
|
||||
|
||||
except Exception as e:
|
||||
print(colorText(f"Error reading file: {e}", "red"))
|
||||
return pd.DataFrame()
|
||||
|
||||
|
||||
def choose_directory():
|
||||
root = tk.Tk()
|
||||
root.withdraw() # Hide the main window
|
||||
directory = filedialog.askdirectory(title="Select a Directory")
|
||||
print("Selected directory:", directory)
|
||||
return directory
|
||||
|
||||
|
||||
def choose_file(initial_directory=None, required_substring=None):
|
||||
"""Open a file dialog and ensure the selected file contains a required substring."""
|
||||
while True:
|
||||
root = tk.Tk()
|
||||
root.withdraw() # Hide the main window
|
||||
file_path = filedialog.askopenfilename(initialdir=initial_directory)
|
||||
|
||||
if not file_path:
|
||||
print("No file selected.")
|
||||
return None
|
||||
|
||||
if required_substring and required_substring not in file_path:
|
||||
print(
|
||||
f"The selected file must contain '{required_substring}' in its path or name. Please try again."
|
||||
)
|
||||
else:
|
||||
return file_path
|
||||
|
||||
|
||||
|
||||
|
||||
def get_sanitized_input(prompt: str) -> str:
|
||||
while True:
|
||||
user_input = input(prompt)
|
||||
if user_input.strip() == "":
|
||||
return user_input # Allow blank lines
|
||||
if re.match(r'^[a-zA-Z0-9_\- .]+$', user_input.strip()):
|
||||
return user_input
|
||||
else:
|
||||
print("Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed.")
|
||||
|
||||
|
||||
def regulator(paths, case_insensitive=True):
|
||||
"""
|
||||
Build a regex pattern that matches any of the given Windows path fragments.
|
||||
"""
|
||||
escaped = [re.escape(p) for p in paths]
|
||||
pattern = "(?:" + "|".join(escaped) + ")"
|
||||
if case_insensitive:
|
||||
pattern = "(?i)" + pattern # Add inline case-insensitive flag
|
||||
print(f"Regulator is providing: {pattern}")
|
||||
return pattern
|
||||
|
||||
def displayIntro():
|
||||
print(
|
||||
colorText(
|
||||
r"""
|
||||
███
|
||||
████ ░████████
|
||||
█████████████ ███████████████
|
||||
█████████████████████ █████████████████████
|
||||
███████████████████ ██████████████████████▓
|
||||
███████████████████ ██████████████████████
|
||||
█████████████████████ ███████████████████████
|
||||
████████████████████████████████████████████████████████
|
||||
█████████ ██ ██ █████████
|
||||
█████████ ██ ███ █ █████████
|
||||
█████████ ██ ████ █████ █████████████
|
||||
█████████ ██ ██████ █████████████
|
||||
████████ ██ ███████ ████████████░
|
||||
███████ ██ ██▓ ██████ ████████████
|
||||
██████ ██ ████ █████ ███████████
|
||||
█████████████████████████████████████████████████
|
||||
▒████████████████████ ██████████████████
|
||||
███████████████████ ███████████████▒
|
||||
███████████████ █████████████
|
||||
██████████ ███████████
|
||||
████████
|
||||
████
|
||||
""",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
r"""
|
||||
_____ .__ .__ __ ___________ .__
|
||||
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
|
||||
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
|
||||
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
|
||||
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
|
||||
\/ \/ \/ \/
|
||||
""",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"=================================================================================",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"======================== Welcome to the Airlock API Tool ========================",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"=================================================================================",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
print(
|
||||
colorText(
|
||||
"\n --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"\nSequentually follow these steps to prepare a policy for enforcement:",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"\n1. Choose which originating policy or policies to move to enforcement",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
if not selected_policies:
|
||||
print(colorText(" [✗] No policies have been chosen", "red"))
|
||||
else:
|
||||
print(colorText("The following policies have been choosen:", "green"))
|
||||
for policy in selected_policies:
|
||||
print(colorText(f" [✓] {policy.name}", "green"))
|
||||
|
||||
print(colorText("2. Choose the destination policy and allowlist", "cyan"))
|
||||
|
||||
if not destination_policy:
|
||||
print(colorText(" [✗] No destination policy has been chosen", "red"))
|
||||
elif destination_policy:
|
||||
print(colorText(f" [✓] {destination_policy[0].name} has been selected as the destination policy", "green"))
|
||||
|
||||
|
||||
|
||||
if not destination_allowlist:
|
||||
print(colorText(" [✗] No allowlist has been chosen", "red"))
|
||||
elif destination_allowlist:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"3. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
if not selected_policies:
|
||||
print(colorText(" [✗] No policies have been chosen", "red"))
|
||||
else:
|
||||
if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
|
||||
print(colorText(" [✓] Data has been fetched", "green"))
|
||||
else:
|
||||
print(colorText(" [✗] Data has not been fetched", "red"))
|
||||
|
||||
print(colorText("4. Manually review the files:", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" 'prepare_policy\\needs_approved\\good_hashes.csv' and 'prepare_policy\\needs_approved\\unknown_hashes.csv'\n",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" If metarules need to be created, please make note of them, and remove the row from the csv.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" When complete, save both csv files to the directory 'approved' and choose this option.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
if os.path.exists(f"{working_dir}\\Approved\\approved_executions.csv"):
|
||||
print(colorText(" [✓] Reviewed hashes have been loaded", "green"))
|
||||
else:
|
||||
print(colorText(" [✗] Reviewed hashes have not been loaded", "red"))
|
||||
|
||||
if os.path.exists(
|
||||
f"{working_dir}\\Needs_Review\\Review_Second\\primary_Paths.csv",
|
||||
):
|
||||
print(colorText(" [✓] Path review list created", "green"))
|
||||
else:
|
||||
print(colorText(" [✗] Path review list has not been created", "red"))
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"5. Manually review the files \n 'prepare_policy\\needs_approved\\primary_Paths.csv'\n 'prepare_policy\\needs_approved\\secondary_Paths.csv'",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing path exclusions you do not approve of. The secondary list can be not added at all if nothing is useful",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" Do the same process with the list of publishers forthe same directories",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" Preflight Lists will be generated", "cyan"))
|
||||
|
||||
if os.path.exists(
|
||||
f"{working_dir}\\Approved\\primary_Paths.csv",
|
||||
):
|
||||
print(colorText(" [✓] Reviewed path list detected", "green"))
|
||||
else:
|
||||
print(colorText(" [✗] Path review list has not been detected", "red"))
|
||||
|
||||
if os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv") and os.path.exists(
|
||||
f"{working_dir}\\Preflight\\approved_hashes.csv"
|
||||
):
|
||||
print(colorText(" [✓] Preflight Path Exclusion List has been generated", "green"))
|
||||
else:
|
||||
print(colorText(" [✗] Preflight Path Exclusion List has not been generated", "red"))
|
||||
|
||||
print(colorText("6. Test ------------------------------------------------------", "cyan"))
|
||||
print(colorText(" Print rather than apply selected data.", "cyan"))
|
||||
|
||||
print(colorText("7. Liftoff ------------------------------------------------------", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" Apply path exclusions according to allowed and approved paths",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan"))
|
||||
print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"R. Remove/Reset Generated data - will prompt to allow keeping execution history",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText("F. 📂 - Open Working Directory", "cyan"))
|
||||
print(colorText("Q. 🔚 - Quit", "cyan"))
|
||||
|
||||
|
||||
def areYouSure():
|
||||
print(
|
||||
colorText(
|
||||
"🛑****************************************************************************************************************************************🛑",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"⚠️=========================================================================================================================================⚠️",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"🛑========================================================================================================================================🛑",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"🛑========================================================================================================================================🛑",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"⚠️=========================================================================================================================================⚠️",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"🛑****************************************************************************************************************************************🛑",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def locked():
|
||||
print(
|
||||
colorText(
|
||||
r"""
|
||||
████████████████████████████████████████████████████████████████
|
||||
███ ██
|
||||
██ ██████ ███
|
||||
██ ████████████ ███
|
||||
██ ████ ███ ███
|
||||
██ ███ ███ ███
|
||||
██ ███ ███ ███
|
||||
██ ▒████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ███
|
||||
███ ███
|
||||
████████████████████████████████████████████████████████████████████
|
||||
▒██████████████████████████████████████████████████████████████████▒
|
||||
▒████
|
||||
▒████
|
||||
▓██████████████████████████████████████████
|
||||
█████████████████████████████████████████████░
|
||||
""",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def printDeviceEnforceChecklist():
|
||||
print(
|
||||
colorText(
|
||||
"\n --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"\nSequentually follow these steps to prepare a policy for enforcement:",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"\n1. Choose which originating policy or policies to move to enforcement",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText("3. Manually review the files:", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" If metarules need to be created, please make note of them, and remove the row from the csv.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" When complete, save both csv files to the directory 'approved' and choose this option.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"4. Manually review the file 'needs_approved\\paths_needing_review.csv'",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing path exclusions you do not approve of",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" Do the same process with the list of publishers forthe same directories",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" Preflight Lists will be generated", "cyan"))
|
||||
|
||||
print(colorText("5. Choose the destination policy and parent and child allow list", "cyan"))
|
||||
|
||||
print(colorText("6. Test ------------------------------------------------------", "cyan"))
|
||||
print(colorText(" Print rather than apply selected data.", "cyan"))
|
||||
|
||||
print(colorText("7. Liftoff ------------------------------------------------------", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" Apply path exclusions according to allowed and approved paths",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan"))
|
||||
print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"R. Remove/Reset Generated data - will prompt to allow keeping execution history",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
print(colorText("B. Back", "cyan"))
|
||||
|
||||
|
||||
def colorText(text, color):
|
||||
colors = {
|
||||
"red": "\033[91m",
|
||||
"green": "\033[92m",
|
||||
"yellow": "\033[93m",
|
||||
"blue": "\033[94m",
|
||||
"magenta": "\033[95m",
|
||||
"cyan": "\033[96m",
|
||||
"white": "\033[97m",
|
||||
"reset": "\033[0m",
|
||||
}
|
||||
return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"
|
||||
|
||||
|
||||
def formatHTML(df, output_html_path=None, overwrite=True):
|
||||
from datetime import datetime
|
||||
|
||||
# Get current date and filename for subtitle
|
||||
today = datetime.now().strftime("%d %B %Y") # Changed to "Day Month Year"
|
||||
filename = output_html_path.replace(".html", "") if output_html_path else "Report"
|
||||
|
||||
dark_css = """
|
||||
<style>
|
||||
body {
|
||||
background-color: #000000;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
|
||||
color: #f8f8f2;
|
||||
}
|
||||
.header {
|
||||
text-align: center;
|
||||
margin: 20px auto;
|
||||
padding: 10px;
|
||||
border-bottom: 2px solid #ffd700;
|
||||
max-width: 95%;
|
||||
}
|
||||
.header h1 {
|
||||
color: #ffd700;
|
||||
margin: 0;
|
||||
font-size: 32px;
|
||||
}
|
||||
.header p {
|
||||
color: #00bfff;
|
||||
margin: 5px 0 0 0;
|
||||
font-size: 18px;
|
||||
}
|
||||
.table-container {
|
||||
overflow-y: scroll;
|
||||
margin: 0 auto;
|
||||
width: 95%;
|
||||
max-height: calc(80vh - 100px);
|
||||
display: block;
|
||||
border: 1px solid #3a3a4d;
|
||||
margin-bottom: 0;
|
||||
}
|
||||
table {
|
||||
border-collapse: collapse;
|
||||
font-size: 14px;
|
||||
background-color: #1e1e2f;
|
||||
color: #f8f8f2;
|
||||
width: max-content;
|
||||
}
|
||||
th, td {
|
||||
border: 1px solid #3a3a4d;
|
||||
text-align: left;
|
||||
padding: 10px;
|
||||
max-width: 300px;
|
||||
word-wrap: break-word;
|
||||
overflow-wrap: break-word;
|
||||
}
|
||||
/* First column: no wrap */
|
||||
td:nth-child(1), th:nth-child(1) {
|
||||
white-space: nowrap;
|
||||
max-width: none !important;
|
||||
word-wrap: normal !important;
|
||||
}
|
||||
th {
|
||||
background-color: #2e2e40;
|
||||
color: #ffd700;
|
||||
position: sticky;
|
||||
top: 0;
|
||||
z-index: 10;
|
||||
}
|
||||
tr:nth-child(even) {
|
||||
background-color: #262638;
|
||||
}
|
||||
tr:hover {
|
||||
background-color: #33334d;
|
||||
color: #00bfff;
|
||||
}
|
||||
/* Custom scrollbar styling */
|
||||
.table-container::-webkit-scrollbar {
|
||||
width: 12px;
|
||||
}
|
||||
.table-container::-webkit-scrollbar-track {
|
||||
background: #1e1e2f;
|
||||
}
|
||||
.table-container::-webkit-scrollbar-thumb {
|
||||
background-color: #3a3a4d;
|
||||
border-radius: 6px;
|
||||
}
|
||||
</style>
|
||||
"""
|
||||
|
||||
header = f"""
|
||||
<div class="header">
|
||||
<h1>Airlock Tools</h1>
|
||||
<p>{filename} - {today}</p>
|
||||
</div>
|
||||
"""
|
||||
|
||||
html_table = df.to_html(index=False, escape=False)
|
||||
styled_html = (
|
||||
f"<html>\n"
|
||||
f"<head><title>Airlock Tools Report</title></head>\n"
|
||||
f"<body>\n"
|
||||
f"{dark_css}\n"
|
||||
f"{header}\n"
|
||||
f"<div class='table-container'>\n"
|
||||
f" {html_table}\n"
|
||||
f"</div>\n"
|
||||
f"</body>\n"
|
||||
f"</html>"
|
||||
)
|
||||
if output_html_path:
|
||||
with open(output_html_path, "w", encoding="utf-8") as f:
|
||||
f.write(styled_html)
|
||||
print(f"✅ Styled table saved to '{output_html_path}'")
|
||||
elif overwrite:
|
||||
with tempfile.NamedTemporaryFile(
|
||||
suffix=".html", delete=False, mode="w", encoding="utf-8"
|
||||
) as f:
|
||||
f.write(styled_html)
|
||||
temp_path = f.name
|
||||
|
||||
print(f"✅ Styled table saved to temporary file: {temp_path}")
|
||||
else:
|
||||
return styled_html
|
||||
|
||||
|
||||
|
||||
def open_directory(path):
|
||||
system = platform.system()
|
||||
|
||||
if system == "Windows":
|
||||
os.startfile(path)
|
||||
elif system == "Linux":
|
||||
subprocess.run(["xdg-open", path])
|
||||
else:
|
||||
raise OSError(f"Unsupported operating system: {system}")
|
||||
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Callable, Any
|
||||
import psutil # For CPU and memory monitoring
|
||||
|
||||
# Set up the logger
|
||||
logger = logging.getLogger("AsyncTaskQueue")
|
||||
logger.setLevel(logging.INFO)
|
||||
handler = logging.StreamHandler()
|
||||
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||
handler.setFormatter(formatter)
|
||||
logger.addHandler(handler)
|
||||
|
||||
class AsyncTaskQueue:
|
||||
def __init__(
|
||||
self,
|
||||
min_workers: int = 1,
|
||||
max_workers: int = 10,
|
||||
cpu_threshold: float = 90.0,
|
||||
memory_threshold: float = 90.0,
|
||||
):
|
||||
self.queue = asyncio.Queue()
|
||||
self.min_workers = min_workers
|
||||
self.max_workers = max_workers
|
||||
self.cpu_threshold = cpu_threshold
|
||||
self.memory_threshold = memory_threshold
|
||||
self.workers = []
|
||||
self._stop_event = asyncio.Event()
|
||||
self._scale_lock = asyncio.Lock() # Prevent race conditions
|
||||
|
||||
async def _scale_up(self):
|
||||
"""Add a worker if under max limit and resources allow."""
|
||||
async with self._scale_lock:
|
||||
if (
|
||||
len(self.workers) < self.max_workers
|
||||
and psutil.cpu_percent() < self.cpu_threshold
|
||||
and psutil.virtual_memory().percent < self.memory_threshold
|
||||
):
|
||||
worker = asyncio.create_task(self.worker_loop(f"Worker-{len(self.workers) + 1}"))
|
||||
self.workers.append(worker)
|
||||
logger.info(f"Scaled up. Workers: {len(self.workers)}")
|
||||
|
||||
async def _scale_down(self):
|
||||
"""Remove a worker if above min limit."""
|
||||
async with self._scale_lock:
|
||||
if len(self.workers) > self.min_workers:
|
||||
worker = self.workers.pop()
|
||||
worker.cancel()
|
||||
logger.info(f"Scaled down. Workers: {len(self.workers)}")
|
||||
|
||||
async def _monitor_resources(self):
|
||||
"""Monitor CPU and memory usage, and adjust workers."""
|
||||
while not self._stop_event.is_set():
|
||||
cpu_usage = psutil.cpu_percent()
|
||||
memory_usage = psutil.virtual_memory().percent
|
||||
|
||||
if (
|
||||
cpu_usage > self.cpu_threshold
|
||||
or memory_usage > self.memory_threshold
|
||||
):
|
||||
await self._scale_down()
|
||||
elif (
|
||||
len(self.workers) < self.max_workers
|
||||
and self.queue.qsize() > 2 # Only scale up if there's work
|
||||
):
|
||||
await self._scale_up()
|
||||
|
||||
await asyncio.sleep(2) # Check every 2 seconds
|
||||
|
||||
async def start_workers(self):
|
||||
"""Start initial workers and the resource monitor."""
|
||||
for i in range(self.min_workers):
|
||||
worker = asyncio.create_task(self.worker_loop(f"Worker-{i+1}"))
|
||||
self.workers.append(worker)
|
||||
# Start the resource monitor
|
||||
asyncio.create_task(self._monitor_resources())
|
||||
logger.info(f"Started {self.min_workers} workers and resource monitor.")
|
||||
|
||||
async def stop_workers(self):
|
||||
"""Stop all workers and the resource monitor."""
|
||||
logger.info("Stopping workers...")
|
||||
self._stop_event.set()
|
||||
await self.queue.join() # Wait for all tasks to complete
|
||||
for worker in self.workers:
|
||||
worker.cancel()
|
||||
await asyncio.gather(*self.workers, return_exceptions=True)
|
||||
logger.info("All workers stopped.")
|
||||
|
||||
async def worker_loop(self, name: str):
|
||||
"""Worker loop: Process tasks from the queue."""
|
||||
logger.info(f"{name} started.")
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
task = await self.queue.get()
|
||||
logger.info(f"{name} processing: {task['name']}")
|
||||
await task['func'](*task['args'])
|
||||
except Exception as e:
|
||||
logger.error(f"Error in {name}: {e}", exc_info=True)
|
||||
finally:
|
||||
self.queue.task_done()
|
||||
logger.info(f"{name} exited.")
|
||||
|
||||
async def enqueue(self, name: str, func: Callable, *args: Any):
|
||||
"""Add a task to the queue."""
|
||||
logger.info(f"Enqueuing task: {name}")
|
||||
await self.queue.put({'name': name, 'func': func, 'args': args})
|
||||
|
||||
async def run_sync_task_in_thread(func: Callable, *args: Any):
|
||||
"""Run a synchronous function in a separate thread."""
|
||||
await asyncio.to_thread(func, *args)
|
||||
"""import asyncio
|
||||
|
||||
async def example_async_task(name: str, duration: int):
|
||||
print(f"{name} started, will sleep for {duration} seconds")
|
||||
await asyncio.sleep(duration)
|
||||
print(f"{name} finished")
|
||||
|
||||
async def main():
|
||||
queue = AsyncTaskQueue(
|
||||
min_workers=2,
|
||||
max_workers=10,
|
||||
cpu_threshold=90.0,
|
||||
memory_threshold=90.0,
|
||||
)
|
||||
await queue.start_workers()
|
||||
|
||||
# Enqueue tasks
|
||||
for i in range(20):
|
||||
await queue.enqueue(f"Task{i}", example_async_task, f"Task{i}", 1)
|
||||
|
||||
await asyncio.sleep(10) # Let tasks run
|
||||
await queue.stop_workers()
|
||||
|
||||
asyncio.run(main())
|
||||
"""
|
||||
@@ -0,0 +1,179 @@
|
||||
from textual.app import App, ComposeResult
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Header, Tabs, Tab, Static, Footer, DirectoryTree, Button
|
||||
from textual.containers import Horizontal
|
||||
import logging
|
||||
import dotenv
|
||||
import os
|
||||
|
||||
from Development_Stubs.WIP.localApproval import moveToLocalApproval
|
||||
from flows.prepPolicy import (
|
||||
buildPathsandPublishers,
|
||||
buildPreflights,
|
||||
selectAllowlists,
|
||||
selectPolicies,
|
||||
sortHashes,
|
||||
)
|
||||
from flows.quietAgent import findQuietAgents
|
||||
from services.agenthandler import findAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
import services.policyhandler as policyh
|
||||
from utils.utils import (
|
||||
areYouSure,
|
||||
colorText,
|
||||
displayIntro,
|
||||
load_env,
|
||||
open_directory,
|
||||
printEnforceChecklist,
|
||||
)
|
||||
from screens.agent_results import AgentResultsScreen
|
||||
logger = logging.getLogger(__name__)
|
||||
dotenv.load_dotenv()
|
||||
|
||||
ASCII_ART = displayIntro()
|
||||
|
||||
class BaseScreen(Screen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Static("This is a dummy screen.", id="content")
|
||||
yield Footer()
|
||||
|
||||
class LandingScreen(BaseScreen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Static(ASCII_ART, id="ascii-art", markup=False)
|
||||
yield Horizontal(
|
||||
Button("Get Started", id="get-started-button"),
|
||||
Button("Settings", id="settings-button"),
|
||||
)
|
||||
yield Footer()
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
if event.button.id == "get-started-button":
|
||||
# Go back to the previous screen (MainMenuScreen)
|
||||
self.app.pop_screen()
|
||||
|
||||
class LocalApprovalScreen(BaseScreen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
|
||||
yield Footer()
|
||||
|
||||
class OTPScreen(BaseScreen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Static("OTP: Enter your one-time password here.", id="content")
|
||||
yield Footer()
|
||||
|
||||
class SearchScreen(BaseScreen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Static("Search: Find items in the system.", id="content")
|
||||
yield Horizontal(
|
||||
Button("Search for Agents", id="searchAgent-button"),
|
||||
)
|
||||
yield Footer()
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
if event.button.id == "searchAgent-button":
|
||||
self.app.push_screen(AgentResultsScreen(self.app.api, self.app.working_dir))
|
||||
|
||||
class QuietHostsScreen(BaseScreen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Static("Quiet Hosts: Manage hosts that are not responding.", id="content")
|
||||
yield Footer()
|
||||
|
||||
class PolicyPrepScreen(BaseScreen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Static("Policy Prep: Prepare policies for deployment.", id="content")
|
||||
yield Footer()
|
||||
|
||||
class UpdatePoliciesScreen(BaseScreen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Static("Update Policies: Update existing policies.", id="content")
|
||||
yield Footer()
|
||||
|
||||
class SettingsScreen(BaseScreen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Static("Settings: Configure application settings.", id="content")
|
||||
yield Footer()
|
||||
|
||||
class DirectoryTreeScreen(Screen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield DirectoryTree("./")
|
||||
yield Footer()
|
||||
|
||||
class MainMenuScreen(Screen):
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
yield Tabs(
|
||||
Tab("Home", id="home_screen"),
|
||||
Tab("Local Approval", id="local_approval"),
|
||||
Tab("OTP", id="otp"),
|
||||
Tab("Search", id="search"),
|
||||
Tab("Quiet Hosts", id="quiet"),
|
||||
Tab("Policy Prep", id="policy"),
|
||||
Tab("Update Policies", id="update"),
|
||||
id="tabs"
|
||||
)
|
||||
yield Static("Select a tab to begin.", id="content")
|
||||
yield Footer()
|
||||
|
||||
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
|
||||
match event.tab.id:
|
||||
case "home_screen":
|
||||
self.app.push_screen(LandingScreen())
|
||||
case "local_approval":
|
||||
self.app.push_screen(LocalApprovalScreen())
|
||||
case "otp":
|
||||
self.app.push_screen(OTPScreen())
|
||||
case "search":
|
||||
self.app.push_screen(SearchScreen())
|
||||
case "quiet":
|
||||
self.app.push_screen(QuietHostsScreen())
|
||||
case "policy":
|
||||
self.app.push_screen(PolicyPrepScreen())
|
||||
case "update":
|
||||
self.app.push_screen(UpdatePoliciesScreen())
|
||||
case "settings":
|
||||
self.app.push_screen(SettingsScreen())
|
||||
|
||||
class AirlockTools(App):
|
||||
BINDINGS = [
|
||||
("q", "quit", "Quit"),
|
||||
("d", "open_dir", "Open Directory"),
|
||||
("b", "back", "Go Back"),
|
||||
]
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper, working_dir: str):
|
||||
super().__init__()
|
||||
self.api = api
|
||||
self.working_dir = working_dir
|
||||
|
||||
def on_mount(self) -> None:
|
||||
self.push_screen(MainMenuScreen())
|
||||
|
||||
def action_quit(self) -> None:
|
||||
self.exit()
|
||||
|
||||
def action_open_dir(self) -> None:
|
||||
self.push_screen(DirectoryTreeScreen())
|
||||
|
||||
def action_back(self) -> None:
|
||||
if len(self.screen_stack) > 2:
|
||||
self.pop_screen()
|
||||
else:
|
||||
self.bell()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
api = AirlockAPIWrapper(base_url=str(os.getenv("URL")),api_key = "b82ab7e2b12430fff0cbf0e798f79e16c1aea2c37ffee5dc884de7ba11d3a3a4")
|
||||
working_dir = "./" # Or a dynamic path
|
||||
AirlockTools(api, working_dir).run()
|
||||
|
||||
|
||||
+16
-14
@@ -32,6 +32,8 @@ SALT_SIZE = 16 # 128-bit Salt
|
||||
NONCE_SIZE = 12 # AES-GCM
|
||||
KEY_SIZE = 32 # AES-256
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _derive_key(password: bytes, salt: bytes) -> bytes:
|
||||
kdf = PBKDF2HMAC(
|
||||
@@ -67,20 +69,6 @@ def store_api_key(service: str, username: str, api_key: str, password: str):
|
||||
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()
|
||||
|
||||
logger.debug(f"API key for service '{service}' and user '{username}' stored successfully.")
|
||||
|
||||
print("\n✅ API key stored securely.")
|
||||
@@ -95,6 +83,20 @@ def retrieve_api_key(service: str, username: str, password: str) -> str:
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user