RustImplementation #25

Merged
mysticmomba merged 13 commits from RustImplementation into master 2025-11-07 11:39:51 -05:00
35 changed files with 2145 additions and 807 deletions
+52 -50
View File
@@ -1,50 +1,52 @@
name: Build EXE and Release ##name: Build EXE and Release
run-name: ${{ gitea.actor }} ##run-name: ${{ gitea.actor }}
on: ##on:
push: ## push:
branches: ## branches:
- master ## - master
##
jobs: ##jobs:
Build and Release: ## Build and Release:
runs-on: debian-stable ## runs-on: debian-stable
env: ## env:
DISPLAY: :99 ## DISPLAY: :99
##
steps: ## steps:
- name: Install Prerequisites ## - name: Install Prerequisites
run: | ## run: |
dpkg --add-architecture i386 ## dpkg --add-architecture i386
apt update > /dev/null 2>&1 ## apt update > /dev/null 2>&1
apt install git curl wine32:i386 xvfb -y > /dev/null 2>&1 ## apt install git curl wine32:i386 xvfb -y > /dev/null 2>&1
##
- name: Start X Virtual Framebuffer (Xvfb) ## - name: Start X Virtual Framebuffer (Xvfb)
run: | ## run: |
# Start Xvfb in the background using the defined display number ## # Start Xvfb in the background using the defined display number
Xvfb :99 -screen 0 1024x768x16 & ## Xvfb :99 -screen 0 1024x768x16 &
##
- name: Cloning Repository ## - name: Cloning Repository
run: | ## run: |
git clone https://brotoskyj:${{ secrets.RUNNER_TOKEN }}@git.racooncity.org/brotoskyj/AirlockTools --branch RustImplementation ## git clone https://brotoskyj:${{ secrets.RUNNER_TOKEN }}@git.racooncity.org/brotoskyj/AirlockTools --branch RustImplementation
pwd ## pwd
ls ## ls
##
- name: Setting Up Build Environment ## - name: Setting Up Build Environment
run: | ## run: |
apt install wine64 -y -qq > /dev/null 2>&1 ## apt install wine64 -y -qq > /dev/null 2>&1
curl -L -o python.exe https://www.python.org/ftp/python/3.13.9/python-3.13.9-amd64.exe ## curl -L -o python.exe https://www.python.org/ftp/python/3.13.9/python-3.13.9-amd64.exe
wine python.exe /quiet /NoWeb InstallAllUsers=1 PrependPath=1 TargetDir=C:/Python313 ## wget https://download.visualstudio.microsoft.com/download/pr/0c8b0c6f-3d30-4d06-98d8-3a7a42f3e78a/64c0b0f3b8b2d0e11d9a2303d5e39a22/ndp472-devpack-ENU.exe -O dotnet472.exe
cp -r AirlockTools/ ~/.wine/drive_c/Python313/ ## wget https://aka.ms/vs/16/release/vs_buildtools.exe -O vstools2019.exe
cd ~/.wine/drive_c/Python313 ## wine dotnet472.exe /q /norestart /ChainingPackage ADMINDEPLOYMENT /log dotnet472.log
wine python.exe -m pip install nuitka pywin32 keyring --break-system-packages ## wine python.exe /quiet /NoWeb InstallAllUsers=1 PrependPath=1 TargetDir=C:/Python313
wine python.exe -m pip install -r AirlockTools/requirements.txt ## wine vstools2019.exe --quiet --wait --norestart --nocache --installPath "C:\\BuildTools" --add Microsoft.VisualStudio.Workload.VCTools --includeRecommended --acceptEula
wine python.exe -m pip install \ ## cp -r AirlockTools/ ~/.wine/drive_c/Python313/
--index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ \ ## cd ~/.wine/drive_c/Python313
airlock-libs --break-system-packages ## wine python.exe -m pip install nuitka pywin32 keyring --break-system-packages
## wine python.exe -m pip install -r AirlockTools/requirements.txt
- name: Build Executable ##
env: ## - name: Build Executable
DISPLAY: :99 ## env:
WINEDEBUG: -all ## DISPLAY: :99
run: | ## WINEDEBUG: -all
xvfb-run -a wine cmd /c nuitka --onefile --onefile-windows-splash-screen-image=AirlockTools/loading.png --follow-imports --include-module=keyring.backends.Windows --include-module=win32cred --include-module=pywintypes --include-module=pythoncom --include-module=win32api --include-module=win32security --include-module=win32con --windows-product-name='LoXide' --windows-product-version='1.2.10.0' --windows-company-name='WVUM IRT' --windows-icon-from-ico=AirlockTools/IRT_icon_32-512.ico AirlockTools/AirlockTools_client.py --deployment --assume-yes-for-downloads ## run: |
## xvfb-run -a wine cmd /c nuitka --onefile --onefile-windows-splash-screen-image=AirlockTools/loading.png --mingw64 --follow-imports --include-module=keyring.backends.Windows --include-module=win32cred --include-module=pywintypes --include-module=pythoncom --include-module=win32api --include-module=win32security --include-module=win32con --windows-product-name='LoXide' --windows-product-version='1.2.10.0' --windows-company-name='WVUM IRT' --windows-icon-from-ico=AirlockTools/IRT_icon_32-512.ico AirlockTools/AirlockTools_client.py --deployment --assume-yes-for-downloads
##
+2
View File
@@ -2,6 +2,8 @@ name: Build Library
run-name: ${{ gitea.actor }} run-name: ${{ gitea.actor }}
on: on:
push: push:
branches:
- RustImplementation
paths: paths:
- airlock_libs/** - airlock_libs/**
+26 -31
View File
@@ -14,76 +14,71 @@
# along with this program. If not, see <https://www.gnu.org/licenses/>. # along with this program. If not, see <https://www.gnu.org/licenses/>.
#TODO Continue implementing logger # TODO Continue implementing logger
#TODO Add input sanitation and CSV injection prevention # TODO Add input sanitation and CSV injection prevention
#TODO Continue OTP and Local approval rewrites # TODO Continue OTP and Local approval rewrites
#TODO Explore pywin32 # TODO Explore pywin32
#TODO Fix Requirements.txt # TODO Fix Requirements.txt
#TODO Create Generic system_config.json for gitea # TODO Create Generic system_config.json for gitea
import logging import logging
import os import os
import tempfile import tempfile
import dotenv import dotenv
import urllib3 import urllib3
from services.API import AirlockAPIWrapper from services.API import AirlockAPIWrapper
from services.security import getAPI from services.security import getAPI
from utils.setup import get_base_directory, setup from utils.setup import get_base_directory, setup
from utils.TUI import run_AirlockTools from utils.TUI import run_Loxide
from utils.utils import irtang from utils.utils import irtang
urllib3.disable_warnings( urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
urllib3.exceptions.InsecureRequestWarning
)
def main(): def main():
if "NUITKA_ONEFILE_PARENT" in os.environ: if "NUITKA_ONEFILE_PARENT" in os.environ:
splash_filename = os.path.join( splash_filename = os.path.join(
tempfile.gettempdir(), tempfile.gettempdir(),
f"onefile_{int(os.environ['NUITKA_ONEFILE_PARENT'])}_splash_feedback.tmp" f"onefile_{int(os.environ['NUITKA_ONEFILE_PARENT'])}_splash_feedback.tmp",
) )
if os.path.exists(splash_filename): if os.path.exists(splash_filename):
os.unlink(splash_filename) os.unlink(splash_filename)
irtang() irtang()
#Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
setup() setup()
base_dir = get_base_directory() base_dir = get_base_directory()
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
dotenv.load_dotenv(dotenv_path=base_dir / ".env") dotenv.load_dotenv(dotenv_path=base_dir / ".env")
try: try:
url = os.getenv("URL") url = os.getenv("URL")
username = os.getenv("USERNAME") username = os.getenv("USERNAME")
if not url: if not url:
raise ValueError("Missing URL in environment variables.") raise ValueError("Missing URL in environment variables.")
if not username: if not username:
raise ValueError("Missing USERNAME in environment variables.") raise ValueError("Missing USERNAME in environment variables.")
logger.debug(f"Retrieved URL: {url}") logger.debug(f"Retrieved URL: {url}")
logger.debug(f"Retrieved Username: {username}") logger.debug(f"Retrieved Username: {username}")
except ValueError as e: except ValueError as e:
logger.error(f"Configuration error: {e}", exc_info=True) logger.error(f"Configuration error: {e}", exc_info=True)
raise raise
api_key = getAPI(username, "Loxide")
api_key = getAPI(username, "AirlockTools")
if api_key is None: if api_key is None:
raise ValueError("API key for AirlockTools is missing.") raise ValueError("API key for Loxide is missing.")
api = AirlockAPIWrapper( api = AirlockAPIWrapper(
base_url=str(os.getenv("URL")), base_url=str(os.getenv("URL")),
api_key=api_key, api_key=api_key,
) )
run_AirlockTools(api) run_Loxide(api)
if __name__ == "__main__": if __name__ == "__main__":
+38 -35
View File
@@ -14,12 +14,11 @@
# along with this program. If not, see <https://www.gnu.org/licenses/>. # along with this program. If not, see <https://www.gnu.org/licenses/>.
# TODO Add CSV injection prevention
#TODO Add CSV injection prevention # TODO Continue OTP and Local approval rewrites
#TODO Continue OTP and Local approval rewrites # TODO Explore pywin32
#TODO Explore pywin32 # TODO Fix Requirements.txt
#TODO Fix Requirements.txt # TODO Create Generic system_config.json for gitea
#TODO Create Generic system_config.json for gitea
import logging import logging
@@ -29,62 +28,66 @@ import dotenv
import urllib3 import urllib3
import flows.localApproval as la import flows.localApproval as la
from Server.scheduler_async import recurring_job, register_function, reload_jobs, start_scheduler from Server.scheduler_async import (
recurring_job,
register_function,
reload_jobs,
start_scheduler,
)
from services.API import AirlockAPIWrapper from services.API import AirlockAPIWrapper
from services.policyhandler import updateAuditPoliciesFromEnforcementPolices from services.policyhandler import updateAuditPoliciesFromEnforcementPolices
from services.security import getAPI from services.security import getAPI
from utils.setup import setup from utils.setup import setup
urllib3.disable_warnings( urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
urllib3.exceptions.InsecureRequestWarning
)
def main(): def main():
#Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
working_dir = setup() working_dir = setup()
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
dotenv.load_dotenv(dotenv_path=working_dir / ".env") dotenv.load_dotenv(dotenv_path=working_dir / ".env")
try: try:
url = os.getenv("URL") url = os.getenv("URL")
username = os.getenv("USERNAME") username = os.getenv("USERNAME")
if not url: if not url:
raise ValueError("Missing URL in environment variables.") raise ValueError("Missing URL in environment variables.")
if not username: if not username:
raise ValueError("Missing USERNAME in environment variables.") raise ValueError("Missing USERNAME in environment variables.")
logger.debug(f"Retrieved URL: {url}")
logger.debug(f"Retrieved Username: {username}")
logger.debug(f"Retrieved URL: {url}")
logger.debug(f"Retrieved Username: {username}")
except ValueError as e: except ValueError as e:
logger.error(f"Configuration error: {e}", exc_info=True) logger.error(f"Configuration error: {e}", exc_info=True)
raise raise
api = AirlockAPIWrapper( api = AirlockAPIWrapper(
base_url=str(os.getenv("URL")), base_url=str(os.getenv("URL")),
api_key = getAPI(username, "AirlockTools"), api_key=getAPI(username, "AirlockTools"),
) )
logger.info("Running non-interactively to start monitoring Airlock Changes") logger.info("Running non-interactively to start monitoring Airlock Changes")
register_function("monitorLA", la.scheduleAddingLAHashes) register_function("monitorLA", la.scheduleAddingLAHashes)
register_function("updateAuditPolicies", updateAuditPoliciesFromEnforcementPolices) register_function("updateAuditPolicies", updateAuditPoliciesFromEnforcementPolices)
if not os.path.exists("scheduling\\jobs.json"):
if not os.path.exists("scheduling\\jobs.json"):
recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[api]) recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[api])
recurring_job("updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api]) recurring_job(
"updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api]
)
else: else:
reload_jobs() reload_jobs()
start_scheduler() start_scheduler()
if __name__ == "__main__": if __name__ == "__main__":
main() main()
+3 -3
View File
@@ -1,7 +1,7 @@
# 🛡️ Airlock Tools # 🛡️ Loxide
Python toolkit for secure, auditable, and automated airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction. Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction.
--- ---
@@ -51,7 +51,7 @@ Python toolkit for secure, auditable, and automated airlock agent and policy man
## 📜 License ## 📜 License
**AirlockTools** is licensed under the **GNU Affero General Public License v3.0**. **Loxide** is licensed under the **GNU Affero General Public License v3.0**.
You may copy, distribute, and modify the software under the terms of the AGPL-3.0 license. You may copy, distribute, and modify the software under the terms of the AGPL-3.0 license.
+45 -22
View File
@@ -30,6 +30,7 @@ scheduled_jobs: Dict[str, asyncio.TimerHandle] = {}
# Path to the JSON file for job persistence TODO - pin this to the correct place # Path to the JSON file for job persistence TODO - pin this to the correct place
JOBS_FILE = os.path.join(os.getcwd(), "jobs.json") JOBS_FILE = os.path.join(os.getcwd(), "jobs.json")
def register_function(name: str, func: Callable): def register_function(name: str, func: Callable):
""" """
Register a function so it can be called by name later. Register a function so it can be called by name later.
@@ -38,6 +39,7 @@ def register_function(name: str, func: Callable):
""" """
FUNCTION_MAP[name] = func FUNCTION_MAP[name] = func
def load_jobs() -> List[Dict[str, Any]]: def load_jobs() -> List[Dict[str, Any]]:
""" """
Load jobs from the JSON file, or return [] if none exist. Load jobs from the JSON file, or return [] if none exist.
@@ -47,6 +49,7 @@ def load_jobs() -> List[Dict[str, Any]]:
with open(JOBS_FILE, "r") as f: with open(JOBS_FILE, "r") as f:
return json.load(f) return json.load(f)
def save_jobs(jobs: List[Dict[str, Any]]): def save_jobs(jobs: List[Dict[str, Any]]):
""" """
Save jobs to the JSON file (overwrite). Save jobs to the JSON file (overwrite).
@@ -54,6 +57,7 @@ def save_jobs(jobs: List[Dict[str, Any]]):
with open(JOBS_FILE, "w") as f: with open(JOBS_FILE, "w") as f:
json.dump(jobs, f, indent=4) json.dump(jobs, f, indent=4)
def cancel_job(job_id: str): def cancel_job(job_id: str):
""" """
Cancel a scheduled job by ID and remove it from the registry and persistence. Cancel a scheduled job by ID and remove it from the registry and persistence.
@@ -66,7 +70,15 @@ def cancel_job(job_id: str):
jobs = [j for j in load_jobs() if j.get("id") != job_id] jobs = [j for j in load_jobs() if j.get("id") != job_id]
save_jobs(jobs) save_jobs(jobs)
def run_once_job(job_id: str, func_name: str, delay_seconds: float, args=None, kwargs=None, persist=True):
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). Schedule a job to run once after a delay (in seconds).
""" """
@@ -87,18 +99,25 @@ def run_once_job(job_id: str, func_name: str, delay_seconds: float, args=None, k
if persist: if persist:
jobs = [j for j in load_jobs() if j.get("id") != job_id] jobs = [j for j in load_jobs() if j.get("id") != job_id]
jobs.append({ jobs.append(
"id": job_id, {
"type": "once", "id": job_id,
"delay": delay_seconds, "type": "once",
"function": func_name, "delay": delay_seconds,
"args": args, "function": func_name,
"kwargs": kwargs "args": args,
}) "kwargs": kwargs,
}
)
save_jobs(jobs) save_jobs(jobs)
logger.info(f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds.") 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):
def recurring_job(
job_id: str, func_name: str, interval: float, args=None, kwargs=None, persist=True
):
""" """
Schedule a recurring job. Schedule a recurring job.
""" """
@@ -121,17 +140,20 @@ def recurring_job(job_id: str, func_name: str, interval: float, args=None, kwarg
if persist: if persist:
jobs = [j for j in load_jobs() if j.get("id") != job_id] jobs = [j for j in load_jobs() if j.get("id") != job_id]
jobs.append({ jobs.append(
"id": job_id, {
"type": "recurring", "id": job_id,
"interval": interval, "type": "recurring",
"function": func_name, "interval": interval,
"args": args, "function": func_name,
"kwargs": kwargs "args": args,
}) "kwargs": kwargs,
}
)
save_jobs(jobs) save_jobs(jobs)
logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.") logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.")
def reload_jobs(): def reload_jobs():
""" """
Reload jobs from JSON and reschedule them. Reload jobs from JSON and reschedule them.
@@ -145,7 +167,7 @@ def reload_jobs():
job["delay"], job["delay"],
job.get("args"), job.get("args"),
job.get("kwargs"), job.get("kwargs"),
persist=False persist=False,
) )
elif job["type"] == "recurring": elif job["type"] == "recurring":
recurring_job( recurring_job(
@@ -154,9 +176,10 @@ def reload_jobs():
job["interval"], job["interval"],
job.get("args"), job.get("args"),
job.get("kwargs"), job.get("kwargs"),
persist=False persist=False,
) )
async def start_scheduler(): async def start_scheduler():
""" """
Start the asynchronous scheduler loop. Start the asynchronous scheduler loop.
@@ -189,4 +212,4 @@ async def start_scheduler():
while True: while True:
await asyncio.sleep(3600) # Sleep indefinitely; jobs run via call_later await asyncio.sleep(3600) # Sleep indefinitely; jobs run via call_later
except asyncio.CancelledError: except asyncio.CancelledError:
logger.critical("Scheduler stopped.") logger.critical("Scheduler stopped.")
+1 -1
View File
@@ -17,7 +17,7 @@ dependencies = [
[[package]] [[package]]
name = "airlock_libs" name = "airlock_libs"
version = "1.0.3" version = "2.0.0"
dependencies = [ dependencies = [
"chrono", "chrono",
"indicatif", "indicatif",
+1 -1
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "airlock_libs" name = "airlock_libs"
version = "1.0.3" version = "2.0.0"
edition = "2024" edition = "2024"
[lib] [lib]
+63 -60
View File
@@ -1,6 +1,9 @@
from typing import Dict, List, Optional from typing import Dict, List
def pull_policy_exec_histories(self, type: List[str], checkpoint: str, policy: List[str]) -> str:
"""Retrieve execution history logs.""" def pull_policy_exec_histories(
self, type: List[str], checkpoint: str, policy: List[str]
) -> str:
"""Retrieve execution history logs."""
def api(AirlockAPIWrapper): def api(AirlockAPIWrapper):
""" """
@@ -22,66 +25,66 @@ def api(AirlockAPIWrapper):
""" """
def history_logging( def history_logging(
api, api,
exec_types: str, exec_types: str,
checkpoint_number: str, checkpoint_number: str,
policy_names: str, policy_names: str,
) -> List[Dict[str, Any]]: ) -> List[Dict[str, Any]]:
""" """
Query execution history logs from the Airlock API. Query execution history logs from the Airlock API.
Parameters Parameters
---------- ----------
exec_types : str exec_types : str
A JSON-style string list of execution types to retrieve. A JSON-style string list of execution types to retrieve.
Example: "[3,5,8]" Example: "[3,5,8]"
- 0 = Trusted Execution - 0 = Trusted Execution
- 1 = Blocked Execution - 1 = Blocked Execution
- 2 = Untrusted Execution [Audit] - 2 = Untrusted Execution [Audit]
- 3 = Untrusted Execution [OTP] - 3 = Untrusted Execution [OTP]
- 5 = Trusted Publisher Execution - 5 = Trusted Publisher Execution
- 8 = Trusted Process Execution - 8 = Trusted Process Execution
(etc.) (etc.)
checkpoint_number : str checkpoint_number : str
The checkpoint ID. Used to fetch results after a certain event. The checkpoint ID. Used to fetch results after a certain event.
Example: "601d275487bacb01e3470713" Example: "601d275487bacb01e3470713"
policy_names : str policy_names : str
A comma-separated or JSON-style list of policy group names. A comma-separated or JSON-style list of policy group names.
Example: "Apple Mac" or "["Apple Mac", "Servers London"]" Example: "Apple Mac" or "["Apple Mac", "Servers London"]"
Returns Returns
------- -------
List[Dict[str, Any]] List[Dict[str, Any]]
A list of dictionaries, where each dictionary represents an A list of dictionaries, where each dictionary represents an
execution history record. Each record can include fields like: execution history record. Each record can include fields like:
- checkpoint: str - checkpoint: str
- type: int - type: int
- username: str - username: str
- hostname: str - hostname: str
- filename: str - filename: str
- ppolicy: str - ppolicy: str
- policyname: str - policyname: str
- policyver: str - policyver: str
- commandline: str - commandline: str
- publisher: str - publisher: str
- pprocess: str - pprocess: str
- gprocess: str - gprocess: str
- sha256: str - sha256: str
- datetime: str - datetime: str
- ip: str - ip: str
- localip: str - localip: str
Raises Raises
------ ------
RuntimeError RuntimeError
If the request fails or the response cannot be parsed. If the request fails or the response cannot be parsed.
Example Example
------- -------
>>> histories = await airlock_libs.history_logging("[3,5,8]", "601d275487bacb01e3470713", "Apple Mac") >>> histories = await airlock_libs.history_logging("[3,5,8]", "601d275487bacb01e3470713", "Apple Mac")
>>> print(histories[0]["filename"]) >>> print(histories[0]["filename"])
'chrome.exe' 'chrome.exe'
""" """
... ...
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "maturin"
[project] [project]
name = "airlock_libs" name = "airlock_libs"
version = "1.0.3" version = "2.0.0"
description = "Airlock Digital API Wrapper" description = "Airlock Digital API Wrapper"
readme = "README.md" readme = "README.md"
license = { text = "AGPL-3.0-only" } license = { text = "AGPL-3.0-only" }
+4 -4
View File
@@ -13,7 +13,7 @@ use std::{
env, env,
fmt::Write, fmt::Write,
fs::{self, File}, fs::{self, File},
io::Read, io::{Read, Seek, SeekFrom},
path::PathBuf, path::PathBuf,
str::FromStr, str::FromStr,
}; };
@@ -96,7 +96,9 @@ pub fn pull_policy_exec_histories(
let client = build_client(py, &py_self); let client = build_client(py, &py_self);
let api: Py<PyAny> = py_self; let api: Py<PyAny> = py_self;
let cutoff = Local::now().naive_local() - Duration::days(days); let cutoff = Local::now().naive_local() - Duration::days(days);
let mut f = File::open(&writeable_filepath).unwrap();
loop { loop {
f.seek(SeekFrom::Start(0)).unwrap();
let execution_histories = history_logging( let execution_histories = history_logging(
py, py,
&api, &api,
@@ -110,7 +112,6 @@ pub fn pull_policy_exec_histories(
break; break;
} }
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() { let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() {
let mut f = File::open(&writeable_filepath).unwrap();
let mut contents = String::new(); let mut contents = String::new();
f.read_to_string(&mut contents).unwrap(); f.read_to_string(&mut contents).unwrap();
let existing_data: ApiResponse = let existing_data: ApiResponse =
@@ -177,8 +178,7 @@ pub fn pull_policy_exec_histories(
) )
{ {
let date_diff = Local::now().naive_local().date() - last_date; let date_diff = Local::now().naive_local().date() - last_date;
let percentage_diff = let percentage_diff = (days - date_diff.num_days()) as f64 / days as f64 * 100.0;
((days + 10) - date_diff.num_days()) as f64 / (days + 10) as f64 * 100.0;
progress_bar.set_position(percentage_diff.round() as u64); progress_bar.set_position(percentage_diff.round() as u64);
progress_bar.set_message("Total Percent Complete"); progress_bar.set_message("Total Percent Complete");
} }
+36 -16
View File
@@ -41,7 +41,9 @@ def getLocalApprovals(api: AirlockAPIWrapper):
result = api.otp_find_awaiting() result = api.otp_find_awaiting()
local_approval = pd.DataFrame(result["response"]["otpusage"]) local_approval = pd.DataFrame(result["response"]["otpusage"])
if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"): 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 = pd.read_parquet(
f"{base_dir}\\cache\\newest_local_approval.parquet"
)
previous_run.to_parquet( previous_run.to_parquet(
f"{base_dir}\\cache\\last_local_approval.parquet", index=False f"{base_dir}\\cache\\last_local_approval.parquet", index=False
) )
@@ -66,10 +68,10 @@ def getLocalApprovals(api: AirlockAPIWrapper):
def scheduleAddingLAHashes(api: AirlockAPIWrapper): def scheduleAddingLAHashes(api: AirlockAPIWrapper):
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}") policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]") bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
pups = load_env_json("PUPS", "[]") pups = load_env_json("PUPS", "[]")
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type = int) threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type=int)
try: try:
register_function("add_hash", returnFromLocalApproval) register_function("add_hash", returnFromLocalApproval)
@@ -93,7 +95,9 @@ def scheduleAddingLAHashes(api: AirlockAPIWrapper):
duration_minutes = int(batch_df["duration"].iloc[0]) duration_minutes = int(batch_df["duration"].iloc[0])
start_time = datetime.datetime.now() start_time = datetime.datetime.now()
run_time = start_time + datetime.timedelta(minutes=duration_minutes) run_time = start_time + datetime.timedelta(minutes=duration_minutes)
early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95)) early_time = start_time + datetime.timedelta(
minutes=np.floor(duration_minutes * 0.95)
)
early_timestamp = early_time.timestamp() early_timestamp = early_time.timestamp()
run_timestamp = run_time.timestamp() run_timestamp = run_time.timestamp()
@@ -148,7 +152,13 @@ def scheduleAddingLAHashes(api: AirlockAPIWrapper):
logger.warning(f"Failed to process batch {batchid}: {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 def returnFromLocalApproval(
api,
device_df,
policy_relationship_map,
bad_publisher_list,
pups,
threat_tolerance_constant,
): ):
""" """
# Get unique policy names from device list # Get unique policy names from device list
@@ -166,11 +176,14 @@ def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publish
#TODO finish logic for adding hashes #TODO finish logic for adding hashes
""" """
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}") policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]") bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
pups = load_env_json("PUPS", "[]") pups = load_env_json("PUPS", "[]")
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE") threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE")
print(f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}") print(
f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}"
)
def moveToLocalApproval(api: AirlockAPIWrapper): def moveToLocalApproval(api: AirlockAPIWrapper):
possible_durations = [15, 60, 360, 1440, 10080] possible_durations = [15, 60, 360, 1440, 10080]
@@ -213,10 +226,9 @@ def moveToLocalApproval(api: AirlockAPIWrapper):
def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid): def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid):
purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}" purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
api.otp_generate(agentid, duration_selected, purpose) api.otp_generate(agentid, duration_selected, purpose)
def monitorAuditStatus(api: AirlockAPIWrapper): def monitorAuditStatus(api: AirlockAPIWrapper):
@@ -224,11 +236,13 @@ def monitorAuditStatus(api: AirlockAPIWrapper):
last_agents = [] last_agents = []
if not last_agents: if not last_agents:
last_agents = current_agents last_agents = current_agents
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}") policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
# Reverse map for audit → enforcement # Reverse map for audit → enforcement
reverse_policy_map = {v: k for k, v in policy_relationship_map.items()} 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()) known_transitions = set(policy_relationship_map.items()) | set(
reverse_policy_map.items()
)
# Index last_agents by hostname for quick lookup # Index last_agents by hostname for quick lookup
last_agent_map = {agent.hostname: agent for agent in last_agents} last_agent_map = {agent.hostname: agent for agent in last_agents}
@@ -261,8 +275,8 @@ def monitorAuditStatus(api: AirlockAPIWrapper):
def getNewLocalApprovals(api: AirlockAPIWrapper): def getNewLocalApprovals(api: AirlockAPIWrapper):
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
current_la = getLocalApprovals(api) current_la = getLocalApprovals(api)
# Load old approval list # Load old approval list
@@ -273,15 +287,21 @@ def getNewLocalApprovals(api: AirlockAPIWrapper):
old_la = pd.DataFrame(columns=current_la.columns) old_la = pd.DataFrame(columns=current_la.columns)
# Create composite keys # Create composite keys
current_la["key"] = current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str) 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) old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str)
# Find new entries # Find new entries
new_entries = current_la[~current_la["key"].isin(old_la["key"])] new_entries = current_la[~current_la["key"].isin(old_la["key"])]
# Convert 'granted' to datetime and filter by last 10 minutes # Convert 'granted' to datetime and filter by last 10 minutes
new_entries["granted"] = pd.to_datetime(new_entries["granted"], utc=True, errors="coerce") new_entries["granted"] = pd.to_datetime(
ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10) 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] recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago]
# Save current approvals for next run # Save current approvals for next run
+60 -43
View File
@@ -14,7 +14,6 @@
# along with this program. If not, see <https://www.gnu.org/licenses/>. # along with this program. If not, see <https://www.gnu.org/licenses/>.
from datetime import datetime from datetime import datetime
import logging import logging
import os import os
@@ -33,19 +32,24 @@ logger = logging.getLogger(__name__)
def otp_generate(api: AirlockAPIWrapper): def otp_generate(api: AirlockAPIWrapper):
otp_dict = {} otp_dict = {}
agents = selectAgents(api) agents = selectAgents(api)
print(colorText("Would you like to continue with these devices?","white")) print(colorText("Would you like to continue with these devices?", "white"))
for agent in agents: for agent in agents:
print(agent.hostname) print(agent.hostname)
confirm = Selector.confirm() confirm = Selector.confirm()
if agents and confirm: if agents and confirm:
requester = get_sanitized_input("Who is requesting the OTP: ") requester = get_sanitized_input("Who is requesting the OTP: ")
because = get_sanitized_input("Why/What work are they doing?: ") because = get_sanitized_input("Why/What work are they doing?: ")
purpose = f"Requester: {requester} - for : {because}" purpose = f"Requester: {requester} - for : {because}"
possible_durations = [15, 60, 360, 1440, 10080] possible_durations = [15, 60, 360, 1440, 10080]
print(colorText("Please select a duration in minutes: ", "white")) 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")) 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) duration_selected = Selector.select_int(possible_durations)
if isinstance(duration_selected, list): if isinstance(duration_selected, list):
@@ -60,56 +64,68 @@ def otp_generate(api: AirlockAPIWrapper):
print(colorText("Requested Codes:", "green")) print(colorText("Requested Codes:", "green"))
for key, value in otp_dict.items(): for key, value in otp_dict.items():
print(colorText(f"{key} | {value}","green")) print(colorText(f"{key} | {value}", "green"))
def otp_activities_by_agent(api: AirlockAPIWrapper): def otp_activities_by_agent(api: AirlockAPIWrapper):
activeagents = api.otp_find_active() activeagents = api.otp_find_active()
awaitingagents = api.otp_find_awaiting() awaitingagents = api.otp_find_awaiting()
enforcedagents = api.otp_find_enforced() enforcedagents = api.otp_find_enforced()
revokedagents = api.otp_find_revoked() revokedagents = api.otp_find_revoked()
# Add a 'status' column to each DataFrame # Add a 'status' column to each DataFrame
activeagents['status'] = 'active' activeagents["status"] = "active"
awaitingagents['status'] = 'awaiting' awaitingagents["status"] = "awaiting"
enforcedagents['status'] = 'enforced' enforcedagents["status"] = "enforced"
revokedagents['status'] = 'revoked' revokedagents["status"] = "revoked"
# Combine all into one DataFrame # Combine all into one DataFrame
combined_agents = pd.concat([activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True) combined_agents = pd.concat(
combined_agents = combined_agents.sort_values(by='otpid', ascending=False) [activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True
)
combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
#Optionally, select specific hosts # Optionally, select specific hosts
user_input = get_sanitized_input("\nWould you like to search for a specific device? (y/n): ").strip().lower() user_input = (
if user_input == 'y': get_sanitized_input("\nWould you like to search for a specific device? (y/n): ")
.strip()
.lower()
)
if user_input == "y":
agentnames = [] agentnames = []
agents = selectAgents(api) agents = selectAgents(api)
for agent in agents: for agent in agents:
agentnames.append(agent.hostname) agentnames.append(agent.hostname)
combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)] combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)]
#Present and select rows # Present and select rows
selected_rows = Selector.select_dataframe_with_mode( selected_rows = Selector.select_dataframe_with_mode(
combined_agents, combined_agents,
columns=['otpid', 'hostname', 'status','purpose','granted'], columns=["otpid", "hostname", "status", "purpose", "granted"],
header="OTP Sessions" header="OTP Sessions",
) )
combined_df = pd.DataFrame() combined_df = pd.DataFrame()
for row in selected_rows: for row in selected_rows:
otpid = row['otpid'] otpid = row["otpid"]
hostname = row['hostname'] hostname = row["hostname"]
result = api.otp_get_activities(otpid) result = api.otp_get_activities(otpid)
result['hostname'] = hostname result["hostname"] = hostname
if not result.empty: if not result.empty:
logger.info(f"Activities for {hostname} (otpid: {otpid}):\n{result}") logger.info(f"Activities for {hostname} (otpid: {otpid}):\n{result}")
combined_df = pd.concat([combined_df, result], ignore_index=True) combined_df = pd.concat([combined_df, result], ignore_index=True)
else: else:
logger.info(f"No activities found for {hostname} (otpid: {otpid})") logger.info(f"No activities found for {hostname} (otpid: {otpid})")
user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower() user_input = (
if user_input == 'y': get_sanitized_input(
"\nWould you like to export the results to a CSV file? (y/n): "
)
.strip()
.lower()
)
if user_input == "y":
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = f"otp_activities_{timestamp}.csv" filename = f"otp_activities_{timestamp}.csv"
@@ -128,43 +144,44 @@ def otp_activities_by_agent(api: AirlockAPIWrapper):
logging.debug("User declined to export the DataFrame.") logging.debug("User declined to export the DataFrame.")
def otp_revoke(api: AirlockAPIWrapper): def otp_revoke(api: AirlockAPIWrapper):
activeagents = api.otp_find_active() activeagents = api.otp_find_active()
awaitingagents = api.otp_find_awaiting() awaitingagents = api.otp_find_awaiting()
activeagents['status'] = 'active' activeagents["status"] = "active"
awaitingagents['status'] = 'awaiting' awaitingagents["status"] = "awaiting"
combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True) combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True)
combined_agents = combined_agents.sort_values(by='otpid', ascending=False) combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
# Combine all into one DataFrame # Combine all into one DataFrame
combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True) combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True)
combined_agents = combined_agents.sort_values(by='otpid', ascending=False) combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
#Optionally, select specific hosts # Optionally, select specific hosts
user_input = get_sanitized_input("\nWould you like to search for a specific device? (y/n): ").strip().lower() user_input = (
if user_input == 'y': get_sanitized_input("\nWould you like to search for a specific device? (y/n): ")
.strip()
.lower()
)
if user_input == "y":
agentnames = [] agentnames = []
agents = selectAgents(api) agents = selectAgents(api)
for agent in agents: for agent in agents:
agentnames.append(agent.hostname) agentnames.append(agent.hostname)
combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)] combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)]
#Present and select rows # Present and select rows
selected_rows = Selector.select_dataframe_with_mode( selected_rows = Selector.select_dataframe_with_mode(
combined_agents, combined_agents,
columns=['otpid', 'hostname', 'status','purpose','granted'], columns=["otpid", "hostname", "status", "purpose", "granted"],
header="OTP Sessions" header="OTP Sessions",
) )
for row in selected_rows: for row in selected_rows:
otpid = row['otpid'] otpid = row["otpid"]
hostname = row['hostname'] hostname = row["hostname"]
result = api.otp_revoke(otpid) result = api.otp_revoke(otpid)
logger.info(f"{hostname} (otpid: {otpid}):\n{result}") logger.info(f"{hostname} (otpid: {otpid}):\n{result}")
+387 -144
View File
@@ -44,7 +44,6 @@ logger = logging.getLogger(__name__)
dotenv.load_dotenv() dotenv.load_dotenv()
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]: def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()] policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
@@ -60,9 +59,18 @@ def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
return selected if isinstance(selected, list) else [selected] return selected if isinstance(selected, list) else [selected]
def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) -> List[Allowlist]: def selectAllowlists(
if policy == "all": allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()] api: AirlockAPIWrapper, policy=all, allow_multiple=True
else: allowlists = [Allowlist(**row.to_dict()) for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()] ) -> 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)") logger.debug("Prompting for Allowlist(s)")
print(colorText("Please select allowlist(s)", "white")) print(colorText("Please select allowlist(s)", "white"))
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True) selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
@@ -76,9 +84,7 @@ def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True)
def sortHashes( def sortHashes(
api: AirlockAPIWrapper, api: AirlockAPIWrapper, selected_policies: List[Policy], type=[1, 2, 6, 7]
selected_policies: List[Policy],
type=[1, 2, 6, 7]
): ):
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
history_days = Selector.select_value( history_days = Selector.select_value(
@@ -86,34 +92,41 @@ def sortHashes(
value_type=int, value_type=int,
valid_range=(1, 150), valid_range=(1, 150),
) )
logger.debug(f"{history_days} day selected for history") logger.debug(f"{history_days} day selected for history")
if history_days is None: if history_days is None:
logging.warning("No history range selected. Aborting.") logging.warning("No history range selected. Aborting.")
return return
policy_executions = ExecutionHistoryRecord.from_policies( policy_executions = ExecutionHistoryRecord.from_policies(
api, selected_policies, type_=type, history_days=history_days api, selected_policies, type_=type, history_days=history_days
) )
logger.debug(f"Executions contains {policy_executions}") logger.debug(f"Executions contains {policy_executions}")
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(api, policy_executions) enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(
categorized_executions = ExecutionHistoryRecord.categorize_executions_by_hash_decision(enriched_executions) api, policy_executions
approved, unapproved, needs_review, unknown = ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions) )
categorized_executions = (
ExecutionHistoryRecord.categorize_executions_by_hash_decision(
enriched_executions
)
)
approved, unapproved, needs_review, unknown = (
ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
)
categories = { categories = {
"needs_review": needs_review, "needs_review": needs_review,
"approved": approved, "approved": approved,
"unapproved": unapproved, "unapproved": unapproved,
"leftover" : unknown "leftover": unknown,
} }
for label, records in categories.items(): for label, records in categories.items():
if not records: if not records:
continue # Skip empty or falsy categories continue # Skip empty or falsy categories
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv" csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html" html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
@@ -122,9 +135,9 @@ def sortHashes(
df = pd.DataFrame([r.__dict__ for r in records]) df = pd.DataFrame([r.__dict__ for r in records])
# Optional: flatten hash_obj if needed # Optional: flatten hash_obj if needed
if not df.empty and 'hash_obj' in df.columns: if not df.empty and "hash_obj" in df.columns:
hash_df = df['hash_obj'].apply(lambda h: h.to_dict() if h else {}) hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {})
df = pd.concat([df.drop(columns=['hash_obj']), hash_df], axis=1) df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1)
# Save to CSV # Save to CSV
df.to_csv(csv_path, index=False) df.to_csv(csv_path, index=False)
@@ -140,9 +153,11 @@ def buildPathsandPublishers(selected_policies: List[Policy], split):
df1 = pd.DataFrame() df1 = pd.DataFrame()
df2 = pd.DataFrame() df2 = pd.DataFrame()
all_approved_hashes = pd.DataFrame() all_approved_hashes = pd.DataFrame()
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv" path1 = (
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
)
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv" path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type= int) path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type=int)
if os.path.exists(path1): if os.path.exists(path1):
df1 = pd.read_csv(path1) df1 = pd.read_csv(path1)
@@ -163,37 +178,48 @@ def buildPathsandPublishers(selected_policies: List[Policy], split):
if "filename" in all_approved_hashes.columns: if "filename" in all_approved_hashes.columns:
all_approved_hashes = all_approved_hashes.sort_values(by="filename") all_approved_hashes = all_approved_hashes.sort_values(by="filename")
else: else:
logger.warning("Warning: 'filename' column not found in concatenated DataFrame.") logger.warning(
"Warning: 'filename' column not found in concatenated DataFrame."
)
if not all_approved_hashes.empty and path_exclusion_constant: if not all_approved_hashes.empty and path_exclusion_constant:
primary_path_exclusions = calculatePath( primary_path_exclusions = calculatePath(
all_approved_hashes, path_exclusion_constant, all_approved_hashes,
path_exclusion_constant,
split, split,
) )
remaining_hashes = all_approved_hashes[ remaining_hashes = all_approved_hashes[
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"]) ~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
] ]
secondary_path_exclusions = calculatePath( secondary_path_exclusions = calculatePath(
remaining_hashes,(path_exclusion_constant - 1), split remaining_hashes, (path_exclusion_constant - 1), split
) )
remaining_hashes = remaining_hashes[ remaining_hashes = remaining_hashes[
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"]) ~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
] ]
dataframes = { dataframes = {
"all_approved_hashes" : all_approved_hashes, "all_approved_hashes": all_approved_hashes,
"primary_Paths": primary_path_exclusions, "primary_Paths": primary_path_exclusions,
"secondary_Paths": secondary_path_exclusions, "secondary_Paths": secondary_path_exclusions,
"hashes_not_approvable_by_path": remaining_hashes "hashes_not_approvable_by_path": remaining_hashes,
} }
logger.debug("Preparing to sort dataframes") logger.debug("Preparing to sort dataframes")
for name, df in dataframes.items(): for name, df in dataframes.items():
logger.debug(f" DataFrame headers: {list(df.columns)}") logger.debug(f" DataFrame headers: {list(df.columns)}")
if "hashes" in name : df.sort_values(by="filename", inplace=True) if "hashes" in name:
else: df.sort_values(by="longestcfp", inplace=True) df.sort_values(by="filename", inplace=True)
else:
df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv", index=False) df.sort_values(by="longestcfp", inplace=True)
formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html")
df.to_csv(
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv",
index=False,
)
formatHTML(
df,
f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html",
)
if not all_approved_hashes.empty: if not all_approved_hashes.empty:
# Drop all not signed, only keep unique values # Drop all not signed, only keep unique values
@@ -201,14 +227,18 @@ def buildPathsandPublishers(selected_policies: List[Policy], split):
all_approved_hashes["publisher"] != "Not Signed" all_approved_hashes["publisher"] != "Not Signed"
].drop_duplicates(subset=["publisher"]) ].drop_duplicates(subset=["publisher"])
# Remove Bad publisher if somehow they made it this far # Remove Bad publisher if somehow they made it this far
pattern = regulator(load_env_json("BAD_PUBLISHERS","[]")) pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
publist = publist[~publist["publisher"].str.contains(pattern, na=False)] publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
publist = publist[["publisher"]] publist = publist[["publisher"]]
publist.sort_values(by="publisher", inplace=True) publist.sort_values(by="publisher", inplace=True)
publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv", index=False) publist.to_csv(
else: f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv",
index=False,
)
else:
logger.debug("Approved Hashes list appears empty") logger.debug("Approved Hashes list appears empty")
def buildPreflights(selected_policies: List[Policy]): def buildPreflights(selected_policies: List[Policy]):
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
@@ -222,8 +252,7 @@ def buildPreflights(selected_policies: List[Policy]):
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv" path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv"
publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv" publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv"
# Read in and combine the two path generations
#Read in and combine the two path generations
if os.path.exists(path1): if os.path.exists(path1):
df1 = pd.read_csv(path1) df1 = pd.read_csv(path1)
else: else:
@@ -239,19 +268,18 @@ def buildPreflights(selected_policies: List[Policy]):
approved_paths = pd.DataFrame() approved_paths = pd.DataFrame()
else: else:
approved_paths = pd.concat([df1, df2], ignore_index=True) approved_paths = pd.concat([df1, df2], ignore_index=True)
approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep ="first")
#We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions. approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first")
# We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions.
if os.path.exists(hash): if os.path.exists(hash):
hashes = pd.read_csv(hash) hashes = pd.read_csv(hash)
approved_hashes = hashes[~hashes['filename'].isin(approved_paths['longestcfp'])] approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])]
approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep ="first") approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first")
else: else:
logger.warning(f"File not found: {hash}") logger.warning(f"File not found: {hash}")
if os.path.exists(publishers): if os.path.exists(publishers):
approved_publishers = pd.read_csv(publishers) approved_publishers = pd.read_csv(publishers)
@@ -259,19 +287,33 @@ def buildPreflights(selected_policies: List[Policy]):
else: else:
logger.warning(f"File not found: {publishers}") logger.warning(f"File not found: {publishers}")
dataframes = {"approved_paths": approved_paths, "approved_hashes": approved_hashes, "approved_publishers": approved_publishers} dataframes = {
"approved_paths": approved_paths,
"approved_hashes": approved_hashes,
"approved_publishers": approved_publishers,
}
for name, df in dataframes.items(): for name, df in dataframes.items():
logger.debug(f" DataFrame headers: {list(df.columns)}") logger.debug(f" DataFrame headers: {list(df.columns)}")
if name == "approved_paths":df.sort_values(by="longestcfp", inplace=True) if name == "approved_paths":
elif name == "approved_hashes":df.sort_values(by="filename", inplace=True) df.sort_values(by="longestcfp", inplace=True)
elif name == "approved_publishers" : df.sort_values(by="publisher", inplace=True) elif name == "approved_hashes":
df.sort_values(by="filename", inplace=True)
df.to_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv", index=False) elif name == "approved_publishers":
formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html") df.sort_values(by="publisher", inplace=True)
df.to_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv",
index=False,
)
formatHTML(
df,
f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html",
)
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type= int) min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type=int)
def clean_split(path): def clean_split(path):
if not isinstance(path, (str, bytes, os.PathLike)): if not isinstance(path, (str, bytes, os.PathLike)):
@@ -281,7 +323,9 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
return parts return parts
# Diagnostic: log any non-string entries # Diagnostic: log any non-string entries
non_string_entries = df[~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))] non_string_entries = df[
~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))
]
if not non_string_entries.empty: if not non_string_entries.empty:
print(f"[WARNING] Non-string entries found in column '{col}':") print(f"[WARNING] Non-string entries found in column '{col}':")
print(non_string_entries) print(non_string_entries)
@@ -290,10 +334,14 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
split_paths = df[col].apply(clean_split) split_paths = df[col].apply(clean_split)
if min_files_for_path is not None: if min_files_for_path is not None:
df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy() df = df[
split_paths.apply(lambda parts: len(parts) >= min_files_for_path)
].copy()
split_paths = split_paths[df.index] split_paths = split_paths[df.index]
df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant])) df["group_key"] = split_paths.apply(
lambda parts: os.sep.join(parts[:path_exclusion_constant])
)
grouped = df.groupby("group_key") grouped = df.groupby("group_key")
new_rows = [] new_rows = []
@@ -317,7 +365,7 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
for i, parts in enumerate(split_parts): for i, parts in enumerate(split_parts):
filename = parts[-1] filename = parts[-1]
middle = ( middle = (
os.sep.join(parts[len(common_prefix):-1]) os.sep.join(parts[len(common_prefix) : -1])
if len(parts) > len(common_prefix) + 1 if len(parts) > len(common_prefix) + 1
else "" else ""
) )
@@ -330,6 +378,7 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
return pd.DataFrame(new_rows).drop(columns=["group_key"]) return pd.DataFrame(new_rows).drop(columns=["group_key"])
def calculatePath(approved_hashes, path_exclusion_constant, split): def calculatePath(approved_hashes, path_exclusion_constant, split):
if split: if split:
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")] dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
@@ -337,7 +386,7 @@ def calculatePath(approved_hashes, path_exclusion_constant, split):
dfs_by_policy = [approved_hashes] dfs_by_policy = [approved_hashes]
badpathparts = load_env_json("BAD_PATH_PARTS", "[]") badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type = int) min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type=int)
processed_dfs = [] processed_dfs = []
@@ -364,7 +413,9 @@ def calculatePath(approved_hashes, path_exclusion_constant, split):
] ]
unique_sha_counts = ( unique_sha_counts = (
lcp_not_forbidden_review.groupby("longestcfp")["sha256"].nunique().reset_index() lcp_not_forbidden_review.groupby("longestcfp")["sha256"]
.nunique()
.reset_index()
) )
unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"] unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
@@ -380,51 +431,64 @@ def calculatePath(approved_hashes, path_exclusion_constant, split):
return pathExclusions return pathExclusions
def testChange(selected_policies, destination_policy, destination_allowlist): def testChange(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
logger.info("These path exclusions would be added to:") logger.info("These path exclusions would be added to:")
logger.info(destination_policy) logger.info(destination_policy)
pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv") pathexclusions = pd.read_csv(
hashes = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv") f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
)
hashes = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
)
unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates() unique_combinations = pathexclusions[
["longestcfp", "file_extension"]
].drop_duplicates()
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\") drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
processed_paths = [ processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}" (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
for path, ext in unique_combinations.itertuples(index=False, name=None) for path, ext in unique_combinations.itertuples(index=False, name=None)
] ]
for path in processed_paths: for path in processed_paths:
logger.info(path) logger.info(path)
print(colorText("These publishers would added", "yellow")) print(colorText("These publishers would added", "yellow"))
processed_publishers = [] processed_publishers = []
if os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"): if os.path.exists(
publishers = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv") f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
if publishers.empty: ):
print(colorText("The publishers list is empty.", "red")) publishers = pd.read_csv(
else: f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
processed_publishers = ( )
publishers[publishers["publisher"] != "Not Signed"] if publishers.empty:
["publisher"] print(colorText("The publishers list is empty.", "red"))
.drop_duplicates() else:
.tolist() processed_publishers = (
) publishers[publishers["publisher"] != "Not Signed"]["publisher"]
for publisher in processed_publishers: .drop_duplicates()
print(publisher) .tolist()
)
for publisher in processed_publishers:
print(publisher)
print(colorText("These hashes would be added to:", "yellow")) print(colorText("These hashes would be added to:", "yellow"))
print(destination_allowlist) print(destination_allowlist)
processed_hashes = hashes["sha256"].unique().tolist() processed_hashes = hashes["sha256"].unique().tolist()
print_x_wide(processed_hashes, 3) print_x_wide(processed_hashes, 3)
return processed_paths, processed_hashes, processed_publishers return processed_paths, processed_hashes, processed_publishers
def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 into functions
def menu_policy_enforce(
api: AirlockAPIWrapper,
): # TODO Need to clean up 6 and 7 into functions
selected_policies = [] selected_policies = []
destination_policy = [] destination_policy = []
destination_allowlist = [] destination_allowlist = []
@@ -434,16 +498,22 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
while True: while True:
printEnforceChecklist(selected_policies, destination_policy, destination_allowlist) printEnforceChecklist(
selected_policies, destination_policy, destination_allowlist
)
choice = get_sanitized_input("\nEnter your choice: ") choice = get_sanitized_input("\nEnter your choice: ")
if choice == "1": if choice == "1":
clear_screen() clear_screen()
selected_policies = selectPolicies(api,True) selected_policies = selectPolicies(api, True)
elif choice == "2": elif choice == "2":
clear_screen() clear_screen()
print(colorText("Please choose destination_name Policy for Path Exclusions", "white")) print(
colorText(
"Please choose destination_name Policy for Path Exclusions", "white"
)
)
destination_policy = selectPolicies(api, False) destination_policy = selectPolicies(api, False)
@@ -461,35 +531,53 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
elif choice == "4": elif choice == "4":
clear_screen() clear_screen()
if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"): if os.path.exists(
f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"
):
buildPathsandPublishers(selected_policies, False) buildPathsandPublishers(selected_policies, False)
else: else:
print("File not found. Please make sure it's saved correctly and try again.") print(
"File not found. Please make sure it's saved correctly and try again."
)
elif choice == "5": elif choice == "5":
clear_screen() clear_screen()
if os.path.exists(f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv") and os.path.exists( if os.path.exists(
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
) and os.path.exists(
f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv" f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
): ):
buildPreflights(selected_policies) buildPreflights(selected_policies)
else: else:
print("File not found. Please make sure it's saved correctly and try again.") print(
"File not found. Please make sure it's saved correctly and try again."
)
elif choice == "6": elif choice == "6":
clear_screen() clear_screen()
if ( if (
os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv") os.path.exists(
and os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv") f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
)
and os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
)
and destination_policy and destination_policy
and destination_allowlist and destination_allowlist
): ):
processed_paths, processed_hashes, processed_publishers = testChange(selected_policies, destination_policy, destination_allowlist) processed_paths, processed_hashes, processed_publishers = testChange(
selected_policies, destination_policy, destination_allowlist
)
else: else:
# Log which condition(s) failed # Log which condition(s) failed
missing_items = [] missing_items = []
if not os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"): if not os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
):
missing_items.append("approved_paths.csv not found") missing_items.append("approved_paths.csv not found")
if not os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"): if not os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
):
missing_items.append("approved_hashes.csv not found") missing_items.append("approved_hashes.csv not found")
if not destination_policy: if not destination_policy:
missing_items.append("destination_policy is empty or None") missing_items.append("destination_policy is empty or None")
@@ -513,13 +601,19 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
and confirmation.strip() == "I AGREE" and confirmation.strip() == "I AGREE"
): ):
print(colorText("Proceeding with the code...", "yellow")) print(colorText("Proceeding with the code...", "yellow"))
api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes) api.hash_add_to_allowlist(
api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths) destination_allowlist[0].applicationid, processed_hashes
)
api.policy_add_path_exclusions(
destination_policy[0].groupid, processed_paths
)
if processed_publishers: if processed_publishers:
api.policy_add_publishers(destination_policy[0].groupid, processed_publishers) api.policy_add_publishers(
destination_policy[0].groupid, processed_publishers
)
locked() locked()
else: else:
logger.error("Confirmation block failed. Reasons:") logger.error("Confirmation block failed. Reasons:")
if not processed_publishers or processed_hashes or processed_paths: if not processed_publishers or processed_hashes or processed_paths:
@@ -529,30 +623,52 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
if not destination_allowlist: if not destination_allowlist:
logger.error(" - `destination_allowlist` is missing or invalid.") logger.error(" - `destination_allowlist` is missing or invalid.")
if confirmation.strip() != "I AGREE": if confirmation.strip() != "I AGREE":
logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip()) logger.error(
" - User did not confirm with 'I AGREE'. Received: '%s'",
confirmation.strip(),
)
elif choice.upper() == "F": elif choice.upper() == "F":
open_directory(working_dir) open_directory(working_dir)
elif choice.upper() == "B": elif choice.upper() == "B":
break break
else: else:
print(colorText("Invalid choice. Please try again.", "red")) print(colorText("Invalid choice. Please try again.", "red"))
def section_header(title): def section_header(title):
print(colorText("\n --------------------------------------------------------------------", "cyan")) print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(colorText(f" ------------- {title} -------------", "cyan")) print(colorText(f" ------------- {title} -------------", "cyan"))
print(colorText(" --------------------------------------------------------------------", "cyan")) print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist): def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
section_header("🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒") section_header("🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒")
print(colorText("\nSequentially follow these steps to prepare a policy for enforcement:", "white")) print(
colorText(
"\nSequentially follow these steps to prepare a policy for enforcement:",
"white",
)
)
# Step 1: Originating Policies # Step 1: Originating Policies
print(colorText("\n1. Choose which policy or policies to gather execution info from", "cyan")) print(
colorText(
"\n1. Choose which policy or policies to gather execution info from", "cyan"
)
)
if not selected_policies: if not selected_policies:
print(colorText(" [✗] No policies have been chosen", "red")) print(colorText(" [✗] No policies have been chosen", "red"))
else: else:
@@ -561,67 +677,194 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print(colorText(f" [✓] {policy.name}", "green")) print(colorText(f" [✓] {policy.name}", "green"))
# Step 2: Destination Policy and Allowlist # Step 2: Destination Policy and Allowlist
print(colorText("2. Choose the destination policy and associated allowlist", "cyan")) print(
colorText("2. Choose the destination policy and associated allowlist", "cyan")
)
if destination_policy: if destination_policy:
print(colorText(f" [✓] {destination_policy[0].name} has been selected as the destination policy", "green")) print(
colorText(
f" [✓] {destination_policy[0].name} has been selected as the destination policy",
"green",
)
)
else: else:
print(colorText(" [✗] No destination policy has been chosen", "red")) print(colorText(" [✗] No destination policy has been chosen", "red"))
if destination_allowlist: if destination_allowlist:
print(colorText(f" [✓] {destination_allowlist[0].name} has been selected as allowlist", "green")) print(
colorText(
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
"green",
)
)
else: else:
print(colorText(" [✗] No allowlist has been chosen", "red")) print(colorText(" [✗] No allowlist has been chosen", "red"))
# Step 3: Data Preparation # Step 3: Data Preparation
print(colorText(f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review", "cyan")) print(
colorText(
f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review",
"cyan",
)
)
if selected_policies: if selected_policies:
policy_id = selected_policies[0].name policy_id = selected_policies[0].name
review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv" review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv"
print(colorText(" [✓] Data has been fetched" if os.path.exists(review_path) else " [✗] Data has not been fetched", "green" if os.path.exists(review_path) else "red")) print(
colorText(
(
" [✓] Data has been fetched"
if os.path.exists(review_path)
else " [✗] Data has not been fetched"
),
"green" if os.path.exists(review_path) else "red",
)
)
else: else:
print(colorText(" [✗] No policies selected, cannot check data fetch status", "red")) print(
colorText(
" [✗] No policies selected, cannot check data fetch status", "red"
)
)
# Step 4: Manual Review # Step 4: Manual Review
print(colorText("4. Manually review the files:", "cyan")) print(colorText("4. Manually review the files:", "cyan"))
print(colorText(" Remove the rows containing hashes you do not approve of", "cyan")) print(
print(colorText(f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.", "cyan")) colorText(
print(colorText(" This will start the process to generate possible filepath approvals", "cyan")) " Remove the rows containing hashes you do not approve of", "cyan"
)
)
print(
colorText(
f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.",
"cyan",
)
)
print(
colorText(
" This will start the process to generate possible filepath approvals",
"cyan",
)
)
if selected_policies: if selected_policies:
policy_id = selected_policies[0].name policy_id = selected_policies[0].name
approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv" approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv"
second_review_path = f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv" second_review_path = (
print(colorText(" [✓] Reviewed hashes have been loaded" if os.path.exists(approved_path) else " [✗] Reviewed hashes have not been loaded", "green" if os.path.exists(approved_path) else "red")) f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv"
print(colorText(" [✓] Path review list created" if os.path.exists(second_review_path) else " [✗] Path review list has not been created", "green" if os.path.exists(second_review_path) else "red")) )
print(
colorText(
(
" [✓] Reviewed hashes have been loaded"
if os.path.exists(approved_path)
else " [✗] Reviewed hashes have not been loaded"
),
"green" if os.path.exists(approved_path) else "red",
)
)
print(
colorText(
(
" [✓] Path review list created"
if os.path.exists(second_review_path)
else " [✗] Path review list has not been created"
),
"green" if os.path.exists(second_review_path) else "red",
)
)
else: else:
print(colorText(" [✗] No policies selected, cannot check reviewed hashes or path list", "red")) print(
colorText(
" [✗] No policies selected, cannot check reviewed hashes or path list",
"red",
)
)
# Step 5: Path Review # Step 5: Path Review
print(colorText(f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\", "cyan")) print(
print(colorText(" Remove the rows containing path exclusions or publishers you do not approve of.", "cyan")) colorText(
print(colorText(f" When complete, save the files to {working_dir}\\data\\Approved", "cyan")) f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\",
print(colorText(" Choose this option when done to build your preflights", "cyan")) "cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions or publishers you do not approve of.",
"cyan",
)
)
print(
colorText(
f" When complete, save the files to {working_dir}\\data\\Approved",
"cyan",
)
)
print(
colorText(" Choose this option when done to build your preflights", "cyan")
)
if selected_policies: if selected_policies:
policy_id = selected_policies[0].name policy_id = selected_policies[0].name
reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv" reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv"
preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv" preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv"
preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv" preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv"
print(colorText(" [✓] Reviewed path list detected" if os.path.exists(reviewed_path) else " [✗] Path review list has not been detected", "green" if os.path.exists(reviewed_path) else "red")) print(
preflight_ready = os.path.exists(preflight_paths) and os.path.exists(preflight_hashes) colorText(
print(colorText(" [✓] Preflight Path Exclusion List has been generated" if preflight_ready else " [✗] Preflight Path Exclusion List has not been generated", "green" if preflight_ready else "red")) (
" [✓] Reviewed path list detected"
if os.path.exists(reviewed_path)
else " [✗] Path review list has not been detected"
),
"green" if os.path.exists(reviewed_path) else "red",
)
)
preflight_ready = os.path.exists(preflight_paths) and os.path.exists(
preflight_hashes
)
print(
colorText(
(
" [✓] Preflight Path Exclusion List has been generated"
if preflight_ready
else " [✗] Preflight Path Exclusion List has not been generated"
),
"green" if preflight_ready else "red",
)
)
else: else:
print(colorText(" [✗] No policies selected, cannot check preflight status", "red")) print(
colorText(
" [✗] No policies selected, cannot check preflight status", "red"
)
)
# Final Steps # Final Steps
print(colorText("6. Test ------------------------------------------------------", "cyan")) print(
print(colorText(" Prints to console the changes that would be made, must be done to proceed. ", "cyan")) colorText(
"6. Test ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Prints to console the changes that would be made, must be done to proceed. ",
"cyan",
)
)
print(colorText("7. Liftoff ------------------------------------------------------", "cyan")) print(
print(colorText(" Apply path exclusions and approved publishers to selected policy", "cyan")) colorText(
"7. Liftoff ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Apply path exclusions and approved publishers to selected policy",
"cyan",
)
)
print(colorText(" Apply approved hashes to allowlist", "cyan")) print(colorText(" Apply approved hashes to allowlist", "cyan"))
# Utility Options # Utility Options
print(colorText("F. 📂 - Open Working Directory", "cyan")) print(colorText("F. 📂 - Open Working Directory", "cyan"))
print(colorText("B. 🔚 - Back", "cyan")) print(colorText("B. 🔚 - Back", "cyan"))
+19 -14
View File
@@ -51,20 +51,22 @@ def findQuietAgents(api: AirlockAPIWrapper):
valid_range=(1, 150), valid_range=(1, 150),
) )
confirm = Selector.confirm(f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N : ") confirm = Selector.confirm(
f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N : "
)
# Get execution history as a DataFrame # Get execution history as a DataFrame
if confirm: if confirm:
policy_exec_history = getPolicyInfo( policy_exec_history = getPolicyInfo(
api, selected_policy[0], [1, 2, 6, 7], history_days api, selected_policy[0], [1, 2, 6, 7], history_days
) )
if policy_exec_history.empty: if policy_exec_history.empty:
logging.info("No execution history found for the selected policy and time range.") logging.info(
"No execution history found for the selected policy and time range."
)
get_sanitized_input("Press enter to continue") get_sanitized_input("Press enter to continue")
return return
# Convert 'datetime' column to timezone-aware datetime objects # Convert 'datetime' column to timezone-aware datetime objects
policy_exec_history["datetime"] = pd.to_datetime( policy_exec_history["datetime"] = pd.to_datetime(
policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
@@ -82,12 +84,14 @@ def findQuietAgents(api: AirlockAPIWrapper):
hostname_counts = policy_exec_history["hostname"].value_counts() hostname_counts = policy_exec_history["hostname"].value_counts()
# Map execution counts to agents # Map execution counts to agents
agents["execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int) agents["execution_count"] = (
agents["hostname"].map(hostname_counts).fillna(0).astype(int)
)
# Find most recent execution per hostname # Find most recent execution per hostname
most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates( most_recent_exec = policy_exec_history.sort_values(
subset="hostname", keep="first" by="days_ago"
) ).drop_duplicates(subset="hostname", keep="first")
# Map most recent execution age to agents # Map most recent execution age to agents
agents["days_since"] = agents["hostname"].map( agents["days_since"] = agents["hostname"].map(
@@ -101,7 +105,9 @@ def findQuietAgents(api: AirlockAPIWrapper):
) )
# Sort agents by execution count and hostname # Sort agents by execution count and hostname
agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True]) agents = agents.sort_values(
by=["execution_count", "hostname"], ascending=[True, True]
)
# Save to CSV # Save to CSV
filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv" filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
@@ -116,8 +122,7 @@ def findQuietAgents(api: AirlockAPIWrapper):
ready_percentage = (ready_agents / total_agents) * 100 ready_percentage = (ready_agents / total_agents) * 100
# Print results # Print results
message = ( message = (
f"Total agents: {total_agents}\n" f"Total agents: {total_agents}\n"
f"Agents marked as 'enforce_ready': {ready_agents}\n" f"Agents marked as 'enforce_ready': {ready_agents}\n"
@@ -125,5 +130,5 @@ def findQuietAgents(api: AirlockAPIWrapper):
f"Percentage ready for enforcement: {ready_percentage:.2f}%" f"Percentage ready for enforcement: {ready_percentage:.2f}%"
) )
logger.debug(message) logger.debug(message)
colorText(message,"green") colorText(message, "green")
get_sanitized_input("Press enter to continue") get_sanitized_input("Press enter to continue")
BIN
View File
Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.6 MiB

After

Width:  |  Height:  |  Size: 1.8 MiB

+14 -18
View File
@@ -21,40 +21,36 @@ from models.policy import Policy
@dataclass @dataclass
class Agent: class Agent:
hostname: str
agentid: str agentid: str
clientversion: str clientversion: str
domain: str domain: str
freespace: int freespace: int
groupid: str # Changed to str to match UUID-style IDs groupid: str
hostname: str
ip: str ip: str
localip: str localip: str
lastcheckin: str lastcheckin: str
os: str os: str
policyversion: str policyversion: str
status: int # raw status code status: int
username: str username: str
groupname: Optional[str] = field(default=None) groupname: Optional[str] = field(default=None)
status_text: Optional[str] = field(default=None) status_text: Optional[str] = field(default=None)
# Class-level status map # Class-level status map
status_map: ClassVar[dict] = { status_map: ClassVar[dict] = {0: "Offline", 1: "Online", 2: "Hidden", 3: "Safemode"}
0: "Offline",
1: "Online",
2: "Hidden",
3: "Safemode"
}
def enrich_with_policies(self, policies: List[Policy]): def enrich_with_policies(self, policies: List[Policy]):
"""Enrich the agent with groupname and human-readable status.""" """Enrich the agent with groupname and human-readable status."""
self.status_text = self.status_map.get(self.status, "Unknown") self.status_text = self.status_map.get(self.status, "Unknown")
for policy in policies: for policy in policies:
if policy.groupid == self.groupid: if policy.groupid == self.groupid:
self.groupname = policy.name self.groupname = policy.name
break break
if not self.groupname: if not self.groupname:
self.groupname = "Unknown" self.groupname = "Unknown"
""" """
from models.agent import Agent from models.agent import Agent
+71 -37
View File
@@ -28,7 +28,6 @@ import pandas as pd
import airlock_libs import airlock_libs
from services.API import AirlockAPIWrapper from services.API import AirlockAPIWrapper
from services.policyhandler import pullPolicyExechistories
from utils.configmanager import get_protected_value, load_env_json from utils.configmanager import get_protected_value, load_env_json
from utils.utils import colorText, regulator from utils.utils import colorText, regulator
@@ -36,11 +35,13 @@ logger = logging.getLogger(__name__)
dotenv.load_dotenv() dotenv.load_dotenv()
@dataclass @dataclass
class Hash: class Hash:
""" """
Hash model representing Hash data Hash model representing Hash data
""" """
sha256: str sha256: str
applications: str applications: str
baselines: str baselines: str
@@ -62,7 +63,7 @@ class Hash:
sha384: str sha384: str
sha512: str sha512: str
at_decision: Optional[str] = None at_decision: Optional[str] = None
def to_dict(self): def to_dict(self):
return asdict(self) return asdict(self)
@@ -100,11 +101,15 @@ class Hash:
for hash_obj in hashes: for hash_obj in hashes:
publisher = hash_obj.publisher or "" publisher = hash_obj.publisher or ""
description = hash_obj.description or "" description = hash_obj.description or ""
reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} reputation = (
hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
)
scannermatch = reputation.get("scannermatch") scannermatch = reputation.get("scannermatch")
logger.debug(f"Evaluating hash: {hash_obj}") logger.debug(f"Evaluating hash: {hash_obj}")
logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}") logger.debug(
f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}"
)
# 1. Unapproved: bad publisher or PUP # 1. Unapproved: bad publisher or PUP
if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): if re.search(bad_publishers_pattern, publisher, re.IGNORECASE):
@@ -128,9 +133,9 @@ class Hash:
# 3. Approved or Unapproved based on threat level # 3. Approved or Unapproved based on threat level
try: try:
score = int(scannermatch) # pyright: ignore[reportArgumentType] score = int(scannermatch) # pyright: ignore[reportArgumentType]
logger.debug(f"Parsed scannermatch score: {score}") logger.debug(f"Parsed scannermatch score: {score}")
if score > threat_tolerance: # pyright: ignore[reportOperatorIssue] if score > threat_tolerance: # pyright: ignore[reportOperatorIssue]
logger.debug("Unapproved: Unsigned file with high threat score.") logger.debug("Unapproved: Unsigned file with high threat score.")
hash_obj.at_decision = "unapproved" hash_obj.at_decision = "unapproved"
unapproved_count += 1 unapproved_count += 1
@@ -139,15 +144,17 @@ class Hash:
hash_obj.at_decision = "approved" hash_obj.at_decision = "approved"
approved_count += 1 approved_count += 1
except (ValueError, TypeError): except (ValueError, TypeError):
logger.debug("Needs Review: Scannermatch score is missing or invalid. — {e}") logger.debug(
"Needs Review: Scannermatch score is missing or invalid. — {e}"
)
hash_obj.at_decision = "needs_review" hash_obj.at_decision = "needs_review"
needs_review_count += 1 needs_review_count += 1
logger.debug(
logger.debug(f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}") f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}"
)
return hashes return hashes
@classmethod @classmethod
def export_to_csv(cls, hash_list, directory_path): def export_to_csv(cls, hash_list, directory_path):
""" """
@@ -204,11 +211,12 @@ class ExecutionHistoryRecord:
localip: Optional[str] = None localip: Optional[str] = None
extid: Optional[str] = None extid: Optional[str] = None
extname: Optional[str] = None extname: Optional[str] = None
exttype: Optional[int] = None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension exttype: Optional[int] = (
None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension
)
extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge
hash_obj: Optional[Hash] = None hash_obj: Optional[Hash] = None
@classmethod @classmethod
def from_dict(cls, data: dict): def from_dict(cls, data: dict):
mandatory_fields = [ mandatory_fields = [
@@ -225,7 +233,9 @@ class ExecutionHistoryRecord:
"datetime", "datetime",
] ]
missing_fields = [ missing_fields = [
field for field in mandatory_fields if field not in data or data[field] is None field
for field in mandatory_fields
if field not in data or data[field] is None
] ]
if missing_fields: if missing_fields:
raise ValueError(f"Missing mandatory fields: {missing_fields}") raise ValueError(f"Missing mandatory fields: {missing_fields}")
@@ -255,7 +265,7 @@ class ExecutionHistoryRecord:
extname=data.get("extname"), extname=data.get("extname"),
exttype=data.get("exttype"), exttype=data.get("exttype"),
extbrowser=data.get("extbrowser"), extbrowser=data.get("extbrowser"),
hash_obj=data.get("hash_obj") hash_obj=data.get("hash_obj"),
) )
@classmethod @classmethod
@@ -264,7 +274,9 @@ class ExecutionHistoryRecord:
) -> List["ExecutionHistoryRecord"]: ) -> List["ExecutionHistoryRecord"]:
executions = [] executions = []
for policy in selected_policies: for policy in selected_policies:
execs = airlock_libs.pull_policy_exec_histories(api, policy.name, str([1,2,6,7]), history_days) execs = airlock_libs.pull_policy_exec_histories(
api, policy.name, str([1, 2, 6, 7]), history_days
)
if execs: if execs:
data = json.loads(execs) data = json.loads(execs)
exechistories = data.get("response", {}).get("exechistories", []) exechistories = data.get("response", {}).get("exechistories", [])
@@ -275,8 +287,12 @@ class ExecutionHistoryRecord:
df = df.drop_duplicates(subset=["sha256", "filename", "hostname"]) df = df.drop_duplicates(subset=["sha256", "filename", "hostname"])
df = df.sort_values(by=["sha256", "filename"]) df = df.sort_values(by=["sha256", "filename"])
executions.extend([cls.from_dict(row.to_dict()) for _, row in df.iterrows()]) executions.extend(
logger.debug(f"Staging of Execution history for policy: {policy.name} is complete") [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( print(
colorText( colorText(
f"Staging of Execution history for policy: {policy.name} is complete", f"Staging of Execution history for policy: {policy.name} is complete",
@@ -285,20 +301,23 @@ class ExecutionHistoryRecord:
) )
return executions return executions
@staticmethod @staticmethod
def enrich_with_hashes( def enrich_with_hashes(
api: AirlockAPIWrapper, api: AirlockAPIWrapper, executions: List["ExecutionHistoryRecord"]
executions: List["ExecutionHistoryRecord"]
) -> List["ExecutionHistoryRecord"]: ) -> List["ExecutionHistoryRecord"]:
""" """
Enriches each ExecutionHistoryRecord with a matching Hash object by querying the API. Enriches each ExecutionHistoryRecord with a matching Hash object by querying the API.
""" """
sha256_list = list({e.sha256.strip().lower() for e in executions if e.sha256}) sha256_list = list({e.sha256.strip().lower() for e in executions if e.sha256})
logger.info(f"Extracted {len(sha256_list)} unique sha256 values from {len(executions)} execution records.") logger.info(
f"Extracted {len(sha256_list)} unique sha256 values from {len(executions)} execution records."
)
if not sha256_list: if not sha256_list:
logger.warning("No sha256 values found in execution records. Skipping enrichment.") logger.warning(
"No sha256 values found in execution records. Skipping enrichment."
)
return executions return executions
logger.debug("Querying hash data from API...") logger.debug("Querying hash data from API...")
@@ -307,8 +326,10 @@ class ExecutionHistoryRecord:
hash_objects = [] hash_objects = []
required_fields = { required_fields = {
f.name for f in dataclasses.fields(Hash) f.name
if f.default == dataclasses.MISSING and f.default_factory == dataclasses.MISSING for f in dataclasses.fields(Hash)
if f.default == dataclasses.MISSING
and f.default_factory == dataclasses.MISSING
} }
for sha256, (_, row) in zip(sha256_list, hash_df.iterrows()): for sha256, (_, row) in zip(sha256_list, hash_df.iterrows()):
@@ -346,11 +367,15 @@ class ExecutionHistoryRecord:
exec_record.hash_obj = hash_obj exec_record.hash_obj = hash_obj
enriched_count += 1 enriched_count += 1
logger.info(f"Enriched {enriched_count} out of {len(executions)} execution records with hash data.") logger.info(
f"Enriched {enriched_count} out of {len(executions)} execution records with hash data."
)
return executions return executions
@staticmethod @staticmethod
def categorize_executions_by_hash_decision(executions: List["ExecutionHistoryRecord"]) -> List["ExecutionHistoryRecord"]: def categorize_executions_by_hash_decision(
executions: List["ExecutionHistoryRecord"],
) -> List["ExecutionHistoryRecord"]:
""" """
Categorizes the hash_obj of each ExecutionHistoryRecord based on publisher, description, and reputation. Categorizes the hash_obj of each ExecutionHistoryRecord based on publisher, description, and reputation.
@@ -374,11 +399,15 @@ class ExecutionHistoryRecord:
publisher = hash_obj.publisher or "" publisher = hash_obj.publisher or ""
description = hash_obj.description or "" description = hash_obj.description or ""
reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} reputation = (
hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
)
scannermatch = reputation.get("scannermatch") scannermatch = reputation.get("scannermatch")
logger.debug(f"Evaluating hash: {hash_obj}") logger.debug(f"Evaluating hash: {hash_obj}")
logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}") logger.debug(
f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}"
)
# 1. Unapproved: bad publisher or PUP # 1. Unapproved: bad publisher or PUP
if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): if re.search(bad_publishers_pattern, publisher, re.IGNORECASE):
@@ -402,7 +431,7 @@ class ExecutionHistoryRecord:
# 3. Approved or Unapproved based on threat level # 3. Approved or Unapproved based on threat level
try: try:
score = int(scannermatch) # pyright: ignore[reportArgumentType] score = int(scannermatch) # pyright: ignore[reportArgumentType]
logger.debug(f"Parsed scannermatch score: {score}") logger.debug(f"Parsed scannermatch score: {score}")
if threat_tolerance is not None and score >= threat_tolerance: if threat_tolerance is not None and score >= threat_tolerance:
logger.debug("Unapproved: Unsigned file with high threat score.") logger.debug("Unapproved: Unsigned file with high threat score.")
@@ -413,7 +442,9 @@ class ExecutionHistoryRecord:
hash_obj.at_decision = "approved" hash_obj.at_decision = "approved"
approved_count += 1 approved_count += 1
except (ValueError, TypeError) as e: except (ValueError, TypeError) as e:
logger.debug(f"Needs Review: Scannermatch score is missing or invalid. — {e}") logger.debug(
f"Needs Review: Scannermatch score is missing or invalid. — {e}"
)
hash_obj.at_decision = "needs_review" hash_obj.at_decision = "needs_review"
needs_review_count += 1 needs_review_count += 1
@@ -423,12 +454,14 @@ class ExecutionHistoryRecord:
) )
return executions return executions
@classmethod @classmethod
def sort_by_hash_decision( def sort_by_hash_decision(cls, executions: List["ExecutionHistoryRecord"]) -> Tuple[
cls, executions: List["ExecutionHistoryRecord"] List["ExecutionHistoryRecord"],
) -> Tuple[List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"]]: List["ExecutionHistoryRecord"],
List["ExecutionHistoryRecord"],
List["ExecutionHistoryRecord"],
]:
""" """
Sorts ExecutionHistoryRecord objects into approved, unapproved, needs_review, and unknown groups Sorts ExecutionHistoryRecord objects into approved, unapproved, needs_review, and unknown groups
based on the value of hash_obj.at_decision. based on the value of hash_obj.at_decision.
@@ -454,7 +487,9 @@ class ExecutionHistoryRecord:
else: else:
unknown.append(record) unknown.append(record)
logger.info(f"[ExecutionHistoryRecord] Sorted {len(sorted_executions)} records by hash_obj.at_decision:") logger.info(
f"[ExecutionHistoryRecord] Sorted {len(sorted_executions)} records by hash_obj.at_decision:"
)
logger.info(f" Approved: {len(approved)}") logger.info(f" Approved: {len(approved)}")
logger.info(f" Unapproved: {len(unapproved)}") logger.info(f" Unapproved: {len(unapproved)}")
logger.info(f" Needs Review: {len(needs_review)}") logger.info(f" Needs Review: {len(needs_review)}")
@@ -463,7 +498,6 @@ class ExecutionHistoryRecord:
return approved, unapproved, needs_review, unknown return approved, unapproved, needs_review, unknown
""" """
executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30) executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30)
+6 -2
View File
@@ -29,7 +29,9 @@ class Policy:
def __repr__(self): def __repr__(self):
# Show all current attributes, including dynamically added ones # Show all current attributes, including dynamically added ones
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) attrs = ", ".join(
f"{key}={repr(value)}" for key, value in self.__dict__.items()
)
return f"<Execution({attrs})>" return f"<Execution({attrs})>"
def to_dict(self): def to_dict(self):
@@ -53,7 +55,9 @@ class Allowlist:
def __repr__(self): def __repr__(self):
# Show all current attributes, including dynamically added ones # Show all current attributes, including dynamically added ones
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) attrs = ", ".join(
f"{key}={repr(value)}" for key, value in self.__dict__.items()
)
return f"<Execution({attrs})>" return f"<Execution({attrs})>"
def to_dict(self): def to_dict(self):
+14 -14
View File
@@ -1,14 +1,14 @@
cryptography==46.0.1 cryptography==46.0.3
keyring==25.6.0 keyring==25.6.0
numpy==2.3.2 numpy==2.3.4
pandas==2.3.1 pandas==2.3.3
python-dotenv==1.1.1 pymongo==4.15.3
pymongo python-dotenv==1.2.1
requests==2.32.5 Requests==2.32.5
schedule==1.2.2 textual==6.5.0
tqdm==4.67.1 tqdm==4.67.1
urllib3==2.5.0 urllib3==2.5.0
bson==0.5.10 pyperclip==1.11.0
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ --extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
airlock_libs==1.0.3 airlock_libs==2.0.0
+41
View File
@@ -0,0 +1,41 @@
from typing import List
from textual.app import ComposeResult
from textual.screen import Screen
from models.agent import Agent
from widgets.multiagentselector import MultiAgentSelector
from widgets.OTP_generate import OTPGenerator
class OTPWorkflowScreen(Screen):
"""Screen that handles the OTP generation workflow."""
def __init__(self, all_agents: List[Agent]):
super().__init__()
self.all_agents = all_agents
self.selected_devices = None
def compose(self) -> ComposeResult:
"""Start with the multi-agent selector."""
yield MultiAgentSelector(self.all_agents)
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents - switch to OTP generator."""
self.selected_devices = message.selected_agents
# Remove the MultiAgentSelector
selector = self.query_one(MultiAgentSelector)
selector.remove()
# Mount the OTPGenerator with the selected Agent objects
# No need to pass API - it will access self.app.api directly
self.mount(OTPGenerator(self.selected_devices))
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
"""Handle OTP generation request - call the actual OTP generation function."""
# This will be handled by the main app, but we can also do it here
# For now, just pass it up to the app level
pass
+38 -34
View File
@@ -23,7 +23,6 @@ import requests
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class AirlockAPIWrapper: class AirlockAPIWrapper:
""" """
A wrapper class for interacting with the Airlock API. A wrapper class for interacting with the Airlock API.
@@ -141,25 +140,25 @@ class AirlockAPIWrapper:
payload = {"status": "0"} payload = {"status": "0"}
result = self._post("/v1/otp/usage", payload) result = self._post("/v1/otp/usage", payload)
return pd.DataFrame(result["response"]["otpusage"]) return pd.DataFrame(result["response"]["otpusage"])
def otp_find_enforced(self) -> pd.DataFrame: def otp_find_enforced(self) -> pd.DataFrame:
"""Find OTPs that are awaiting activation.""" """Find OTPs that are awaiting activation."""
payload = {"status": "2"} payload = {"status": "2"}
result = self._post("/v1/otp/usage", payload) result = self._post("/v1/otp/usage", payload)
return pd.DataFrame(result["response"]["otpusage"]) return pd.DataFrame(result["response"]["otpusage"])
def otp_find_revoked(self) -> pd.DataFrame: def otp_find_revoked(self) -> pd.DataFrame:
"""Find OTPs that are awaiting activation.""" """Find OTPs that are awaiting activation."""
payload = {"status": "3"} payload = {"status": "3"}
result = self._post("/v1/otp/usage", payload) result = self._post("/v1/otp/usage", payload)
return pd.DataFrame(result["response"]["otpusage"]) return pd.DataFrame(result["response"]["otpusage"])
def otp_find_by_agent(self, agentid) -> pd.DataFrame: def otp_find_by_agent(self, agentid) -> pd.DataFrame:
"""Find OTP by agent.""" """Find OTP by agent."""
payload = {"agentid": agentid} payload = {"agentid": agentid}
result = self._post("/v1/otp/usage", payload) result = self._post("/v1/otp/usage", payload)
return pd.DataFrame(result["response"]["otpusage"]) return pd.DataFrame(result["response"]["otpusage"])
def otp_generate(self, agentid: str, duration: int, purpose: str) -> str: def otp_generate(self, agentid: str, duration: int, purpose: str) -> str:
"""Generate a new OTP for an agent.""" """Generate a new OTP for an agent."""
payload = { payload = {
@@ -175,7 +174,7 @@ class AirlockAPIWrapper:
payload = {"otpid": otpid} payload = {"otpid": otpid}
result = self._post("/v1/otp/activities", payload) result = self._post("/v1/otp/activities", payload)
return pd.DataFrame(result["response"]["otpactivities"]) return pd.DataFrame(result["response"]["otpactivities"])
def otp_revoke(self, otpid: str) -> dict: def otp_revoke(self, otpid: str) -> dict:
""" """
Revoke an active OTP. Revoke an active OTP.
@@ -186,18 +185,17 @@ class AirlockAPIWrapper:
""" """
payload = {"otpid": otpid} payload = {"otpid": otpid}
return self._post("/v1/otp/revoke", payload) return self._post("/v1/otp/revoke", payload)
def otp_validate(self, otpcode: str) -> dict:
"""
Validate an OTP code.
Parameters:
- otpcode (str): The OTP code to validate.
Returns:
- dict: JSON response indicating validity.
"""
payload = {"otpcode": otpcode}
return self._post("/v1/otp/validate", payload)
def otp_validate(self, otpcode: str) -> dict:
"""
Validate an OTP code.
Parameters:
- otpcode (str): The OTP code to validate.
Returns:
- dict: JSON response indicating validity.
"""
payload = {"otpcode": otpcode}
return self._post("/v1/otp/validate", payload)
# Policy Management # Policy Management
def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict: def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict:
@@ -225,7 +223,7 @@ class AirlockAPIWrapper:
payload = {"groupid": groupid} payload = {"groupid": groupid}
result = self._post("/v1/group/agents", payload) result = self._post("/v1/group/agents", payload)
return pd.DataFrame(result["response"]["agents"]) return pd.DataFrame(result["response"]["agents"])
def policy_list_allowlists(self, groupid: str) -> pd.DataFrame: def policy_list_allowlists(self, groupid: str) -> pd.DataFrame:
"""List allowlists assigned to a specific policy group.""" """List allowlists assigned to a specific policy group."""
payload = {"groupid": groupid} payload = {"groupid": groupid}
@@ -236,31 +234,37 @@ class AirlockAPIWrapper:
"""Set audit mode for a policy group. 1=Audit, 0=Enforcement""" """Set audit mode for a policy group. 1=Audit, 0=Enforcement"""
payload = {"groupid": groupid, "auditmode": auditmode} payload = {"groupid": groupid, "auditmode": auditmode}
return self._post("/v1/group/settings/auditmode", payload) return self._post("/v1/group/settings/auditmode", payload)
def policy_set_script_custom(self, def policy_set_script_custom(
groupid: str, self,
script_custom: int, groupid: str,
scripts_audit: List[str], script_custom: int,
scripts_disabled: List[str], scripts_audit: List[str],
scripts_respect: List[str], scripts_disabled: List[str],
) -> dict: scripts_respect: List[str],
) -> dict:
"""Set audit mode for a policy group. 1=Audit, 0=Enforcement""" """Set audit mode for a policy group. 1=Audit, 0=Enforcement"""
payload = {"groupid": groupid, payload = {
"script_custom": script_custom, "groupid": groupid,
"scripts_audit": scripts_audit, "script_custom": script_custom,
"scripts_disabled": scripts_disabled, "scripts_audit": scripts_audit,
"scripts_respect": scripts_respect "scripts_disabled": scripts_disabled,
} "scripts_respect": scripts_respect,
}
return self._post("/v1/group/settings/script_custom", payload) return self._post("/v1/group/settings/script_custom", payload)
# Execution History # Execution History
def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> str: def history_logging(
self, type: List[str], checkpoint: str, policy: List[str]
) -> str:
"""Retrieve execution history logs.""" """Retrieve execution history logs."""
payload = {"type": type, "checkpoint": checkpoint, "policy": policy} payload = {"type": type, "checkpoint": checkpoint, "policy": policy}
result = self._post("/v1/logging/exechistories", payload) result = self._post("/v1/logging/exechistories", payload)
return result["response"]["exechistories"] return result["response"]["exechistories"]
def history_execution(self, today: str, date_selected: str, agent_name: str) -> List[Dict]: def history_execution(
self, today: str, date_selected: str, agent_name: str
) -> List[Dict]:
""" """
Retrieve execution history logs. Retrieve execution history logs.
+93 -28
View File
@@ -47,7 +47,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
print(colorText("No agents selected or invalid history range.", "red")) print(colorText("No agents selected or invalid history range.", "red"))
return return
historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d") historical_date = (datetime.now() - timedelta(days=history_days)).strftime(
"%Y-%m-%d"
)
today = datetime.now().strftime("%Y-%m-%d") today = datetime.now().strftime("%Y-%m-%d")
all_history = [] all_history = []
@@ -56,7 +58,11 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
try: try:
exechistory = api.history_execution(today, historical_date, agent.hostname) exechistory = api.history_execution(today, historical_date, agent.hostname)
except Exception as e: except Exception as e:
print(colorText(f"❌ Error retrieving history for {agent.hostname}: {e}", "red")) print(
colorText(
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
)
)
continue continue
if isinstance(exechistory, list): if isinstance(exechistory, list):
@@ -76,7 +82,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
print(colorText(f"{key}: {value}", "green")) print(colorText(f"{key}: {value}", "green"))
print("\n") print("\n")
else: else:
print(colorText(f"No execution history found for {agent.hostname}.", "yellow")) print(
colorText(f"No execution history found for {agent.hostname}.", "yellow")
)
if outputjson: if outputjson:
print(json.dumps(all_history, indent=2)) print(json.dumps(all_history, indent=2))
@@ -92,6 +100,7 @@ def findAllAgents(api):
return agents return agents
def findAgents(api, return_dataframe): def findAgents(api, return_dataframe):
agents = selectAgents(api) agents = selectAgents(api)
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
@@ -113,8 +122,14 @@ def findAgents(api, return_dataframe):
print(agent_df) print(agent_df)
logging.debug("Displayed DataFrame to console.") 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() user_input = (
if user_input == 'y': 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") timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = f"agentsearch_{timestamp}.csv" filename = f"agentsearch_{timestamp}.csv"
file_path = os.path.join(str(working_dir), filename) file_path = os.path.join(str(working_dir), filename)
@@ -131,17 +146,27 @@ def findAgents(api, return_dataframe):
else: else:
logging.debug("User declined to export the DataFrame.") logging.debug("User declined to export the DataFrame.")
def collect_device_names() -> List[str]: def collect_device_names() -> List[str]:
print(colorText("🔍 Device Search", "cyan")) print(colorText("🔍 Device Search", "cyan"))
print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan")) print(
print(colorText("When you're done, press Enter twice (Three times if you have a single device).\n", "cyan")) colorText(
"Enter the device hostnames you'd like to search for, one per line.", "cyan"
)
)
print(
colorText(
"When you're done, press Enter twice (Three times if you have a single device).\n",
"cyan",
)
)
print(colorText("Example:", "cyan")) print(colorText("Example:", "cyan"))
print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan")) print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan"))
print(colorText("Paste or type your device names below:", "white")) print(colorText("Paste or type your device names below:", "white"))
device_input_lines = [] device_input_lines = []
empty_line_count = 0 empty_line_count = 0
valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$') valid_line_pattern = re.compile(r"^[a-zA-Z0-9_\- ]+$")
while True: while True:
line = get_sanitized_input("") line = get_sanitized_input("")
@@ -158,7 +183,12 @@ def collect_device_names() -> List[str]:
if valid_line_pattern.match(stripped_line): if valid_line_pattern.match(stripped_line):
device_input_lines.append(stripped_line) device_input_lines.append(stripped_line)
else: else:
print(colorText(f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow")) print(
colorText(
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
"yellow",
)
)
return [name for name in device_input_lines if name] return [name for name in device_input_lines if name]
@@ -168,10 +198,13 @@ def choose_match_type() -> bool:
return get_sanitized_input("").strip().lower() in ["y", "yes"] return get_sanitized_input("").strip().lower() in ["y", "yes"]
def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool) -> List['Agent']: def match_agents(
device_names: List[str], agents: List["Agent"], use_exact: bool
) -> List["Agent"]:
if use_exact: if use_exact:
return [ return [
agent for agent in agents agent
for agent in agents
if agent.hostname.lower() in [name.lower() for name in device_names] if agent.hostname.lower() in [name.lower() for name in device_names]
] ]
else: else:
@@ -180,23 +213,38 @@ def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool
return [agent for agent in agents if regex.search(agent.hostname)] return [agent for agent in agents if regex.search(agent.hostname)]
def show_unmatched(device_names: List[str], matched_agents: List['Agent'], use_exact: bool): def show_unmatched(
device_names: List[str], matched_agents: List["Agent"], use_exact: bool
):
if use_exact: if use_exact:
unmatched = [name for name in device_names if not any(agent.hostname.lower() == name.lower() for agent in matched_agents)] unmatched = [
name
for name in device_names
if not any(
agent.hostname.lower() == name.lower() for agent in matched_agents
)
]
else: else:
unmatched = [name for name in device_names if not any(re.search(re.escape(name), agent.hostname, re.IGNORECASE) for agent in matched_agents)] unmatched = [
name
for name in device_names
if not any(
re.search(re.escape(name), agent.hostname, re.IGNORECASE)
for agent in matched_agents
)
]
if unmatched: if unmatched:
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}") logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow")) print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
def enrich_agents(agents: List['Agent'], policies: List['Policy']): def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
for agent in agents: for agent in agents:
agent.enrich_with_policies(policies) agent.enrich_with_policies(policies)
def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']: def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
device_names = collect_device_names() device_names = collect_device_names()
if not device_names: if not device_names:
logger.debug("No device names entered") logger.debug("No device names entered")
@@ -227,11 +275,11 @@ def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']:
if idx < len(matched_agents): if idx < len(matched_agents):
line += f"{matched_agents[idx].hostname:<30}" line += f"{matched_agents[idx].hostname:<30}"
logger.info(line) logger.info(line)
matched_agents = Selector.select_with_mode( matched_agents = Selector.select_with_mode(
matched_agents, matched_agents,
label_func=lambda agent: agent.hostname, label_func=lambda agent: agent.hostname,
header="Matched Devices:" header="Matched Devices:",
) )
if not matched_agents: if not matched_agents:
@@ -263,11 +311,17 @@ def moveAgentToRelatedPolicy(
if agent.groupid in policy_relationship_map: if agent.groupid in policy_relationship_map:
target_policy = policy_relationship_map[agent.groupid] target_policy = policy_relationship_map[agent.groupid]
elif agent.groupid in policy_relationship_map.values(): elif agent.groupid in policy_relationship_map.values():
logger.debug(f"Agent {agent.hostname} is already in an audit group. No action needed.") logger.debug(
print(f"Agent {agent.hostname} is already in an audit group. No action needed.") 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 return
else: else:
logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.") logger.warning(
f"Error: No corresponding audit policy found for groupid: {agent.groupid}."
)
return return
elif mode == "enforcement": elif mode == "enforcement":
@@ -275,10 +329,14 @@ def moveAgentToRelatedPolicy(
if agent.groupid in inverse_map: if agent.groupid in inverse_map:
target_policy = inverse_map[agent.groupid] target_policy = inverse_map[agent.groupid]
elif agent.groupid in inverse_map.values(): elif agent.groupid in inverse_map.values():
logger.info(f"Agent {agent.hostname} is already in an enforcement group. No action needed.") logger.info(
f"Agent {agent.hostname} is already in an enforcement group. No action needed."
)
return return
else: else:
logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.") logger.warning(
f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}."
)
return return
else: else:
@@ -299,25 +357,32 @@ def toggleEnforcement(api: AirlockAPIWrapper):
devices = selectAgents(api) devices = selectAgents(api)
for device in devices: for device in devices:
print(device.hostname) print(device.hostname)
confirm = Selector.confirm("Would you like to continue with these devices? Y/N: ") confirm = Selector.confirm(
"Would you like to continue with these devices? Y/N: "
)
if direction and devices and confirm: if direction and devices and confirm:
for device in devices: for device in devices:
result = moveAgentToRelatedPolicy(api,device, str(direction).lower()) result = moveAgentToRelatedPolicy(api, device, str(direction).lower())
logger.info(f"{device.hostname}: result: {result}") logger.info(f"{device.hostname}: result: {result}")
get_sanitized_input("Press enter to continue") get_sanitized_input("Press enter to continue")
def moveAgents(api: AirlockAPIWrapper): def moveAgents(api: AirlockAPIWrapper):
devices = selectAgents(api) devices = selectAgents(api)
for device in devices: for device in devices:
print(device.hostname) print(device.hostname)
confirm_devices = Selector.confirm("Would you like to continue with these devices? Y/N: ") confirm_devices = Selector.confirm(
"Would you like to continue with these devices? Y/N: "
)
if devices and confirm_devices: if devices and confirm_devices:
policies = selectPolicies(api, False) policies = selectPolicies(api, False)
confirm_move = Selector.confirm(f"Would you like to move these devices to {policies[0].name}?") confirm_move = Selector.confirm(
f"Would you like to move these devices to {policies[0].name}?"
)
if confirm_move: if confirm_move:
for device in devices: for device in devices:
result = api.agent_move(device.agentid, policies[0].groupid) result = api.agent_move(device.agentid, policies[0].groupid)
logger.info(f"{device.hostname}: result: {result}") logger.info(f"{device.hostname}: result: {result}")
else: else:
logger.info("Exiting without change") logger.info("Exiting without change")
get_sanitized_input("Press enter to continue") get_sanitized_input("Press enter to continue")
+13 -8
View File
@@ -34,7 +34,6 @@ from utils.utils import areYouSure, colorText, get_sanitized_input
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def pullPolicyExechistories( def pullPolicyExechistories(
api: AirlockAPIWrapper, api: AirlockAPIWrapper,
policy: Policy, policy: Policy,
@@ -72,7 +71,7 @@ def pullPolicyExechistories(
) as pbar: ) as pbar:
while True: while True:
histories = api.history_logging( histories = api.history_logging(
type=type, checkpoint=checkpoint, policy= [policy.name] type=type, checkpoint=checkpoint, policy=[policy.name]
) )
# Ensure histories is a list of dictionaries # Ensure histories is a list of dictionaries
@@ -98,13 +97,17 @@ def pullPolicyExechistories(
# Update checkpoint on last item # Update checkpoint on last item
if index == len(histories) - 1: if index == len(histories) - 1:
checkpoint = history_item["checkpoint"] # pyright: ignore[reportArgumentType] checkpoint = history_item[
"checkpoint"
] # pyright: ignore[reportArgumentType]
filebar.desc = f"Checkpoint Progress: {checkpoint}" filebar.desc = f"Checkpoint Progress: {checkpoint}"
break break
try: try:
history_date = datetime.datetime.strptime( history_date = datetime.datetime.strptime(
history_item["datetime"].replace(" +0000 UTC", ""), # pyright: ignore[reportArgumentType] history_item["datetime"].replace(
" +0000 UTC", ""
), # pyright: ignore[reportArgumentType]
"%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%SZ",
).date() ).date()
except ValueError: except ValueError:
@@ -177,8 +180,10 @@ def pullPolicyExechistories(
def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days): def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
import airlock_libs
executionhist_policy = pd.DataFrame() executionhist_policy = pd.DataFrame()
exehist = pullPolicyExechistories(api, policy, type, days, True) exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days)
if exehist is not None: if exehist is not None:
data = json.loads(exehist) data = json.loads(exehist)
executionhist_policy = pd.DataFrame(data["response"]["exechistories"]) executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
@@ -203,7 +208,7 @@ def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
executionhist_policy = executionhist_policy.sort_values( executionhist_policy = executionhist_policy.sort_values(
by=["sha256", "filename"] by=["sha256", "filename"]
) )
logger.debug( f"Staging of Execution history for policy: {policy} is complete") logger.debug(f"Staging of Execution history for policy: {policy} is complete")
print( print(
colorText( colorText(
f"Staging of Execution history for policy: {policy} is complete", f"Staging of Execution history for policy: {policy} is complete",
@@ -235,10 +240,10 @@ def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
for enforcement_policy, audit_policy in policy_relationship_map.items(): for enforcement_policy, audit_policy in policy_relationship_map.items():
api.policy_clone(enforcement_policy, audit_policy) api.policy_clone(enforcement_policy, audit_policy)
api.policy_set_auditmode(audit_policy, "1") api.policy_set_auditmode(audit_policy, "1")
def confirmUpdateAfromE(api: AirlockAPIWrapper): def confirmUpdateAfromE(api: AirlockAPIWrapper):
areYouSure() areYouSure()
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
if confirmation.strip() == "I AGREE": if confirmation.strip() == "I AGREE":
updateAuditPoliciesFromEnforcementPolices(api) updateAuditPoliciesFromEnforcementPolices(api)
+25 -12
View File
@@ -28,8 +28,8 @@ import keyring
# Constants # Constants
KDF_ITERATIONS = 200_000 KDF_ITERATIONS = 200_000
SALT_SIZE = 16 # 128-bit Salt SALT_SIZE = 16 # 128-bit Salt
NONCE_SIZE = 12 # AES-GCM NONCE_SIZE = 12 # AES-GCM
KEY_SIZE = 32 # AES-256 KEY_SIZE = 32 # AES-256
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -49,9 +49,11 @@ def configure_keyring_backend():
system = platform.system() system = platform.system()
if system == "Windows": if system == "Windows":
import keyring.backends.Windows import keyring.backends.Windows
keyring.set_keyring(keyring.backends.Windows.WinVaultKeyring()) keyring.set_keyring(keyring.backends.Windows.WinVaultKeyring())
elif system == "Linux": elif system == "Linux":
import keyring.backends.kwallet import keyring.backends.kwallet
keyring.set_keyring(keyring.backends.kwallet.DBusKeyring()) keyring.set_keyring(keyring.backends.kwallet.DBusKeyring())
else: else:
raise EnvironmentError(f"Unsupported OS: {system}") raise EnvironmentError(f"Unsupported OS: {system}")
@@ -68,8 +70,9 @@ def store_api_key(service: str, username: str, api_key: str, password: str):
b64 = base64.b64encode(blob).decode() b64 = base64.b64encode(blob).decode()
keyring.set_password(service, username, b64) keyring.set_password(service, username, b64)
logger.debug(
logger.debug(f"API key for service '{service}' and user '{username}' stored successfully.") f"API key for service '{service}' and user '{username}' stored successfully."
)
print("\n✅ API key stored securely.") print("\n✅ API key stored securely.")
print("The program will now exit. Press Enter to continue...") print("The program will now exit. Press Enter to continue...")
@@ -90,8 +93,8 @@ def retrieve_api_key(service: str, username: str, password: str) -> str:
raise ValueError("No stored secret for this service/username.") raise ValueError("No stored secret for this service/username.")
blob = base64.b64decode(b64) blob = base64.b64decode(b64)
salt = blob[:SALT_SIZE] salt = blob[:SALT_SIZE]
nonce = blob[SALT_SIZE:SALT_SIZE + NONCE_SIZE] nonce = blob[SALT_SIZE : SALT_SIZE + NONCE_SIZE]
ct = blob[SALT_SIZE + NONCE_SIZE:] ct = blob[SALT_SIZE + NONCE_SIZE :]
key = _derive_key(password.encode(), salt) key = _derive_key(password.encode(), salt)
aesgcm = AESGCM(key) aesgcm = AESGCM(key)
pt = aesgcm.decrypt(nonce, ct, associated_data=None) pt = aesgcm.decrypt(nonce, ct, associated_data=None)
@@ -124,7 +127,9 @@ def getAPI(USERNAME, SERVICE_NAME):
if api_key_exists(SERVICE_NAME, USERNAME): if api_key_exists(SERVICE_NAME, USERNAME):
for attempt in range(1, 4): for attempt in range(1, 4):
password = getpass(f"Attempt {attempt}/3 - Enter password to unlock your API key: ") password = getpass(
f"Attempt {attempt}/3 - Enter password to unlock your API key: "
)
try: try:
apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password) apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password)
logging.debug("API key successfully retrieved.") logging.debug("API key successfully retrieved.")
@@ -134,9 +139,15 @@ def getAPI(USERNAME, SERVICE_NAME):
logging.error("Failed to retrieve API key after 3 incorrect attempts.") logging.error("Failed to retrieve API key after 3 incorrect attempts.")
raise ValueError("Failed to retrieve API key after 3 incorrect attempts.") raise ValueError("Failed to retrieve API key after 3 incorrect attempts.")
else: else:
logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.") logging.warning(
api_key = getpass(f"No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip() f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'."
print("Please exit and relaunch program after saving your credential to avoid errors") )
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: while True:
password = getpass("Create a password to encrypt your API key: ") password = getpass("Create a password to encrypt your API key: ")
@@ -155,7 +166,9 @@ def getAPI(USERNAME, SERVICE_NAME):
logging.error(f"Failed to store API key: {e}") logging.error(f"Failed to store API key: {e}")
break break
else: else:
logging.warning("Password does not meet complexity requirements. Try again.") logging.warning(
"Password does not meet complexity requirements. Try again."
)
class APIKeyManager: class APIKeyManager:
@@ -169,4 +182,4 @@ class APIKeyManager:
def get(cls) -> str: def get(cls) -> str:
if cls._api_key is None: if cls._api_key is None:
raise ValueError("API key not loaded. Call APIKeyManager.load() first.") raise ValueError("API key not loaded. Call APIKeyManager.load() first.")
return cls._api_key return cls._api_key
+24 -15
View File
@@ -29,14 +29,15 @@ PROTECTED_KEYS = [
"PATH_EXCLUSION_CONST", "PATH_EXCLUSION_CONST",
"MIN_FILES_FOR_PATH", "MIN_FILES_FOR_PATH",
"VT_THREAT_TOLERANCE", "VT_THREAT_TOLERANCE",
"POLICY_MAP_ENF_AUD" "POLICY_MAP_ENF_AUD",
] ]
_protected_config = {} _protected_config = {}
def get_system_config_path() -> Path: def get_system_config_path() -> Path:
# Check inside bundled EXE directory first # Check inside bundled EXE directory first
bundled_dir = Path(getattr(sys, '_MEIPASS', '')) bundled_dir = Path(getattr(sys, "_MEIPASS", ""))
bundled_path = bundled_dir / "system_config.json" bundled_path = bundled_dir / "system_config.json"
if bundled_path.exists(): if bundled_path.exists():
return bundled_path return bundled_path
@@ -44,6 +45,7 @@ def get_system_config_path() -> Path:
# Fallback to external location # Fallback to external location
return Path(__file__).parent.parent / "system_config.json" return Path(__file__).parent.parent / "system_config.json"
def load_protected_config() -> dict: def load_protected_config() -> dict:
global _protected_config global _protected_config
try: try:
@@ -52,19 +54,20 @@ def load_protected_config() -> dict:
except FileNotFoundError: except FileNotFoundError:
logging.warning("⚠️ system_config.json not found. Using built-in defaults.") logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
system_config = { system_config = {
"APPNAME": "AirlockTools", "APPNAME": "Loxide",
"PATH_EXCLUSION_CONST": 4, "PATH_EXCLUSION_CONST": 4,
"MIN_FILES_FOR_PATH": 4, "MIN_FILES_FOR_PATH": 4,
"VT_THREAT_TOLERANCE": 4, "VT_THREAT_TOLERANCE": 4,
"POLICY_MAP_ENF_AUD": { "POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"},
"enforced_id": "audit_id"
}
} }
_protected_config = {key: system_config[key] for key in PROTECTED_KEYS} _protected_config = {key: system_config[key] for key in PROTECTED_KEYS}
return _protected_config return _protected_config
def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
def get_protected_value(
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
) -> Optional[T]:
value = _protected_config.get(key) value = _protected_config.get(key)
if value is None: if value is None:
logging.warning(f"Protected config key '{key}' not found.") logging.warning(f"Protected config key '{key}' not found.")
@@ -74,9 +77,12 @@ def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default:
value = value.strip("'\"") value = value.strip("'\"")
return cast_type(value) return cast_type(value)
except (ValueError, TypeError): except (ValueError, TypeError):
logging.warning(f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}.") logging.warning(
f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}."
)
return default return default
def get_protected_json(key: str, default: str = "{}") -> dict: def get_protected_json(key: str, default: str = "{}") -> dict:
raw = _protected_config.get(key, default) raw = _protected_config.get(key, default)
if isinstance(raw, dict): if isinstance(raw, dict):
@@ -85,13 +91,11 @@ def get_protected_json(key: str, default: str = "{}") -> dict:
return json.loads(raw) return json.loads(raw)
except json.JSONDecodeError: except json.JSONDecodeError:
try: try:
escaped = raw.encode('unicode_escape').decode('utf-8') escaped = raw.encode("unicode_escape").decode("utf-8")
return json.loads(escaped) return json.loads(escaped)
except Exception as e: except Exception as e:
logging.error(f"Failed to parse protected JSON key '{key}': {e}") logging.error(f"Failed to parse protected JSON key '{key}': {e}")
return json.loads(default) return json.loads(default)
def load_env_json(key: str, default: str): def load_env_json(key: str, default: str):
@@ -100,13 +104,16 @@ def load_env_json(key: str, default: str):
return json.loads(raw) return json.loads(raw)
except json.JSONDecodeError: except json.JSONDecodeError:
try: try:
escaped = raw.encode('unicode_escape').decode('utf-8') escaped = raw.encode("unicode_escape").decode("utf-8")
return json.loads(escaped) return json.loads(escaped)
except Exception as e: except Exception as e:
logging.error(f"Failed to parse {key}: {e}") logging.error(f"Failed to parse {key}: {e}")
return json.loads(default) return json.loads(default)
def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
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. Safely retrieves an environment variable and casts it to the desired type.
@@ -126,5 +133,7 @@ def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T]
value = value.strip("'\"") # Strip surrounding quotes value = value.strip("'\"") # Strip surrounding quotes
return cast_type(value) return cast_type(value)
except (ValueError, TypeError): except (ValueError, TypeError):
logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.") logger.warning(
return default f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}."
)
return default
+53 -34
View File
@@ -21,9 +21,12 @@ from utils.utils import colorText, get_sanitized_input
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class Selector: class Selector:
@staticmethod @staticmethod
def _get_sorted_items(items: List[Any], label_func: Callable[[Any], str]) -> List[Any]: def _get_sorted_items(
items: List[Any], label_func: Callable[[Any], str]
) -> List[Any]:
return sorted(items, key=lambda item: label_func(item).lower()) return sorted(items, key=lambda item: label_func(item).lower())
@staticmethod @staticmethod
@@ -31,10 +34,10 @@ class Selector:
items: List[Any], items: List[Any],
label_func: Callable[[Any], str], label_func: Callable[[Any], str],
num_columns: int = 4, num_columns: int = 4,
header: str = "Available Choices:" header: str = "Available Choices:",
) -> None: ) -> None:
# Force single column if items are DataFrame rows # Force single column if items are DataFrame rows
if items and isinstance(items[0], (pd.Series, dict)): if items and isinstance(items[0], (pd.Series, dict)):
num_columns = 1 num_columns = 1
@@ -51,9 +54,7 @@ class Selector:
@staticmethod @staticmethod
def _display_selected_items( def _display_selected_items(
selected: List[Any], selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 4
label_func: Callable[[Any], str],
num_columns: int = 4
) -> None: ) -> None:
print(colorText("\nCurrent selections:", "cyan")) print(colorText("\nCurrent selections:", "cyan"))
if not selected: if not selected:
@@ -92,7 +93,7 @@ class Selector:
allow_multiple: bool = False, allow_multiple: bool = False,
prompt_each: bool = False, prompt_each: bool = False,
header: str = "Available Choices:", header: str = "Available Choices:",
num_columns: int = 4 num_columns: int = 4,
) -> Union[Optional[Any], List[Any]]: ) -> Union[Optional[Any], List[Any]]:
if not items: if not items:
logger.warning("No items available for selection.") logger.warning("No items available for selection.")
@@ -104,9 +105,19 @@ class Selector:
if allow_multiple: if allow_multiple:
while True: while True:
Selector._display_choices(remaining_items, label_func, num_columns=num_columns, header=header) Selector._display_choices(
Selector._display_selected_items(selected, label_func, num_columns=num_columns) remaining_items, label_func, num_columns=num_columns, header=header
choice = get_sanitized_input("Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: ").strip().lower() )
Selector._display_selected_items(
selected, label_func, num_columns=num_columns
)
choice = (
get_sanitized_input(
"Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: "
)
.strip()
.lower()
)
if choice == "q": if choice == "q":
break break
elif choice == "r": elif choice == "r":
@@ -125,10 +136,14 @@ class Selector:
logger.info(f"Selected: {label_func(item)}") logger.info(f"Selected: {label_func(item)}")
else: else:
logger.warning("Item already selected.") logger.warning("Item already selected.")
remaining_items = [item for item in remaining_items if item not in newly_selected] remaining_items = [
item for item in remaining_items if item not in newly_selected
]
return selected if selected else None return selected if selected else None
else: else:
Selector._display_choices(full_sorted_items, label_func, num_columns=num_columns, header=header) Selector._display_choices(
full_sorted_items, label_func, num_columns=num_columns, header=header
)
try: try:
choice = int(get_sanitized_input("Select one item by number: ")) choice = int(get_sanitized_input("Select one item by number: "))
if 1 <= choice <= len(full_sorted_items): if 1 <= choice <= len(full_sorted_items):
@@ -145,9 +160,14 @@ class Selector:
def select_with_mode( def select_with_mode(
items: List[Any], items: List[Any],
label_func: Callable[[Any], str], label_func: Callable[[Any], str],
header: str = "Available Choices:" header: str = "Available Choices:",
) -> List[Any]: ) -> List[Any]:
print(colorText("Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white")) print(
colorText(
"Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):",
"white",
)
)
mode = get_sanitized_input("").strip().lower() mode = get_sanitized_input("").strip().lower()
if mode == "a": if mode == "a":
return items return items
@@ -156,7 +176,7 @@ class Selector:
label_func=label_func, label_func=label_func,
allow_multiple=True, allow_multiple=True,
prompt_each=False, prompt_each=False,
header=header header=header,
) )
if not selected: if not selected:
return items return items
@@ -172,44 +192,38 @@ class Selector:
@staticmethod @staticmethod
def select_objects( def select_objects(
objects: List[Any], objects: List[Any], allow_multiple: bool = False, prompt_each: bool = False
allow_multiple: bool = False,
prompt_each: bool = False
) -> Union[Optional[Any], List[Any]]: ) -> Union[Optional[Any], List[Any]]:
return Selector._select_from_list( return Selector._select_from_list(
objects, objects,
label_func=lambda obj: getattr(obj, "name", str(obj)), label_func=lambda obj: getattr(obj, "name", str(obj)),
allow_multiple=allow_multiple, allow_multiple=allow_multiple,
prompt_each=prompt_each, prompt_each=prompt_each,
header="Available Objects:" header="Available Objects:",
) )
@staticmethod @staticmethod
def select_string( def select_string(
options: List[str], options: List[str], allow_multiple: bool = False, prompt_each: bool = False
allow_multiple: bool = False,
prompt_each: bool = False
) -> Union[Optional[str], List[str]]: ) -> Union[Optional[str], List[str]]:
return Selector._select_from_list( return Selector._select_from_list(
options, options,
label_func=str, label_func=str,
allow_multiple=allow_multiple, allow_multiple=allow_multiple,
prompt_each=prompt_each, prompt_each=prompt_each,
header="Available Options:" header="Available Options:",
) )
@staticmethod @staticmethod
def select_int( def select_int(
options: List[int], options: List[int], allow_multiple: bool = False, prompt_each: bool = False
allow_multiple: bool = False,
prompt_each: bool = False
) -> Union[Optional[int], List[int]]: ) -> Union[Optional[int], List[int]]:
return Selector._select_from_list( return Selector._select_from_list(
options, options,
label_func=lambda x: str(x), label_func=lambda x: str(x),
allow_multiple=allow_multiple, allow_multiple=allow_multiple,
prompt_each=prompt_each, prompt_each=prompt_each,
header="Available Integers:" header="Available Integers:",
) )
@staticmethod @staticmethod
@@ -217,7 +231,7 @@ class Selector:
prompt: str, prompt: str,
value_type: type = int, value_type: type = int,
valid_range: Optional[tuple] = None, valid_range: Optional[tuple] = None,
allow_quit: bool = False allow_quit: bool = False,
) -> Optional[Any]: ) -> Optional[Any]:
while True: while True:
user_input = get_sanitized_input(prompt).strip().lower() user_input = get_sanitized_input(prompt).strip().lower()
@@ -255,7 +269,7 @@ class Selector:
columns: Optional[List[str]] = None, columns: Optional[List[str]] = None,
allow_multiple: bool = False, allow_multiple: bool = False,
prompt_each: bool = False, prompt_each: bool = False,
header: str = "Available Rows:" header: str = "Available Rows:",
) -> List[pd.Series]: ) -> List[pd.Series]:
if df.empty: if df.empty:
print("DataFrame is empty.") print("DataFrame is empty.")
@@ -272,7 +286,7 @@ class Selector:
label_func=label_func, label_func=label_func,
allow_multiple=allow_multiple, allow_multiple=allow_multiple,
prompt_each=prompt_each, prompt_each=prompt_each,
header=header header=header,
) )
if isinstance(result, pd.Series): if isinstance(result, pd.Series):
@@ -286,7 +300,7 @@ class Selector:
def select_dataframe_with_mode( def select_dataframe_with_mode(
df: pd.DataFrame, df: pd.DataFrame,
columns: Optional[List[str]] = None, columns: Optional[List[str]] = None,
header: str = "Available Rows:" header: str = "Available Rows:",
) -> List[pd.Series]: ) -> List[pd.Series]:
if df.empty: if df.empty:
print("⚠️ DataFrame is empty.") print("⚠️ DataFrame is empty.")
@@ -305,7 +319,12 @@ class Selector:
print(f"{i}: {label_func(row)}") print(f"{i}: {label_func(row)}")
# Prompt for mode once # Prompt for mode once
print(colorText("\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white")) print(
colorText(
"\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):",
"white",
)
)
mode = get_sanitized_input("").strip().lower() mode = get_sanitized_input("").strip().lower()
if mode == "a": if mode == "a":
@@ -317,7 +336,7 @@ class Selector:
label_func=label_func, label_func=label_func,
allow_multiple=True, allow_multiple=True,
prompt_each=False, prompt_each=False,
header=header header=header,
) )
if not selected: if not selected:
@@ -331,4 +350,4 @@ class Selector:
return [pd.Series(row) for row in items if row not in selected] return [pd.Series(row) for row in items if row not in selected]
else: else:
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow")) print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
return [] return []
+39 -35
View File
@@ -30,16 +30,16 @@ from utils.configmanager import PROTECTED_KEYS, load_protected_config
def get_base_directory() -> Path: def get_base_directory() -> Path:
system = platform.system() system = platform.system()
home = Path.home() home = Path.home()
if system == 'Windows': if system == "Windows":
return Path(os.getenv('APPDATA', home / 'AppData' / 'Roaming')) / "AirlockTools" return Path(os.getenv("APPDATA", home / "AppData" / "Roaming")) / "Loxide"
elif system == 'Darwin': elif system == "Darwin":
return home / 'Library' / 'Application Support' / "AirlockTools" return home / "Library" / "Application Support" / "Loxide"
else: else:
return home / '.local' / 'share' / "AirlockTools" return home / ".local" / "share" / "Loxide"
def configure_logging(log_dir: Path, log_level: str = "DEBUG"): def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
log_file = log_dir / "airlocktools.log" log_file = log_dir / "Loxide.log"
config = { config = {
"version": 1, # Required key for dictConfig format version "version": 1, # Required key for dictConfig format version
@@ -58,17 +58,17 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
"file": { "file": {
"class": "logging.handlers.TimedRotatingFileHandler", "class": "logging.handlers.TimedRotatingFileHandler",
"filename": str(log_file), "filename": str(log_file),
"when": "midnight", # Rotate logs at midnight "when": "midnight", # Rotate logs at midnight
"interval": 1, # Every 1 day "interval": 1, # Every 1 day
"backupCount": 7, # Keep 7 days of logs "backupCount": 7, # Keep 7 days of logs
"encoding": "utf-8", # Ensure UTF-8 encoding "encoding": "utf-8", # Ensure UTF-8 encoding
"level": "DEBUG", # Always log DEBUG and above "level": "DEBUG", # Always log DEBUG and above
"formatter": "detailed", # Use detailed format "formatter": "detailed", # Use detailed format
}, },
"console": { "console": {
"class": "logging.StreamHandler", "class": "logging.StreamHandler",
"level": log_level.upper(), # Configurable log level "level": log_level.upper(), # Configurable log level
"formatter": "simple", # Use simple format "formatter": "simple", # Use simple format
}, },
}, },
"root": { "root": {
@@ -82,9 +82,9 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
try: try:
config["handlers"]["eventlog"] = { config["handlers"]["eventlog"] = {
"class": "logging.handlers.NTEventLogHandler", "class": "logging.handlers.NTEventLogHandler",
"appname": "AirlockTools", # Event log source name "appname": "Loxide", # Event log source name
"level": "CRITICAL", # Only log critical errors "level": "CRITICAL", # Only log critical errors
"formatter": "simple", # Use simple format "formatter": "simple", # Use simple format
} }
config["root"]["handlers"].append("eventlog") config["root"]["handlers"].append("eventlog")
except Exception as e: except Exception as e:
@@ -94,9 +94,11 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
logging.config.dictConfig(config) logging.config.dictConfig(config)
logging.getLogger().debug("✅ Logging configured.") logging.getLogger().debug("✅ Logging configured.")
def get_system_config_path() -> Path: def get_system_config_path() -> Path:
base_path = Path(getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__)))) base_path = Path(
getattr(sys, "_MEIPASS", os.path.dirname(os.path.abspath(__file__)))
)
return base_path.parent / "system_config.json" return base_path.parent / "system_config.json"
@@ -108,45 +110,45 @@ def load_system_config() -> dict:
except FileNotFoundError: except FileNotFoundError:
logging.warning("⚠️ system_config.json not found. Using built-in defaults.") logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
return { return {
"APPNAME": "AirlockTools", "APPNAME": "Loxide",
"LOG_LEVEL": "DEBUG", "LOG_LEVEL": "DEBUG",
"PATH_EXCLUSION_CONST": 4, "PATH_EXCLUSION_CONST": 4,
"MIN_FILES_FOR_PATH": 4, "MIN_FILES_FOR_PATH": 4,
"VT_THREAT_TOLERANCE": 4, "VT_THREAT_TOLERANCE": 4,
"POLICY_MAP_ENF_AUD": { "POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"},
"enforced_id": "audit_id"
}
} }
def load_user_config(config_dir: Path) -> dict: def load_user_config(config_dir: Path) -> dict:
user_config_path = config_dir / "user_config.json" user_config_path = config_dir / "user_config.json"
if not user_config_path.exists(): if not user_config_path.exists():
default_user_config = { default_user_config = {"URL": "", "LOG_LEVEL": "INFO"}
"URL": "",
"LOG_LEVEL": "INFO"
}
with open(user_config_path, "w") as f: with open(user_config_path, "w") as f:
json.dump(default_user_config, f, indent=4) json.dump(default_user_config, f, indent=4)
logging.debug(f"Created user config at {user_config_path}") logging.debug(f"Created user config at {user_config_path}")
with open(user_config_path, "r") as f: with open(user_config_path, "r") as f:
return json.load(f) return json.load(f)
def write_config_to_env(config: dict, env_path: Path): def write_config_to_env(config: dict, env_path: Path):
for key, value in config.items(): for key, value in config.items():
if key in PROTECTED_KEYS: if key in PROTECTED_KEYS:
continue # Skip protected keys continue # Skip protected keys
try: try:
serialized = json.dumps(value) if isinstance(value, (list, dict)) else str(value) serialized = (
json.dumps(value) if isinstance(value, (list, dict)) else str(value)
)
set_key(env_path, key, serialized) set_key(env_path, key, serialized)
except Exception as e: except Exception as e:
logging.warning(f"Failed to write {key} to .env: {e}") logging.warning(f"Failed to write {key} to .env: {e}")
def setup(): def setup():
base_dir = get_base_directory() base_dir = get_base_directory()
dirs = { dirs = {
'config': base_dir / 'config', "config": base_dir / "config",
'cache': base_dir / 'cache', "cache": base_dir / "cache",
'logs': base_dir / 'logs', "logs": base_dir / "logs",
} }
for name, path in dirs.items(): for name, path in dirs.items():
@@ -154,7 +156,7 @@ def setup():
logging.debug(f"{name.capitalize()} directory ensured at: {path}") logging.debug(f"{name.capitalize()} directory ensured at: {path}")
system_config = load_system_config() system_config = load_system_config()
configure_logging(dirs['logs'], system_config.get("LOG_LEVEL", "DEBUG")) configure_logging(dirs["logs"], system_config.get("LOG_LEVEL", "DEBUG"))
env_path = base_dir / ".env" env_path = base_dir / ".env"
if not env_path.exists(): if not env_path.exists():
@@ -171,7 +173,7 @@ def setup():
"Approved": [], "Approved": [],
"Needs_Review": ["Review_First", "Review_Second", "HTML"], "Needs_Review": ["Review_First", "Review_Second", "HTML"],
"Preflight": ["HTML"], "Preflight": ["HTML"],
"Archived": [] "Archived": [],
} }
for folder_name, subfolders in folders_structure.items(): for folder_name, subfolders in folders_structure.items():
@@ -183,7 +185,7 @@ def setup():
subfolder_path.mkdir(parents=True, exist_ok=True) subfolder_path.mkdir(parents=True, exist_ok=True)
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}") logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
user_config = load_user_config(dirs['config']) user_config = load_user_config(dirs["config"])
merged_config = {**system_config, **user_config} merged_config = {**system_config, **user_config}
protected_config = load_protected_config() protected_config = load_protected_config()
@@ -194,10 +196,12 @@ def setup():
if not url: if not url:
url = os.getenv("URL") url = os.getenv("URL")
if not url: if not url:
url = input("🌐 Enter the service URL (e.g., https://example.com/api): ").strip() url = input(
"🌐 Enter the service URL (e.g., https://example.com/api): "
).strip()
merged_config["URL"] = url merged_config["URL"] = url
set_key(env_path, "URL", url) set_key(env_path, "URL", url)
os.environ["URL"] = url os.environ["URL"] = url
logging.debug(f"Service URL set to: {url}") logging.debug(f"Service URL set to: {url}")
write_config_to_env(merged_config, env_path) write_config_to_env(merged_config, env_path)
View File
+126 -123
View File
@@ -5,7 +5,7 @@ import sys
import dotenv import dotenv
from dotenv import set_key from dotenv import set_key
from textual.app import App, ComposeResult from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical from textual.containers import Vertical
from textual.reactive import reactive from textual.reactive import reactive
from textual.screen import Screen from textual.screen import Screen
from textual.widgets import ( from textual.widgets import (
@@ -16,18 +16,24 @@ from textual.widgets import (
Static, Static,
Tab, Tab,
Tabs, Tabs,
Tree,
) )
from flows.otp import otp_activities_by_agent, otp_generate, otp_revoke from flows.otp import otp_activities_by_agent, otp_revoke
from flows.prepPolicy import menu_policy_enforce from flows.prepPolicy import menu_policy_enforce
from flows.quietAgent import findQuietAgents from flows.quietAgent import findQuietAgents
from models.agent import Agent
from models.policy import Policy
from screens.otpworkflowscreen import OTPWorkflowScreen
from services.agenthandler import findAgents, moveAgents, toggleEnforcement from services.agenthandler import findAgents, moveAgents, toggleEnforcement
from services.API import AirlockAPIWrapper from services.API import AirlockAPIWrapper
from services.policyhandler import confirmUpdateAfromE from services.policyhandler import confirmUpdateAfromE
from utils.configmanager import load_env from utils.configmanager import load_env
from utils.setup import get_base_directory, load_user_config from utils.setup import get_base_directory, load_user_config
from utils.utils import open_directory from utils.utils import open_directory
from widgets.multiagentselector import MultiAgentSelector
from widgets.OTP_generate import OTPGenerator
from widgets.policytreewidget import PolicyTreeWidget
from widgets.themeselector import ThemeSelector
dotenv.load_dotenv() dotenv.load_dotenv()
@@ -58,7 +64,9 @@ def _persist_user_theme(theme_name: str) -> None:
config_dir.mkdir(parents=True, exist_ok=True) config_dir.mkdir(parents=True, exist_ok=True)
if not user_config_path.exists(): if not user_config_path.exists():
# minimal default like your load_user_config does # minimal default like your load_user_config does
user_config_path.write_text('{"URL": "", "LOG_LEVEL": "INFO"}\n', encoding="utf-8") user_config_path.write_text(
'{"URL": "", "LOG_LEVEL": "INFO"}\n', encoding="utf-8"
)
# load existing user config # load existing user config
user_conf = load_user_config(config_dir) user_conf = load_user_config(config_dir)
@@ -86,8 +94,6 @@ def _persist_user_theme(theme_name: str) -> None:
logger.debug("Reloaded .env from %s", env_path) logger.debug("Reloaded .env from %s", env_path)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# 1) SCREEN # 1) SCREEN
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -105,7 +111,7 @@ class MainMenuScreen(Screen):
("🔀 - Move - Other", "move_other_button"), ("🔀 - Move - Other", "move_other_button"),
], ],
"otp": [ "otp": [
("🔐 - Generate OTPs", "otp_generate_button"), ("🎫 - Generate OTPs", "otp_generate_button"),
("📊 - OTP Activities By Agent", "otp_activities_button"), ("📊 - OTP Activities By Agent", "otp_activities_button"),
("❌ - Revoke OTPs", "otp_revoke_button"), ("❌ - Revoke OTPs", "otp_revoke_button"),
], ],
@@ -115,42 +121,24 @@ class MainMenuScreen(Screen):
], ],
} }
# textual themes to expose def __init__(self, api: AirlockAPIWrapper) -> None:
THEME_BUTTONS = [
("textual-dark", "textual-dark"),
("textual-light", "textual-light"),
("nord", "nord"),
("gruvbox", "gruvbox"),
("catppuccin-mocha", "catppuccin-mocha"),
("dracula", "dracula"),
("tokyo-night", "tokyo-night"),
("monokai", "monokai"),
("flexoki", "flexoki"),
("catppuccin-latte", "catppuccin-latte"),
("solarized-light", "solarized-light"),
]
def __init__(self) -> None:
super().__init__() super().__init__()
self.api = api
self.extras = load_env("EXTRAS") self.extras = load_env("EXTRAS")
wd = load_env("WORKING_DIR") or os.getcwd() wd = load_env("WORKING_DIR") or os.getcwd()
if not os.path.isdir(wd): if not os.path.isdir(wd):
wd = os.getcwd() wd = os.getcwd()
self.working_dir = wd self.working_dir = wd
def _make_buttons_for(self, tab_id: str) -> Vertical: def _make_buttons_for(self, tab_id: str) -> Vertical:
defs = self.BUTTON_DEFS.get(tab_id, []) defs = self.BUTTON_DEFS.get(tab_id, [])
buttons = [] buttons = []
for label, btn_id in defs: for label, btn_id in defs:
btn = Button(label, id=btn_id) btn = Button(label, id=btn_id)
btn.styles.width = "100%" # Make button span full width of parent btn.styles.width = "100%"
buttons.append(btn) buttons.append(btn)
return Vertical(*buttons) return Vertical(*buttons)
def compose(self) -> ComposeResult: def compose(self) -> ComposeResult:
yield Header(show_clock=True, icon="") yield Header(show_clock=True, icon="")
@@ -208,8 +196,7 @@ class MainMenuScreen(Screen):
new_index = current + direction new_index = current + direction
if 0 <= new_index < len(buttons): if 0 <= new_index < len(buttons):
buttons[new_index].focus() buttons[new_index].focus()
def switch_tab(self, tab_id: str) -> None: def switch_tab(self, tab_id: str) -> None:
self.current_tab = tab_id self.current_tab = tab_id
content = self.query_one("#content", Vertical) content = self.query_one("#content", Vertical)
@@ -221,88 +208,62 @@ class MainMenuScreen(Screen):
elif tab_id == "dir": elif tab_id == "dir":
content.mount(DirectoryTree(self.working_dir, id="dir_tree")) content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
elif tab_id == "p_tree": elif tab_id == "p_tree":
layout = Horizontal() content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
content.mount(layout)
# Left: Policy Tree
policy_tree = Tree("Policies", id="policy_tree")
policy_tree.styles.width = "2fr"
layout.mount(policy_tree)
# Right: Details pane
details_pane = Static("Select a policy or device to view details", id="details-pane")
details_pane.styles.width = "3fr"
layout.mount(details_pane)
# Build the tree
node_map = {}
# Top-level policies
for _, policy in self.app.policies.iterrows():
if policy["parent"] == "global-policy-settings":
node = policy_tree.root.add(label=policy["name"], data=policy.to_dict())
node_map[policy["groupid"]] = node
# Child policies
for _, policy in self.app.policies.iterrows():
parent_id = policy["parent"]
if parent_id in node_map:
parent_node = node_map[parent_id]
node = parent_node.add(label=policy["name"], data=policy.to_dict())
node_map[policy["groupid"]] = node
# Devices under policies
for _, device in self.app.devices.iterrows():
group_id = device["groupid"]
if group_id in node_map:
parent_node = node_map[group_id]
label = device["hostname"] # Keep tree clean
parent_node.add(label=label, data=device.to_dict())
elif tab_id == "settings": elif tab_id == "settings":
# Create and mount the horizontal container content.mount(ThemeSelector())
horizontal_container = Horizontal(id="settings_grid")
horizontal_container.styles.layout = "horizontal"
horizontal_container.styles.height = "auto"
content.mount(Static("Theme Options"))
content.mount(horizontal_container) # Mount the horizontal container first
# Create 3 columns
for i in range(1):
column = Vertical()
column.styles.width = "1fr"
column.styles.height = "auto"
horizontal_container.mount(column) # Mount each column
for j in range(i, len(self.THEME_BUTTONS), 1):
if j < len(self.THEME_BUTTONS):
label, btn_id = self.THEME_BUTTONS[j]
button = Button(label, id=f"set_theme_{btn_id}", compact=True)
#button.styles.width = "100%"
column.mount(button) # Mount each button
else: else:
content.mount(Static(f"Unknown tab: {tab_id}")) content.mount(Static(f"Unknown tab: {tab_id}"))
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
self.switch_tab(event.tab.id) self.switch_tab(event.tab.id)
def on_tree_node_selected(self, message: Tree.NodeSelected) -> None: def on_multi_agent_selector_agents_selected(
node = message.node self, message: MultiAgentSelector.AgentsSelected
data = node.data ) -> None:
"""Handle selected agents from MultiAgentSelector."""
global _PENDING_JOB
selected_agents = message.selected_agents
logger.info("Selected agents: %s", selected_agents)
# TODO: Implement actual handling of selected agents
_PENDING_JOB = ("multi_agent_action", selected_agents)
self.app.exit()
details_pane = self.query_one("#details-pane", Static) def on_theme_selector_theme_selected(
self, message: ThemeSelector.ThemeSelected
) -> None:
"""Handle theme selection from ThemeSelector."""
global _PENDING_JOB
_persist_user_theme(message.theme_name)
_PENDING_JOB = ("restart",)
self.app.exit()
if data: def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
details = "\n".join(f"{key}: {value}" for key, value in data.items()) """Handle OTP generation request from the workflow."""
else: global _PENDING_JOB
details = f"Selected: {node.label}"
details_pane.update(details) # Log what we received
logger.info(
"OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d",
len(message.devices),
message.requestor,
message.reasoning,
message.duration,
)
def on_directory_tree_file_selected(self, event: DirectoryTree.FileSelected) -> None: # Set up the job to run the OTP generation
_PENDING_JOB = (
"otp_workflow",
message.devices,
message.requestor,
message.reasoning,
message.duration,
)
self.app.exit()
def on_directory_tree_file_selected(
self, event: DirectoryTree.FileSelected
) -> None:
path = event.path path = event.path
logger.debug("Directory file selected: %s", path) logger.debug("Directory file selected: %s", path)
try: try:
@@ -316,14 +277,6 @@ class MainMenuScreen(Screen):
button_id = event.button.id button_id = event.button.id
logger.debug("Button pressed: %s", button_id) logger.debug("Button pressed: %s", button_id)
# theme selection → user config
if button_id.startswith("set_theme_"):
theme_name = button_id.replace("set_theme_", "")
_persist_user_theme(theme_name)
_PENDING_JOB = ("restart",)
self.app.exit()
return
match button_id: match button_id:
case "find_device_button": case "find_device_button":
_PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {}) _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {})
@@ -341,7 +294,10 @@ class MainMenuScreen(Screen):
case "move_other_button": case "move_other_button":
_PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {}) _PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {})
case "otp_generate_button": case "otp_generate_button":
_PENDING_JOB = ("legacy", otp_generate, (self.app.api,), {}) # NEW: Push OTP workflow screen instead of legacy function
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
event.stop()
return # Don't exit the app
case "otp_activities_button": case "otp_activities_button":
_PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {}) _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {})
case "otp_revoke_button": case "otp_revoke_button":
@@ -359,14 +315,10 @@ class MainMenuScreen(Screen):
self.app.exit() self.app.exit()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# 2) APP # 2) APP
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
class AirlockTools(App): class Loxide(App):
CSS = """ CSS = """
#logo { #logo {
width: 100%; width: 100%;
@@ -388,12 +340,23 @@ class AirlockTools(App):
if not os.path.isdir(wd): if not os.path.isdir(wd):
wd = os.getcwd() wd = os.getcwd()
self.working_dir = wd self.working_dir = wd
self.policies = api.policy_find_all()
self.devices = api.agent_find_all()
def on_mount(self) -> None: # Add error handling for API calls
try:
self.policies = [
Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()
]
self.devices = [
Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()
]
except Exception as exc:
logger.error("Failed to load policies/devices: %s", exc)
self.policies = None
self.devices = None
def on_mount(self, api: AirlockAPIWrapper) -> None:
self.theme = self._textual_theme self.theme = self._textual_theme
self.push_screen(MainMenuScreen()) self.push_screen(MainMenuScreen(api))
def action_quit(self) -> None: def action_quit(self) -> None:
global _PENDING_JOB global _PENDING_JOB
@@ -407,7 +370,6 @@ class AirlockTools(App):
screen.switch_tab("dir") screen.switch_tab("dir")
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# 3) TERMINAL + LEGACY # 3) TERMINAL + LEGACY
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -422,6 +384,7 @@ def _restore_terminal_for_legacy() -> None:
if os.name == "nt": if os.name == "nt":
try: try:
import ctypes import ctypes
kernel32 = ctypes.windll.kernel32 kernel32 = ctypes.windll.kernel32
handle = kernel32.GetStdHandle(-11) handle = kernel32.GetStdHandle(-11)
mode = ctypes.c_ulong() mode = ctypes.c_ulong()
@@ -447,7 +410,7 @@ def _run_legacy_job(func, args, kwargs) -> None:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# 4) PUBLIC ENTRYPOINT # 4) PUBLIC ENTRYPOINT
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def run_AirlockTools(api: AirlockAPIWrapper) -> None: def run_Loxide(api: AirlockAPIWrapper) -> None:
global _PENDING_JOB global _PENDING_JOB
while True: while True:
@@ -456,7 +419,7 @@ def run_AirlockTools(api: AirlockAPIWrapper) -> None:
dotenv.load_dotenv(dotenv_path=env_path, override=True) dotenv.load_dotenv(dotenv_path=env_path, override=True)
_PENDING_JOB = None _PENDING_JOB = None
app = AirlockTools(api) app = Loxide(api)
try: try:
app.run() app.run()
@@ -478,6 +441,46 @@ def run_AirlockTools(api: AirlockAPIWrapper) -> None:
# just loop again; fresh .env was already loaded at the top # just loop again; fresh .env was already loaded at the top
continue continue
if job[0] == "multi_agent_action":
# Handle multi-agent selection
logger.info("Multi-agent action with selected agents: %s", job[1])
continue
# NEW: Handle OTP workflow
if job[0] == "otp_workflow":
_, devices, requestor, reasoning, duration = job
# Call your OTP generation with the parameters
def otp_generate_with_params():
print(f"\n{'='*60}")
print("OTP GENERATION")
print(f"{'='*60}")
print(f"Requestor: {requestor}")
print(f"Reasoning: {reasoning}")
print(f"Duration: {duration} minutes")
print(f"\nGenerating OTPs for {len(devices)} devices:")
print(f"{'='*60}\n")
# Call your actual OTP generation function
# You'll need to adapt otp_generate to accept these parameters
# For now, this is a placeholder showing the structure
for device in devices:
print(f"Device: {device}")
print(f" Requestor: {requestor}")
print(f" Reason: {reasoning}")
print(f" Duration: {duration} minutes")
# TODO: Actually call your API to generate OTP
# result = api.generate_otp(device, requestor, reasoning, duration)
print()
print(f"{'='*60}")
print("OTP Generation Complete!")
print(f"{'='*60}")
_run_legacy_job(otp_generate_with_params, (), {})
continue
break break
@@ -486,4 +489,4 @@ def run_AirlockTools(api: AirlockAPIWrapper) -> None:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
if __name__ == "__main__": if __name__ == "__main__":
api = AirlockAPIWrapper() api = AirlockAPIWrapper()
run_AirlockTools(api) run_Loxide(api)
+57 -22
View File
@@ -28,8 +28,6 @@ import pandas as pd
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def import_to_dataframe(file_path: str) -> pd.DataFrame: def import_to_dataframe(file_path: str) -> pd.DataFrame:
df = pd.DataFrame() df = pd.DataFrame()
@@ -96,17 +94,17 @@ def choose_file(initial_directory=None, required_substring=None):
return file_path return file_path
def get_sanitized_input(prompt: str) -> str: def get_sanitized_input(prompt: str) -> str:
while True: while True:
user_input = input(prompt) user_input = input(prompt)
if user_input.strip() == "": if user_input.strip() == "":
return user_input # Allow blank lines return user_input # Allow blank lines
if re.match(r'^[a-zA-Z0-9_\- .]+$', user_input.strip()): if re.match(r"^[a-zA-Z0-9_\- .]+$", user_input.strip()):
return user_input return user_input
else: else:
print("Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed.") print(
"Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed."
)
def regulator(paths, case_insensitive=True): def regulator(paths, case_insensitive=True):
@@ -119,8 +117,10 @@ def regulator(paths, case_insensitive=True):
pattern = "(?i)" + pattern # Add inline case-insensitive flag pattern = "(?i)" + pattern # Add inline case-insensitive flag
print(f"Regulator is providing: {pattern}") print(f"Regulator is providing: {pattern}")
return pattern return pattern
def irtang(): def irtang():
print( print(
colorText( colorText(
r""" r"""
@@ -149,6 +149,8 @@ def irtang():
"yellow", "yellow",
) )
) )
def displayIntro(): def displayIntro():
print( print(
@@ -164,6 +166,8 @@ def displayIntro():
"cyan", "cyan",
) )
) )
def welcome(): def welcome():
print( print(
colorText( colorText(
@@ -184,11 +188,21 @@ def welcome():
) )
) )
def section_header(title):
print(colorText("\n --------------------------------------------------------------------", "cyan"))
print(colorText(f" ------------- {title} -------------", "cyan"))
print(colorText(" --------------------------------------------------------------------", "cyan"))
def section_header(title):
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(colorText(f" ------------- {title} -------------", "cyan"))
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
def areYouSure(): def areYouSure():
@@ -348,7 +362,11 @@ def printDeviceEnforceChecklist():
"cyan", "cyan",
) )
) )
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan")) print(
colorText(
" When complete, save the csv file to the directory 'approved'", "cyan"
)
)
print( print(
colorText( colorText(
" Do the same process with the list of publishers forthe same directories", " Do the same process with the list of publishers forthe same directories",
@@ -357,20 +375,38 @@ def printDeviceEnforceChecklist():
) )
print(colorText(" Preflight Lists will be generated", "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(
"5. Choose the destination policy and parent and child allow list", "cyan"
)
)
print(colorText("6. Test ------------------------------------------------------", "cyan")) print(
colorText(
"6. Test ------------------------------------------------------", "cyan"
)
)
print(colorText(" Print rather than apply selected data.", "cyan")) print(colorText(" Print rather than apply selected data.", "cyan"))
print(colorText("7. Liftoff ------------------------------------------------------", "cyan")) print(
colorText(
"7. Liftoff ------------------------------------------------------", "cyan"
)
)
print( print(
colorText( colorText(
" Apply path exclusions according to allowed and approved paths", " Apply path exclusions according to allowed and approved paths",
"cyan", "cyan",
) )
) )
print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan")) print(
print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan")) 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( print(
colorText( colorText(
@@ -523,10 +559,9 @@ def formatHTML(df, output_html_path=None, overwrite=True):
return styled_html return styled_html
def open_directory(path): def open_directory(path):
system = platform.system() system = platform.system()
if system == "Windows": if system == "Windows":
os.startfile(path) os.startfile(path)
elif system == "Linux": elif system == "Linux":
@@ -537,9 +572,9 @@ def open_directory(path):
def print_x_wide(items: list, width: int): def print_x_wide(items: list, width: int):
for i in range(0, len(items), width): for i in range(0, len(items), width):
row = items[i:i+width] row = items[i : i + width]
print(" | ".join(row)) print(" | ".join(row))
def clear_screen(): def clear_screen():
os.system('cls' if os.name == 'nt' else 'clear') os.system("cls" if os.name == "nt" else "clear")
+346
View File
@@ -0,0 +1,346 @@
import logging
from typing import List
from textual.containers import Horizontal, Vertical
from textual.css.query import NoMatches
from textual.message import Message
from textual.reactive import reactive
from textual.widget import Widget
from textual.widgets import Button, Input, RadioButton, RadioSet, Static, TextArea
from models.agent import Agent
logger = logging.getLogger(__name__)
class OTPGenerator(Widget):
# Reactive properties to track form completion
requestor_filled = reactive(False)
reasoning_filled = reactive(False)
duration_selected = reactive(True) # Default is selected
otp_generated = reactive(False)
class OTPInfo(Message):
def __init__(
self, devices: List[Agent], requestor: str, reasoning: str, duration: int
):
super().__init__()
self.devices = devices
self.requestor = requestor
self.reasoning = reasoning
self.duration = duration
# Duration options in minutes
DURATION_OPTIONS = [
(15, "15 minutes"),
(60, "1 hour"),
(360, "6 hours"),
(1440, "1 day"),
(10080, "7 days"),
]
def __init__(self, devices: List[Agent]):
"""Initialize with a list of Agent objects."""
super().__init__()
self.devices = devices
def watch_requestor_filled(self, old_value: bool, new_value: bool) -> None:
"""Update button state when requestor changes."""
self._update_button_state()
def watch_reasoning_filled(self, old_value: bool, new_value: bool) -> None:
"""Update button state when reasoning changes."""
self._update_button_state()
def watch_otp_generated(self, old_value: bool, new_value: bool) -> None:
"""Update button state when OTP is generated."""
self._update_button_state()
def _update_button_state(self) -> None:
"""Enable/disable the generate button based on form state."""
try:
button = self.query_one("#generate_button", Button)
# Enable only if all fields filled and OTP not yet generated
button.disabled = not (
self.requestor_filled
and self.reasoning_filled
and not self.otp_generated
)
except NoMatches:
pass
def compose(self):
title_text = Static(
f"🎫 Generate One Time Passes for {len(self.devices)} device(s)",
id="otpgen_title",
)
title_text.styles.margin = (0, 0, 1, 0)
yield title_text
with Horizontal() as main_layout:
main_layout.styles.height = "auto"
# Left side - Inputs and controls
with Vertical() as left_side:
left_side.styles.width = "1fr"
left_side.styles.height = "auto"
# Requestor input
requestor_label = Static("Who is requesting OTP?")
requestor_label.styles.margin = (0, 0, 0, 0)
yield requestor_label
requestor_box = Input(
placeholder="Enter requestor name", id="requestor_input"
)
requestor_box.styles.margin = (0, 0, 1, 0)
yield requestor_box
# Reasoning input
reasoning_label = Static("What work are they doing?")
reasoning_label.styles.margin = (0, 0, 0, 0)
yield reasoning_label
reasoning_box = Input(
placeholder="Enter reason for OTP", id="reasoning_input"
)
reasoning_box.styles.margin = (0, 0, 1, 0)
yield reasoning_box
# Duration selection
duration_label = Static("Duration:")
duration_label.styles.margin = (0, 0, 0, 0)
yield duration_label
with RadioSet(id="duration_radio") as radio_set:
radio_set.styles.margin = (0, 0, 1, 0)
for minutes, label in self.DURATION_OPTIONS:
radio = RadioButton(label, id=f"duration_{minutes}")
if minutes == 360: # Default to 6 hours
radio.value = True
yield radio
# Buttons in a horizontal layout
with Horizontal() as button_row:
button_row.styles.height = "auto"
button_row.styles.margin = (1, 0, 0, 0)
back_button = Button("← Back", id="back_button")
back_button.styles.width = "1fr"
yield back_button
generate_button = Button(
"Generate OTP", id="generate_button", variant="primary"
)
generate_button.styles.width = "2fr"
yield generate_button
# Right side - Show device list initially, then output after generation
with Vertical() as right_side:
right_side.styles.width = "2fr"
right_side.styles.height = "100%"
output_label = Static(
f"Selected Devices ({len(self.devices)}):", id="output_label"
)
output_label.styles.margin = (0, 0, 0, 0)
yield output_label
# Container for either device list or output
with Vertical(id="output_container") as output_container:
output_container.styles.height = "1fr"
output_container.styles.margin = (1, 0, 0, 0)
output_container.styles.overflow_y = "auto"
output_container.styles.border = ("round", "green")
# Show device list initially
device_list_text = "\n".join(
f"{device.hostname}" for device in self.devices
)
device_display = Static(device_list_text, id="device_display")
yield device_display
# Copy to clipboard button (hidden initially)
copy_button = Button("📋 Copy to Clipboard", id="copy_clipboard_button")
copy_button.styles.margin = (1, 0, 0, 0)
copy_button.styles.display = "none"
yield copy_button
def on_mount(self) -> None:
"""Set initial button state."""
self._update_button_state()
def on_input_changed(self, event: Input.Changed) -> None:
"""Handle input field changes."""
input_id = event.input.id
if input_id == "requestor_input":
self.requestor_filled = bool(event.value.strip())
elif input_id == "reasoning_input":
self.reasoning_filled = bool(event.value.strip())
def on_button_pressed(self, event: Button.Pressed):
btn_id = event.button.id
if btn_id == "back_button":
self.app.pop_screen()
event.stop()
elif btn_id == "copy_clipboard_button":
try:
output_area = self.query_one("#otp_output", TextArea)
text_to_copy = output_area.text
import pyperclip
pyperclip.copy(text_to_copy)
self.app.notify(
"✅ Copied to clipboard!", severity="information", timeout=2
)
except ImportError:
self.app.notify(
"⚠️ pyperclip not installed. Run: pip install pyperclip",
severity="warning",
)
except Exception as e:
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
event.stop()
elif btn_id == "generate_button":
try:
requestor = self.query_one("#requestor_input", Input).value.strip()
reasoning = self.query_one("#reasoning_input", Input).value.strip()
radio_set = self.query_one("#duration_radio", RadioSet)
selected_button_id = (
radio_set.pressed_button.id if radio_set.pressed_button else None
)
if not selected_button_id:
self._show_error("Please select a duration")
return
duration = int(selected_button_id.replace("duration_", ""))
if not requestor or not reasoning:
self._show_error("Please fill in all fields")
return
self.otp_generated = True
# Access API from the app - this is the key change!
api = self.app.api
output_lines = [
"=" * 60,
"OTP GENERATION RESULTS",
"=" * 60,
]
otp_dict = {}
for device in self.devices:
try:
otp_code = api.otp_generate(device.agentid, duration, reasoning)
otp_dict[device.hostname] = otp_code
logger.debug(f"Generated OTP for {device.hostname}: {otp_code}")
except Exception as e:
otp_dict[device.hostname] = f"ERROR: {str(e)}"
logger.error(
f"Failed to generate OTP for {device.hostname}: {e}"
)
for hostname, otp_code in otp_dict.items():
output_lines.append(f"{hostname:30} | {otp_code}")
output_lines.append("=" * 60)
result_text = "\n".join(output_lines)
self._show_result(result_text)
# Post message with the OTP info
self.post_message(
self.OTPInfo(self.devices, requestor, reasoning, duration)
)
event.stop()
except NoMatches:
self._show_error("UI elements not found")
except Exception as e:
self._show_error(f"Error: {str(e)}")
logger.exception("Error generating OTP")
def _show_error(self, message: str):
"""Display error message in output area."""
try:
container = self.query_one("#output_container", Vertical)
try:
device_display = self.query_one("#device_display", Static)
device_display.remove()
except NoMatches:
pass
try:
output_area = self.query_one("#otp_output", TextArea)
except NoMatches:
output_area = TextArea(id="otp_output", read_only=True)
container.mount(output_area)
output_area.text = f"❌ ERROR: {message}"
except Exception as e:
logger.debug(f"Error showing error message: {e}")
def _show_result(self, message: str):
"""Display result message in output area."""
try:
container = self.query_one("#output_container", Vertical)
try:
device_display = self.query_one("#device_display", Static)
device_display.remove()
except NoMatches:
pass
try:
output_area = self.query_one("#otp_output", TextArea)
except NoMatches:
output_area = TextArea(id="otp_output", read_only=True)
container.mount(output_area)
output_area.text = message
output_label = self.query_one("#output_label", Static)
output_label.update("Generated OTP Details:")
copy_button = self.query_one("#copy_clipboard_button", Button)
copy_button.styles.display = "block"
except Exception as e:
logger.debug(f"Error showing result: {e}")
def display_otp_result(self, result_text: str):
"""Display OTP generation result in the output area."""
try:
container = self.query_one("#output_container", Vertical)
try:
device_display = self.query_one("#device_display", Static)
device_display.remove()
except NoMatches:
pass
try:
output_area = self.query_one("#otp_output", TextArea)
except NoMatches:
output_area = TextArea(id="otp_output", read_only=True)
container.mount(output_area)
output_area.text = result_text
except Exception as e:
logger.debug(f"Error displaying OTP result: {e}")
def clear_form(self):
"""Clear all input fields and reset state."""
try:
self.query_one("#requestor_input", Input).value = ""
self.query_one("#reasoning_input", Input).value = ""
self.query_one("#otp_output", TextArea).text = ""
self.otp_generated = False
self.requestor_filled = False
self.reasoning_filled = False
self._update_button_state()
except NoMatches:
pass
+196
View File
@@ -0,0 +1,196 @@
import difflib
import re
from typing import List
from textual.containers import Horizontal, Vertical
from textual.css.query import NoMatches
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Button, SelectionList, Static, Switch, TextArea
from models.agent import Agent
class MultiAgentSelector(Widget):
class AgentsSelected(Message):
def __init__(self, selected_agents: List[Agent]):
super().__init__()
self.selected_agents = selected_agents
def __init__(self, all_agents: List[Agent]):
super().__init__()
self.all_agents = all_agents
self._match_type = "exact"
@property
def match_type(self):
return self._match_type
@match_type.setter
def match_type(self, value):
self._match_type = value
def compose(self):
title_text = Static("🖧 Multi-Agent Selector", id="selector_title")
title_text.styles.margin = (0, 0, 0, 1)
yield title_text
with Horizontal() as main_layout:
main_layout.styles.height = "auto"
# Left side - Input and controls
with Vertical() as left_pane:
left_pane.styles.width = "1fr"
left_pane.styles.height = "auto"
text_area = TextArea(
id="device_input",
placeholder="Paste device names here (one per line). Supports wildcards: * and ?",
)
text_area.styles.height = 10
text_area.styles.overflow_y = "auto"
yield text_area
with Horizontal(id="switch_search_container") as switch_search:
switch = Switch(value=False, id="match_switch")
switch.styles.width = "auto"
switch.styles.margin = (1, 0, 0, 0)
switch.styles.padding = (0, 0, 0, 0)
yield switch
switch_label = Static("Match: Exact", id="match_switch_label")
switch_label.styles.width = "auto"
switch_label.styles.margin = (2, 1, 0, 0)
yield switch_label
search = Button("🔍 Search", id="search_button")
search.styles.margin = (1, 0, 0, 0)
yield search
with Horizontal() as select_buttons:
select_buttons.styles.margin = (0, 0, 0, 0)
select_all_button = Button("✅ Select All", id="select_all")
select_all_button.styles.margin = (1, 1, 0, 1)
yield select_all_button
select_none_button = Button("🚫 Select None", id="select_none")
select_none_button.styles.margin = (1, 0, 0, 1)
yield select_none_button
with Horizontal() as button_row:
button_row.styles.height = "auto"
button_row.styles.margin = (1, 0, 0, 0)
back_button = Button("← Back", id="back_button")
back_button.styles.width = "1fr"
yield back_button
submit_button = Button(
"▶ Select & Continue", id="submit_selection", variant="primary"
)
submit_button.styles.margin = (0, 5, 2, 1)
submit_button.styles.padding = (0, 6, 0, 0)
yield submit_button
# Right side - Results
with Vertical() as right_pane:
right_pane.styles.width = "2fr"
yield SelectionList(id="match_results")
yield Static(id="unmatched_label")
def on_switch_changed(self, event: Switch.Changed):
self.match_type = "fuzzy" if event.value else "exact"
self.query_one("#match_switch_label", Static).update(
f"Match: {self.match_type.capitalize()}"
)
def on_button_pressed(self, event: Button.Pressed):
btn_id = event.button.id
try:
match_list = self.query_one("#match_results", SelectionList)
except NoMatches:
return
if btn_id == "back_button":
self.app.pop_screen()
event.stop()
elif btn_id == "select_all":
match_list.select_all()
event.stop()
elif btn_id == "select_none":
match_list.deselect_all()
event.stop()
elif btn_id == "submit_selection":
# Get selected hostnames
selected_hostnames = list(match_list.selected)
# Convert back to Agent objects
selected_agents = [
agent
for agent in self.all_agents
if agent.hostname in selected_hostnames
]
self.post_message(self.AgentsSelected(selected_agents))
event.stop()
elif btn_id == "search_button":
self.update_matches()
event.stop()
def update_matches(self):
raw_input = self.query_one("#device_input", TextArea).text.strip()
device_names = [line.strip() for line in raw_input.split("\n") if line.strip()]
matched, unmatched = self.match_devices(device_names)
match_list = self.query_one("#match_results", SelectionList)
match_list.clear_options()
for name in matched:
match_list.add_option((name, name))
unmatched_label = self.query_one("#unmatched_label", Static)
if unmatched:
unmatched_label.update(f"⚠️ No matches for: {', '.join(unmatched)}")
else:
unmatched_label.update("")
def match_devices(self, device_names: list[str]) -> tuple[list[str], list[str]]:
if not self.all_agents or not device_names:
return [], device_names
agent_names = [agent.hostname for agent in self.all_agents]
matched = set()
unmatched = []
for name in device_names:
# Check if the name contains wildcards
has_wildcards = "*" in name or "?" in name
if has_wildcards:
# Use regex for wildcard matching
pattern = re.escape(name)
pattern = pattern.replace(r"\*", ".*").replace(r"\?", ".")
regex = re.compile(f"^{pattern}$", re.IGNORECASE)
wildcard_matches = [
agent_name for agent_name in agent_names if regex.match(agent_name)
]
if wildcard_matches:
matched.update(wildcard_matches)
else:
unmatched.append(name)
elif self.match_type == "exact":
# Case-insensitive exact match
name_lower = name.lower()
exact_match = None
for agent_name in agent_names:
if agent_name.lower() == name_lower:
exact_match = agent_name
break
if exact_match:
matched.add(exact_match)
else:
unmatched.append(name)
else:
# Fuzzy match
matches = difflib.get_close_matches(name, agent_names, n=5, cutoff=0.5)
if matches:
matched.update(matches)
else:
unmatched.append(name)
return sorted(matched), unmatched
+206
View File
@@ -0,0 +1,206 @@
import logging
from rich.text import Text
from textual.containers import Horizontal, Vertical
from textual.widget import Widget
from textual.widgets import Input, OptionList, Static, Tree
from textual.widgets.option_list import Option
logger = logging.getLogger(__name__)
class PolicyTreeWidget(Widget):
"""Widget for displaying and searching a hierarchical policy tree."""
def __init__(self, policies, devices):
super().__init__()
self.policies = policies
self.devices = devices
self.last_highlighted_node = None
def compose(self):
# Left: Policy Tree
policy_tree = Tree("Policies", id="policy_tree")
policy_tree.styles.width = "2fr"
policy_tree.styles.height = "100%"
# Right: Search + Details
label = Static("Device Search:")
search_box = Input(
placeholder="Search policies or devices...", id="tree_search"
)
details_pane = Static("", id="details_pane")
with Horizontal():
yield policy_tree
with Vertical() as right_pane:
right_pane.styles.width = "3fr"
yield label
yield search_box
yield details_pane
def on_mount(self) -> None:
"""Build the tree after mounting."""
self._build_tree()
def _build_tree(self) -> None:
"""Build the policy tree structure."""
policy_tree = self.query_one("#policy_tree", Tree)
node_map = {}
# Top-level policies
for policy in self.policies:
if policy.parent == "global-policy-settings":
node = policy_tree.root.add(label=policy.name, data=policy)
node_map[policy.groupid] = node
# Child policies
for policy in self.policies:
parent_id = policy.parent
if parent_id in node_map:
parent_node = node_map[parent_id]
node = parent_node.add(label=policy.name, data=policy)
node_map[policy.groupid] = node
# Devices under policies
for device in self.devices:
group_id = device.groupid
if group_id in node_map:
parent_node = node_map[group_id]
label = device.hostname
parent_node.add(label=label, data=device)
def _collect_tree_nodes(self, node, all_nodes):
"""Helper to recursively collect all nodes from a tree."""
all_nodes.append(node)
for child in node.children:
self._collect_tree_nodes(child, all_nodes)
def _remove_match_selector(self):
"""Safely remove match selector widgets."""
try:
existing = self.query("#match_selector")
for widget in existing:
if widget.is_attached:
widget.remove()
except Exception as exc:
logger.debug("Failed to remove match_selector: %s", exc)
def on_tree_node_selected(self, message: Tree.NodeSelected) -> None:
"""Handle tree node selection."""
node = message.node
data = node.data
details_pane = self.query_one("#details_pane", Static)
# Reset previous highlight
if self.last_highlighted_node is not None:
original_label = str(self.last_highlighted_node.label).strip()
# Remove any styling
if isinstance(self.last_highlighted_node.label, Text):
original_label = self.last_highlighted_node.label.plain
self.last_highlighted_node.set_label(original_label)
# Apply highlight to current node
label_text = str(node.label).strip()
if isinstance(node.label, Text):
label_text = node.label.plain
highlighted_label = Text(label_text, style="reverse bold")
node.set_label(highlighted_label)
self.last_highlighted_node = node
# Update details pane
if data:
# Work with dataclass objects using __dict__
details = "\n".join(
f"{key}: {value}" for key, value in data.__dict__.items()
)
else:
details = f"Selected: {node.label}"
details_pane.update(details)
# Stop event from bubbling
message.stop()
def on_input_submitted(self, message: Input.Submitted) -> None:
"""Handle search input submission."""
# Remove existing match selector FIRST
self._remove_match_selector()
query = message.value.strip().lower()
tree = self.query_one("#policy_tree", Tree)
details_pane = self.query_one("#details_pane", Static)
all_nodes = []
self._collect_tree_nodes(tree.root, all_nodes)
label_to_node = {}
for node in all_nodes:
label_text = str(node.label).lower()
label_to_node[label_text] = node
if node.data:
# Use __dict__ for dataclass objects
data_dict = (
node.data.__dict__ if hasattr(node.data, "__dict__") else node.data
)
for key, value in data_dict.items():
if isinstance(value, str):
label_to_node[value.lower()] = node
# Wildcard-style substring match
matches = sorted([label for label in label_to_node if query in label])
if matches:
# Try to reuse existing match_selector or create new one
try:
option_list = self.query_one("#match_selector", OptionList)
option_list.clear_options()
option_list.display = True # Ensure it's visible
except:
option_list = OptionList(id="match_selector")
# Mount to the details pane's parent (the Vertical container)
details_pane.parent.mount(option_list)
for label in matches:
option_list.add_option(Option(label, id=f"match_{label}"))
details_pane.update(f"Found {len(matches)} matches. Select one below.")
else:
# Hide or remove the match_selector when no matches
self._remove_match_selector()
details_pane.update("No matches found.")
def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None:
"""Handle selection from search results."""
selected_id = event.option.id.replace("match_", "")
tree = self.query_one("#policy_tree", Tree)
details_pane = self.query_one("#details_pane", Static)
# Find the node
all_nodes = []
self._collect_tree_nodes(tree.root, all_nodes)
label_to_node = {str(node.label).lower(): node for node in all_nodes}
match_node = label_to_node.get(selected_id.lower())
if match_node:
# Expand path (original working logic)
node = match_node
path = []
while node:
path.insert(0, node)
node = node.parent
for node in path:
node.expand()
tree.select_node(match_node)
tree.scroll_to_node(match_node)
match_node.set_label(Text(str(match_node.label), style="reverse bold"))
details_pane.update(f"Selected: {match_node.label}")
# Remove the match_selector after selection
try:
option_list = self.query_one("#match_selector", OptionList)
option_list.remove()
except:
pass
+45
View File
@@ -0,0 +1,45 @@
from textual.containers import Vertical
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Button, Static
class ThemeSelector(Widget):
"""Widget for selecting and applying Textual themes."""
class ThemeSelected(Message):
"""Message posted when a theme is selected."""
def __init__(self, theme_name: str):
super().__init__()
self.theme_name = theme_name
AVAILABLE_THEMES = [
("textual-dark", "textual-dark"),
("textual-light", "textual-light"),
("nord", "nord"),
("gruvbox", "gruvbox"),
("catppuccin-mocha", "catppuccin-mocha"),
("dracula", "dracula"),
("tokyo-night", "tokyo-night"),
("monokai", "monokai"),
("flexoki", "flexoki"),
("catppuccin-latte", "catppuccin-latte"),
("solarized-light", "solarized-light"),
]
def compose(self):
yield Static("Theme Options", id="theme_title")
with Vertical() as column:
column.styles.width = "1fr"
column.styles.height = "auto"
for label, btn_id in self.AVAILABLE_THEMES:
yield Button(label, id=f"set_theme_{btn_id}", compact=True)
def on_button_pressed(self, event: Button.Pressed) -> None:
button_id = event.button.id
if button_id and button_id.startswith("set_theme_"):
theme_name = button_id.replace("set_theme_", "")
self.post_message(self.ThemeSelected(theme_name))