diff --git a/.gitea/workflows/loxide.yml b/.gitea/workflows/loxide.yml index ca7e614..95e4ac3 100644 --- a/.gitea/workflows/loxide.yml +++ b/.gitea/workflows/loxide.yml @@ -1,50 +1,52 @@ -name: Build EXE and Release -run-name: ${{ gitea.actor }} -on: - push: - branches: - - master - -jobs: - Build and Release: - runs-on: debian-stable - env: - DISPLAY: :99 - - steps: - - name: Install Prerequisites - run: | - dpkg --add-architecture i386 - apt update > /dev/null 2>&1 - apt install git curl wine32:i386 xvfb -y > /dev/null 2>&1 - - - name: Start X Virtual Framebuffer (Xvfb) - run: | - # Start Xvfb in the background using the defined display number - Xvfb :99 -screen 0 1024x768x16 & - - - name: Cloning Repository - run: | - git clone https://brotoskyj:${{ secrets.RUNNER_TOKEN }}@git.racooncity.org/brotoskyj/AirlockTools --branch RustImplementation - pwd - ls - - - name: Setting Up Build Environment - run: | - 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 - wine python.exe /quiet /NoWeb InstallAllUsers=1 PrependPath=1 TargetDir=C:/Python313 - cp -r AirlockTools/ ~/.wine/drive_c/Python313/ - cd ~/.wine/drive_c/Python313 - wine python.exe -m pip install nuitka pywin32 keyring --break-system-packages - wine python.exe -m pip install -r AirlockTools/requirements.txt - wine python.exe -m pip install \ - --index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ \ - airlock-libs --break-system-packages - - - name: Build Executable - env: - DISPLAY: :99 - WINEDEBUG: -all - run: | - 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 +##name: Build EXE and Release +##run-name: ${{ gitea.actor }} +##on: +## push: +## branches: +## - master +## +##jobs: +## Build and Release: +## runs-on: debian-stable +## env: +## DISPLAY: :99 +## +## steps: +## - name: Install Prerequisites +## run: | +## dpkg --add-architecture i386 +## apt update > /dev/null 2>&1 +## apt install git curl wine32:i386 xvfb -y > /dev/null 2>&1 +## +## - name: Start X Virtual Framebuffer (Xvfb) +## run: | +## # Start Xvfb in the background using the defined display number +## Xvfb :99 -screen 0 1024x768x16 & +## +## - name: Cloning Repository +## run: | +## git clone https://brotoskyj:${{ secrets.RUNNER_TOKEN }}@git.racooncity.org/brotoskyj/AirlockTools --branch RustImplementation +## pwd +## ls +## +## - name: Setting Up Build Environment +## run: | +## 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 +## wget https://download.visualstudio.microsoft.com/download/pr/0c8b0c6f-3d30-4d06-98d8-3a7a42f3e78a/64c0b0f3b8b2d0e11d9a2303d5e39a22/ndp472-devpack-ENU.exe -O dotnet472.exe +## wget https://aka.ms/vs/16/release/vs_buildtools.exe -O vstools2019.exe +## wine dotnet472.exe /q /norestart /ChainingPackage ADMINDEPLOYMENT /log dotnet472.log +## wine python.exe /quiet /NoWeb InstallAllUsers=1 PrependPath=1 TargetDir=C:/Python313 +## wine vstools2019.exe --quiet --wait --norestart --nocache --installPath "C:\\BuildTools" --add Microsoft.VisualStudio.Workload.VCTools --includeRecommended --acceptEula +## cp -r AirlockTools/ ~/.wine/drive_c/Python313/ +## cd ~/.wine/drive_c/Python313 +## 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: +## DISPLAY: :99 +## WINEDEBUG: -all +## 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 +## \ No newline at end of file diff --git a/.gitea/workflows/loxide_lib.yml b/.gitea/workflows/loxide_lib.yml index c991875..0a21747 100644 --- a/.gitea/workflows/loxide_lib.yml +++ b/.gitea/workflows/loxide_lib.yml @@ -2,6 +2,8 @@ name: Build Library run-name: ${{ gitea.actor }} on: push: + branches: + - RustImplementation paths: - airlock_libs/** diff --git a/AirlockTools_Client.py b/AirlockTools_Client.py index c2d4093..47c9562 100644 --- a/AirlockTools_Client.py +++ b/AirlockTools_Client.py @@ -14,76 +14,71 @@ # along with this program. If not, see . -#TODO Continue implementing logger -#TODO Add input sanitation and CSV injection prevention -#TODO Continue OTP and Local approval rewrites -#TODO Explore pywin32 -#TODO Fix Requirements.txt -#TODO Create Generic system_config.json for gitea +# TODO Continue implementing logger +# TODO Add input sanitation and CSV injection prevention +# TODO Continue OTP and Local approval rewrites +# TODO Explore pywin32 +# TODO Fix Requirements.txt +# TODO Create Generic system_config.json for gitea import logging import os import tempfile + import dotenv import urllib3 from services.API import AirlockAPIWrapper from services.security import getAPI 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 -urllib3.disable_warnings( - urllib3.exceptions.InsecureRequestWarning -) +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + def main(): - - + if "NUITKA_ONEFILE_PARENT" in os.environ: splash_filename = os.path.join( 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): os.unlink(splash_filename) - - 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() base_dir = get_base_directory() logger = logging.getLogger(__name__) dotenv.load_dotenv(dotenv_path=base_dir / ".env") try: - url = os.getenv("URL") - username = os.getenv("USERNAME") + url = os.getenv("URL") + username = os.getenv("USERNAME") - if not url: - raise ValueError("Missing URL in environment variables.") - if not username: - raise ValueError("Missing USERNAME in environment variables.") + if not url: + raise ValueError("Missing URL in environment variables.") + if not username: + raise ValueError("Missing USERNAME in environment variables.") - logger.debug(f"Retrieved URL: {url}") - logger.debug(f"Retrieved Username: {username}") + logger.debug(f"Retrieved URL: {url}") + logger.debug(f"Retrieved Username: {username}") except ValueError as e: - logger.error(f"Configuration error: {e}", exc_info=True) - raise + logger.error(f"Configuration error: {e}", exc_info=True) + raise - - api_key = getAPI(username, "AirlockTools") + api_key = getAPI(username, "Loxide") if api_key is None: - raise ValueError("API key for AirlockTools is missing.") + raise ValueError("API key for Loxide is missing.") api = AirlockAPIWrapper( base_url=str(os.getenv("URL")), api_key=api_key, ) - run_AirlockTools(api) - + run_Loxide(api) if __name__ == "__main__": diff --git a/AirlockTools_Server.py b/AirlockTools_Server.py index 56a58a0..183635c 100644 --- a/AirlockTools_Server.py +++ b/AirlockTools_Server.py @@ -14,12 +14,11 @@ # along with this program. If not, see . - -#TODO Add CSV injection prevention -#TODO Continue OTP and Local approval rewrites -#TODO Explore pywin32 -#TODO Fix Requirements.txt -#TODO Create Generic system_config.json for gitea +# TODO Add CSV injection prevention +# TODO Continue OTP and Local approval rewrites +# TODO Explore pywin32 +# TODO Fix Requirements.txt +# TODO Create Generic system_config.json for gitea import logging @@ -29,62 +28,66 @@ import dotenv import urllib3 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.policyhandler import updateAuditPoliciesFromEnforcementPolices from services.security import getAPI from utils.setup import setup -urllib3.disable_warnings( - urllib3.exceptions.InsecureRequestWarning -) +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + def main(): - #Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored - + # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored + working_dir = setup() - + logger = logging.getLogger(__name__) dotenv.load_dotenv(dotenv_path=working_dir / ".env") try: - url = os.getenv("URL") - username = os.getenv("USERNAME") + url = os.getenv("URL") + username = os.getenv("USERNAME") - if not url: - raise ValueError("Missing URL in environment variables.") - if not username: - raise ValueError("Missing USERNAME in environment variables.") + if not url: + raise ValueError("Missing URL in environment variables.") + if not username: + raise ValueError("Missing USERNAME in environment variables.") + + logger.debug(f"Retrieved URL: {url}") + logger.debug(f"Retrieved Username: {username}") - logger.debug(f"Retrieved URL: {url}") - logger.debug(f"Retrieved Username: {username}") - except ValueError as e: - logger.error(f"Configuration error: {e}", exc_info=True) - raise - + logger.error(f"Configuration error: {e}", exc_info=True) + raise + api = AirlockAPIWrapper( - base_url=str(os.getenv("URL")), - api_key = getAPI(username, "AirlockTools"), - ) - + base_url=str(os.getenv("URL")), + api_key=getAPI(username, "AirlockTools"), + ) logger.info("Running non-interactively to start monitoring Airlock Changes") - register_function("monitorLA", la.scheduleAddingLAHashes) 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("updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api]) + recurring_job( + "updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api] + ) else: reload_jobs() - + start_scheduler() - + + if __name__ == "__main__": main() diff --git a/README.md b/README.md index d647cdc..8257c1c 100644 --- a/README.md +++ b/README.md @@ -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 -**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. diff --git a/Server/scheduler_async.py b/Server/scheduler_async.py index 6f90d36..87e8c80 100644 --- a/Server/scheduler_async.py +++ b/Server/scheduler_async.py @@ -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 JOBS_FILE = os.path.join(os.getcwd(), "jobs.json") + def register_function(name: str, func: Callable): """ Register a function so it can be called by name later. @@ -38,6 +39,7 @@ def register_function(name: str, func: Callable): """ FUNCTION_MAP[name] = func + def load_jobs() -> List[Dict[str, Any]]: """ 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: return json.load(f) + def save_jobs(jobs: List[Dict[str, Any]]): """ 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: json.dump(jobs, f, indent=4) + def cancel_job(job_id: str): """ Cancel a scheduled job by ID and remove it from the registry and persistence. @@ -66,7 +70,15 @@ def cancel_job(job_id: str): jobs = [j for j in load_jobs() if j.get("id") != job_id] save_jobs(jobs) -def run_once_job(job_id: str, func_name: str, delay_seconds: float, args=None, kwargs=None, persist=True): + +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). """ @@ -87,18 +99,25 @@ def run_once_job(job_id: str, func_name: str, delay_seconds: float, args=None, k if persist: jobs = [j for j in load_jobs() if j.get("id") != job_id] - jobs.append({ - "id": job_id, - "type": "once", - "delay": delay_seconds, - "function": func_name, - "args": args, - "kwargs": kwargs - }) + jobs.append( + { + "id": job_id, + "type": "once", + "delay": delay_seconds, + "function": func_name, + "args": args, + "kwargs": kwargs, + } + ) save_jobs(jobs) - logger.info(f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds.") + 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. """ @@ -121,17 +140,20 @@ def recurring_job(job_id: str, func_name: str, interval: float, args=None, kwarg if persist: jobs = [j for j in load_jobs() if j.get("id") != job_id] - jobs.append({ - "id": job_id, - "type": "recurring", - "interval": interval, - "function": func_name, - "args": args, - "kwargs": kwargs - }) + jobs.append( + { + "id": job_id, + "type": "recurring", + "interval": interval, + "function": func_name, + "args": args, + "kwargs": kwargs, + } + ) save_jobs(jobs) logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.") + def reload_jobs(): """ Reload jobs from JSON and reschedule them. @@ -145,7 +167,7 @@ def reload_jobs(): job["delay"], job.get("args"), job.get("kwargs"), - persist=False + persist=False, ) elif job["type"] == "recurring": recurring_job( @@ -154,9 +176,10 @@ def reload_jobs(): job["interval"], job.get("args"), job.get("kwargs"), - persist=False + persist=False, ) + async def start_scheduler(): """ Start the asynchronous scheduler loop. @@ -189,4 +212,4 @@ async def start_scheduler(): while True: await asyncio.sleep(3600) # Sleep indefinitely; jobs run via call_later except asyncio.CancelledError: - logger.critical("Scheduler stopped.") \ No newline at end of file + logger.critical("Scheduler stopped.") diff --git a/airlock_libs/Cargo.lock b/airlock_libs/Cargo.lock index 70327a0..a3d9b52 100644 --- a/airlock_libs/Cargo.lock +++ b/airlock_libs/Cargo.lock @@ -17,7 +17,7 @@ dependencies = [ [[package]] name = "airlock_libs" -version = "1.0.3" +version = "2.0.0" dependencies = [ "chrono", "indicatif", diff --git a/airlock_libs/Cargo.toml b/airlock_libs/Cargo.toml index 15f8786..3a853d9 100644 --- a/airlock_libs/Cargo.toml +++ b/airlock_libs/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "airlock_libs" -version = "1.0.3" +version = "2.0.0" edition = "2024" [lib] diff --git a/airlock_libs/airlock_libs.pyi b/airlock_libs/airlock_libs.pyi index 2e061f1..137dc92 100644 --- a/airlock_libs/airlock_libs.pyi +++ b/airlock_libs/airlock_libs.pyi @@ -1,6 +1,9 @@ -from typing import Dict, List, Optional -def pull_policy_exec_histories(self, type: List[str], checkpoint: str, policy: List[str]) -> str: - """Retrieve execution history logs.""" +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 api(AirlockAPIWrapper): """ @@ -22,66 +25,66 @@ def api(AirlockAPIWrapper): """ def history_logging( - api, - exec_types: str, - checkpoint_number: str, - policy_names: str, - ) -> List[Dict[str, Any]]: - """ - Query execution history logs from the Airlock API. + api, + exec_types: str, + checkpoint_number: str, + policy_names: str, +) -> List[Dict[str, Any]]: + """ + Query execution history logs from the Airlock API. - Parameters - ---------- - exec_types : str - A JSON-style string list of execution types to retrieve. - Example: "[3,5,8]" - - 0 = Trusted Execution - - 1 = Blocked Execution - - 2 = Untrusted Execution [Audit] - - 3 = Untrusted Execution [OTP] - - 5 = Trusted Publisher Execution - - 8 = Trusted Process Execution - (etc.) + Parameters + ---------- + exec_types : str + A JSON-style string list of execution types to retrieve. + Example: "[3,5,8]" + - 0 = Trusted Execution + - 1 = Blocked Execution + - 2 = Untrusted Execution [Audit] + - 3 = Untrusted Execution [OTP] + - 5 = Trusted Publisher Execution + - 8 = Trusted Process Execution + (etc.) - checkpoint_number : str - The checkpoint ID. Used to fetch results after a certain event. - Example: "601d275487bacb01e3470713" + checkpoint_number : str + The checkpoint ID. Used to fetch results after a certain event. + Example: "601d275487bacb01e3470713" - policy_names : str - A comma-separated or JSON-style list of policy group names. - Example: "Apple Mac" or "["Apple Mac", "Servers London"]" + policy_names : str + A comma-separated or JSON-style list of policy group names. + Example: "Apple Mac" or "["Apple Mac", "Servers London"]" - Returns - ------- - List[Dict[str, Any]] - A list of dictionaries, where each dictionary represents an - execution history record. Each record can include fields like: + Returns + ------- + List[Dict[str, Any]] + A list of dictionaries, where each dictionary represents an + execution history record. Each record can include fields like: - - checkpoint: str - - type: int - - username: str - - hostname: str - - filename: str - - ppolicy: str - - policyname: str - - policyver: str - - commandline: str - - publisher: str - - pprocess: str - - gprocess: str - - sha256: str - - datetime: str - - ip: str - - localip: str - Raises - ------ - RuntimeError - If the request fails or the response cannot be parsed. + - checkpoint: str + - type: int + - username: str + - hostname: str + - filename: str + - ppolicy: str + - policyname: str + - policyver: str + - commandline: str + - publisher: str + - pprocess: str + - gprocess: str + - sha256: str + - datetime: str + - ip: str + - localip: str + Raises + ------ + RuntimeError + If the request fails or the response cannot be parsed. - Example - ------- - >>> histories = await airlock_libs.history_logging("[3,5,8]", "601d275487bacb01e3470713", "Apple Mac") - >>> print(histories[0]["filename"]) - 'chrome.exe' - """ - ... \ No newline at end of file + Example + ------- + >>> histories = await airlock_libs.history_logging("[3,5,8]", "601d275487bacb01e3470713", "Apple Mac") + >>> print(histories[0]["filename"]) + 'chrome.exe' + """ + ... diff --git a/airlock_libs/pyproject.toml b/airlock_libs/pyproject.toml index 2f78abd..e7cd8aa 100644 --- a/airlock_libs/pyproject.toml +++ b/airlock_libs/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "airlock_libs" -version = "1.0.3" +version = "2.0.0" description = "Airlock Digital API Wrapper" readme = "README.md" license = { text = "AGPL-3.0-only" } diff --git a/airlock_libs/src/services.rs b/airlock_libs/src/services.rs index 5056344..8870d5c 100644 --- a/airlock_libs/src/services.rs +++ b/airlock_libs/src/services.rs @@ -13,7 +13,7 @@ use std::{ env, fmt::Write, fs::{self, File}, - io::Read, + io::{Read, Seek, SeekFrom}, path::PathBuf, str::FromStr, }; @@ -96,7 +96,9 @@ pub fn pull_policy_exec_histories( let client = build_client(py, &py_self); let api: Py = py_self; let cutoff = Local::now().naive_local() - Duration::days(days); + let mut f = File::open(&writeable_filepath).unwrap(); loop { + f.seek(SeekFrom::Start(0)).unwrap(); let execution_histories = history_logging( py, &api, @@ -110,7 +112,6 @@ pub fn pull_policy_exec_histories( break; } 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(); f.read_to_string(&mut contents).unwrap(); 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 percentage_diff = - ((days + 10) - date_diff.num_days()) as f64 / (days + 10) as f64 * 100.0; + let percentage_diff = (days - date_diff.num_days()) as f64 / days as f64 * 100.0; progress_bar.set_position(percentage_diff.round() as u64); progress_bar.set_message("Total Percent Complete"); } diff --git a/flows/localApproval.py b/flows/localApproval.py index 789b917..5cc69a2 100644 --- a/flows/localApproval.py +++ b/flows/localApproval.py @@ -41,7 +41,9 @@ def getLocalApprovals(api: AirlockAPIWrapper): result = api.otp_find_awaiting() local_approval = pd.DataFrame(result["response"]["otpusage"]) if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"): - previous_run = pd.read_parquet(f"{base_dir}\\cache\\newest_local_approval.parquet") + previous_run = pd.read_parquet( + f"{base_dir}\\cache\\newest_local_approval.parquet" + ) previous_run.to_parquet( f"{base_dir}\\cache\\last_local_approval.parquet", index=False ) @@ -66,10 +68,10 @@ def getLocalApprovals(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", "[]") 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: register_function("add_hash", returnFromLocalApproval) @@ -93,7 +95,9 @@ def scheduleAddingLAHashes(api: AirlockAPIWrapper): duration_minutes = int(batch_df["duration"].iloc[0]) start_time = datetime.datetime.now() run_time = start_time + datetime.timedelta(minutes=duration_minutes) - early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95)) + early_time = start_time + datetime.timedelta( + minutes=np.floor(duration_minutes * 0.95) + ) early_timestamp = early_time.timestamp() run_timestamp = run_time.timestamp() @@ -148,7 +152,13 @@ def scheduleAddingLAHashes(api: AirlockAPIWrapper): 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 @@ -166,11 +176,14 @@ def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publish #TODO finish logic for adding hashes """ 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", "[]") pups = load_env_json("PUPS", "[]") threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE") - print(f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}") + print( + f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}" + ) + def moveToLocalApproval(api: AirlockAPIWrapper): possible_durations = [15, 60, 360, 1440, 10080] @@ -213,10 +226,9 @@ def moveToLocalApproval(api: AirlockAPIWrapper): def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid): - + purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}" api.otp_generate(agentid, duration_selected, purpose) - def monitorAuditStatus(api: AirlockAPIWrapper): @@ -224,11 +236,13 @@ def monitorAuditStatus(api: AirlockAPIWrapper): last_agents = [] if not last_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_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 last_agent_map = {agent.hostname: agent for agent in last_agents} @@ -261,8 +275,8 @@ def monitorAuditStatus(api: AirlockAPIWrapper): def getNewLocalApprovals(api: AirlockAPIWrapper): - - working_dir = load_env("WORKING_DIR") + + working_dir = load_env("WORKING_DIR") current_la = getLocalApprovals(api) # Load old approval list @@ -273,15 +287,21 @@ def getNewLocalApprovals(api: AirlockAPIWrapper): old_la = pd.DataFrame(columns=current_la.columns) # 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) # Find new entries new_entries = current_la[~current_la["key"].isin(old_la["key"])] # Convert 'granted' to datetime and filter by last 10 minutes - new_entries["granted"] = pd.to_datetime(new_entries["granted"], utc=True, errors="coerce") - ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10) + new_entries["granted"] = pd.to_datetime( + new_entries["granted"], utc=True, errors="coerce" + ) + ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( + minutes=10 + ) recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago] # Save current approvals for next run diff --git a/flows/otp.py b/flows/otp.py index 6c09db2..8b15fab 100644 --- a/flows/otp.py +++ b/flows/otp.py @@ -14,7 +14,6 @@ # along with this program. If not, see . - from datetime import datetime import logging import os @@ -33,19 +32,24 @@ logger = logging.getLogger(__name__) def otp_generate(api: AirlockAPIWrapper): otp_dict = {} 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: print(agent.hostname) confirm = Selector.confirm() if agents and confirm: requester = get_sanitized_input("Who is requesting the OTP: ") because = get_sanitized_input("Why/What work are they doing?: ") - + purpose = f"Requester: {requester} - for : {because}" possible_durations = [15, 60, 360, 1440, 10080] print(colorText("Please select a duration in minutes: ", "white")) - print(colorText("15 mins, 60 mins, 360 mins(6 Hours), 1440 mins (24 Hours), 10080 mins (7 Days):", "white")) + print( + colorText( + "15 mins, 60 mins, 360 mins(6 Hours), 1440 mins (24 Hours), 10080 mins (7 Days):", + "white", + ) + ) duration_selected = Selector.select_int(possible_durations) if isinstance(duration_selected, list): @@ -60,56 +64,68 @@ def otp_generate(api: AirlockAPIWrapper): print(colorText("Requested Codes:", "green")) 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): activeagents = api.otp_find_active() awaitingagents = api.otp_find_awaiting() enforcedagents = api.otp_find_enforced() revokedagents = api.otp_find_revoked() - - + # Add a 'status' column to each DataFrame - activeagents['status'] = 'active' - awaitingagents['status'] = 'awaiting' - enforcedagents['status'] = 'enforced' - revokedagents['status'] = 'revoked' + activeagents["status"] = "active" + awaitingagents["status"] = "awaiting" + enforcedagents["status"] = "enforced" + revokedagents["status"] = "revoked" # Combine all into one DataFrame - combined_agents = pd.concat([activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True) - combined_agents = combined_agents.sort_values(by='otpid', ascending=False) + combined_agents = pd.concat( + [activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True + ) + combined_agents = combined_agents.sort_values(by="otpid", ascending=False) - #Optionally, select specific hosts - user_input = get_sanitized_input("\nWould you like to search for a specific device? (y/n): ").strip().lower() - if user_input == 'y': + # Optionally, select specific hosts + user_input = ( + get_sanitized_input("\nWould you like to search for a specific device? (y/n): ") + .strip() + .lower() + ) + if user_input == "y": agentnames = [] agents = selectAgents(api) for agent in agents: agentnames.append(agent.hostname) - combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)] + combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)] - #Present and select rows + # Present and select rows selected_rows = Selector.select_dataframe_with_mode( combined_agents, - columns=['otpid', 'hostname', 'status','purpose','granted'], - header="OTP Sessions" + columns=["otpid", "hostname", "status", "purpose", "granted"], + header="OTP Sessions", ) combined_df = pd.DataFrame() for row in selected_rows: - otpid = row['otpid'] - hostname = row['hostname'] + otpid = row["otpid"] + hostname = row["hostname"] result = api.otp_get_activities(otpid) - result['hostname'] = hostname + result["hostname"] = hostname if not result.empty: logger.info(f"Activities for {hostname} (otpid: {otpid}):\n{result}") combined_df = pd.concat([combined_df, result], ignore_index=True) else: logger.info(f"No activities found for {hostname} (otpid: {otpid})") - user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower() - if user_input == 'y': + user_input = ( + get_sanitized_input( + "\nWould you like to export the results to a CSV file? (y/n): " + ) + .strip() + .lower() + ) + if user_input == "y": working_dir = load_env("WORKING_DIR") timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") filename = f"otp_activities_{timestamp}.csv" @@ -128,43 +144,44 @@ def otp_activities_by_agent(api: AirlockAPIWrapper): logging.debug("User declined to export the DataFrame.") - def otp_revoke(api: AirlockAPIWrapper): activeagents = api.otp_find_active() awaitingagents = api.otp_find_awaiting() - - activeagents['status'] = 'active' - awaitingagents['status'] = 'awaiting' + + activeagents["status"] = "active" + awaitingagents["status"] = "awaiting" combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True) - combined_agents = combined_agents.sort_values(by='otpid', ascending=False) + 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 = combined_agents.sort_values(by='otpid', ascending=False) + combined_agents = combined_agents.sort_values(by="otpid", ascending=False) - #Optionally, select specific hosts - user_input = get_sanitized_input("\nWould you like to search for a specific device? (y/n): ").strip().lower() - if user_input == 'y': + # Optionally, select specific hosts + user_input = ( + get_sanitized_input("\nWould you like to search for a specific device? (y/n): ") + .strip() + .lower() + ) + if user_input == "y": agentnames = [] agents = selectAgents(api) for agent in agents: agentnames.append(agent.hostname) - combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)] + combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)] - #Present and select rows + # Present and select rows selected_rows = Selector.select_dataframe_with_mode( combined_agents, - columns=['otpid', 'hostname', 'status','purpose','granted'], - header="OTP Sessions" + columns=["otpid", "hostname", "status", "purpose", "granted"], + header="OTP Sessions", ) for row in selected_rows: - otpid = row['otpid'] - hostname = row['hostname'] + otpid = row["otpid"] + hostname = row["hostname"] result = api.otp_revoke(otpid) logger.info(f"{hostname} (otpid: {otpid}):\n{result}") - - diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py index f66eee7..6f4243c 100644 --- a/flows/prepPolicy.py +++ b/flows/prepPolicy.py @@ -44,7 +44,6 @@ logger = logging.getLogger(__name__) dotenv.load_dotenv() - def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]: policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()] @@ -60,9 +59,18 @@ def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]: return selected if isinstance(selected, list) else [selected] -def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) -> List[Allowlist]: - if policy == "all": allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()] - else: allowlists = [Allowlist(**row.to_dict()) for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()] +def selectAllowlists( + api: AirlockAPIWrapper, policy=all, allow_multiple=True +) -> List[Allowlist]: + if policy == "all": + allowlists = [ + Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows() + ] + else: + allowlists = [ + Allowlist(**row.to_dict()) + for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows() + ] logger.debug("Prompting for Allowlist(s)") print(colorText("Please select allowlist(s)", "white")) selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True) @@ -76,9 +84,7 @@ def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) def sortHashes( - api: AirlockAPIWrapper, - selected_policies: List[Policy], - type=[1, 2, 6, 7] + api: AirlockAPIWrapper, selected_policies: List[Policy], type=[1, 2, 6, 7] ): working_dir = load_env("WORKING_DIR") history_days = Selector.select_value( @@ -86,34 +92,41 @@ def sortHashes( value_type=int, valid_range=(1, 150), ) - + logger.debug(f"{history_days} day selected for history") - + if history_days is None: logging.warning("No history range selected. Aborting.") return - + policy_executions = ExecutionHistoryRecord.from_policies( api, selected_policies, type_=type, history_days=history_days ) logger.debug(f"Executions contains {policy_executions}") - enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(api, policy_executions) - categorized_executions = ExecutionHistoryRecord.categorize_executions_by_hash_decision(enriched_executions) - approved, unapproved, needs_review, unknown = ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions) + enriched_executions = ExecutionHistoryRecord.enrich_with_hashes( + api, policy_executions + ) + categorized_executions = ( + ExecutionHistoryRecord.categorize_executions_by_hash_decision( + enriched_executions + ) + ) + approved, unapproved, needs_review, unknown = ( + ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions) + ) categories = { - "needs_review": needs_review, - "approved": approved, - "unapproved": unapproved, - "leftover" : unknown - } + "needs_review": needs_review, + "approved": approved, + "unapproved": unapproved, + "leftover": unknown, + } - for label, records in categories.items(): 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" 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]) # Optional: flatten hash_obj if needed - if not df.empty and 'hash_obj' in df.columns: - 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) + if not df.empty and "hash_obj" in df.columns: + 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) # Save to CSV df.to_csv(csv_path, index=False) @@ -140,9 +153,11 @@ def buildPathsandPublishers(selected_policies: List[Policy], split): df1 = pd.DataFrame() df2 = 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" - 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): df1 = pd.read_csv(path1) @@ -163,37 +178,48 @@ def buildPathsandPublishers(selected_policies: List[Policy], split): if "filename" in all_approved_hashes.columns: all_approved_hashes = all_approved_hashes.sort_values(by="filename") 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: primary_path_exclusions = calculatePath( - all_approved_hashes, path_exclusion_constant, + all_approved_hashes, + path_exclusion_constant, split, ) remaining_hashes = all_approved_hashes[ ~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"]) ] secondary_path_exclusions = calculatePath( - remaining_hashes,(path_exclusion_constant - 1), split + remaining_hashes, (path_exclusion_constant - 1), split ) remaining_hashes = remaining_hashes[ ~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"]) ] dataframes = { - "all_approved_hashes" : all_approved_hashes, + "all_approved_hashes": all_approved_hashes, "primary_Paths": primary_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") for name, df in dataframes.items(): logger.debug(f" DataFrame headers: {list(df.columns)}") - if "hashes" in name : df.sort_values(by="filename", inplace=True) - else: df.sort_values(by="longestcfp", inplace=True) - - 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 "hashes" in name: + df.sort_values(by="filename", inplace=True) + else: + df.sort_values(by="longestcfp", inplace=True) + + 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: # 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" ].drop_duplicates(subset=["publisher"]) # 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[["publisher"]] 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) - else: + publist.to_csv( + f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv", + index=False, + ) + else: logger.debug("Approved Hashes list appears empty") + def buildPreflights(selected_policies: List[Policy]): 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" 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): df1 = pd.read_csv(path1) else: @@ -239,19 +268,18 @@ def buildPreflights(selected_policies: List[Policy]): approved_paths = pd.DataFrame() else: 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): 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: - logger.warning(f"File not found: {hash}") - + logger.warning(f"File not found: {hash}") if os.path.exists(publishers): approved_publishers = pd.read_csv(publishers) @@ -259,19 +287,33 @@ def buildPreflights(selected_policies: List[Policy]): else: 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(): logger.debug(f" DataFrame headers: {list(df.columns)}") - if name == "approved_paths":df.sort_values(by="longestcfp", inplace=True) - elif name == "approved_hashes":df.sort_values(by="filename", inplace=True) - elif name == "approved_publishers" : 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") + if name == "approved_paths": + df.sort_values(by="longestcfp", inplace=True) + elif name == "approved_hashes": + df.sort_values(by="filename", inplace=True) + elif name == "approved_publishers": + 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"): - 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): if not isinstance(path, (str, bytes, os.PathLike)): @@ -281,7 +323,9 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): return parts # 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: print(f"[WARNING] Non-string entries found in column '{col}':") print(non_string_entries) @@ -290,10 +334,14 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): split_paths = df[col].apply(clean_split) 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] - 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") new_rows = [] @@ -317,7 +365,7 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): for i, parts in enumerate(split_parts): filename = parts[-1] middle = ( - os.sep.join(parts[len(common_prefix):-1]) + os.sep.join(parts[len(common_prefix) : -1]) if len(parts) > len(common_prefix) + 1 else "" ) @@ -330,6 +378,7 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): return pd.DataFrame(new_rows).drop(columns=["group_key"]) + def calculatePath(approved_hashes, path_exclusion_constant, split): if split: 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] 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 = [] @@ -364,7 +413,9 @@ def calculatePath(approved_hashes, path_exclusion_constant, split): ] 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"] @@ -380,51 +431,64 @@ def calculatePath(approved_hashes, path_exclusion_constant, split): return pathExclusions + 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(destination_policy) + logger.info("These path exclusions would be added to:") + logger.info(destination_policy) - pathexclusions = pd.read_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") + pathexclusions = pd.read_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]:\\") - processed_paths = [ - (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}" - for path, ext in unique_combinations.itertuples(index=False, name=None) - ] + drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\") + processed_paths = [ + (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}" + for path, ext in unique_combinations.itertuples(index=False, name=None) + ] - for path in processed_paths: - logger.info(path) + for path in processed_paths: + logger.info(path) - print(colorText("These publishers would added", "yellow")) - processed_publishers = [] - if os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"): - publishers = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv") - if publishers.empty: - print(colorText("The publishers list is empty.", "red")) - else: - processed_publishers = ( - publishers[publishers["publisher"] != "Not Signed"] - ["publisher"] - .drop_duplicates() - .tolist() - ) - for publisher in processed_publishers: - print(publisher) + print(colorText("These publishers would added", "yellow")) + processed_publishers = [] + if os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv" + ): + publishers = pd.read_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv" + ) + if publishers.empty: + print(colorText("The publishers list is empty.", "red")) + else: + processed_publishers = ( + publishers[publishers["publisher"] != "Not Signed"]["publisher"] + .drop_duplicates() + .tolist() + ) + for publisher in processed_publishers: + print(publisher) - print(colorText("These hashes would be added to:", "yellow")) - print(destination_allowlist) + print(colorText("These hashes would be added to:", "yellow")) + print(destination_allowlist) - processed_hashes = hashes["sha256"].unique().tolist() - print_x_wide(processed_hashes, 3) + processed_hashes = hashes["sha256"].unique().tolist() + print_x_wide(processed_hashes, 3) - return processed_paths, processed_hashes, processed_publishers - -def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 into functions + return processed_paths, processed_hashes, processed_publishers + + +def menu_policy_enforce( + api: AirlockAPIWrapper, +): # TODO Need to clean up 6 and 7 into functions selected_policies = [] destination_policy = [] 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") while True: - printEnforceChecklist(selected_policies, destination_policy, destination_allowlist) + printEnforceChecklist( + selected_policies, destination_policy, destination_allowlist + ) choice = get_sanitized_input("\nEnter your choice: ") if choice == "1": clear_screen() - selected_policies = selectPolicies(api,True) + selected_policies = selectPolicies(api, True) elif choice == "2": 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) @@ -461,35 +531,53 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 elif choice == "4": 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) 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": 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" ): buildPreflights(selected_policies) 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": clear_screen() if ( - os.path.exists(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") + os.path.exists( + 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_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: # Log which condition(s) failed 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") - 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") if not destination_policy: 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" ): print(colorText("Proceeding with the code...", "yellow")) - api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes) - api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths) + api.hash_add_to_allowlist( + destination_allowlist[0].applicationid, processed_hashes + ) + api.policy_add_path_exclusions( + destination_policy[0].groupid, processed_paths + ) if processed_publishers: - api.policy_add_publishers(destination_policy[0].groupid, processed_publishers) + api.policy_add_publishers( + destination_policy[0].groupid, processed_publishers + ) locked() - + else: logger.error("Confirmation block failed. Reasons:") 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: logger.error(" - `destination_allowlist` is missing or invalid.") if confirmation.strip() != "I AGREE": - logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip()) + logger.error( + " - User did not confirm with 'I AGREE'. Received: '%s'", + confirmation.strip(), + ) elif choice.upper() == "F": open_directory(working_dir) elif choice.upper() == "B": break - else: print(colorText("Invalid choice. Please try again.", "red")) + def section_header(title): - print(colorText("\n --------------------------------------------------------------------", "cyan")) + print( + colorText( + "\n --------------------------------------------------------------------", + "cyan", + ) + ) print(colorText(f" ------------- {title} -------------", "cyan")) - print(colorText(" --------------------------------------------------------------------", "cyan")) + print( + colorText( + " --------------------------------------------------------------------", + "cyan", + ) + ) def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist): working_dir = load_env("WORKING_DIR") 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 - 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: print(colorText(" [✗] No policies have been chosen", "red")) else: @@ -561,67 +677,194 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all print(colorText(f" [✓] {policy.name}", "green")) # 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: - 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: print(colorText(" [✗] No destination policy has been chosen", "red")) 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: print(colorText(" [✗] No allowlist has been chosen", "red")) # 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: policy_id = selected_policies[0].name 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: - 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 print(colorText("4. Manually review the files:", "cyan")) - print(colorText(" 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")) + print( + colorText( + " 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: policy_id = selected_policies[0].name 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" - 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")) + second_review_path = ( + f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv" + ) + 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: - 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 - print(colorText(f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\", "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")) + print( + colorText( + f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\", + "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: policy_id = selected_policies[0].name reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv" preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.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")) - 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")) + 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", + ) + ) + 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: - print(colorText(" [✗] No policies selected, cannot check preflight status", "red")) + print( + colorText( + " [✗] No policies selected, cannot check preflight status", "red" + ) + ) # Final Steps - print(colorText("6. Test ------------------------------------------------------", "cyan")) - print(colorText(" Prints to console the changes that would be made, must be done to proceed. ", "cyan")) + print( + 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(colorText(" Apply path exclusions and approved publishers to selected policy", "cyan")) + print( + colorText( + "7. Liftoff ------------------------------------------------------", "cyan" + ) + ) + print( + colorText( + " Apply path exclusions and approved publishers to selected policy", + "cyan", + ) + ) print(colorText(" Apply approved hashes to allowlist", "cyan")) - # Utility Options print(colorText("F. 📂 - Open Working Directory", "cyan")) print(colorText("B. 🔚 - Back", "cyan")) diff --git a/flows/quietAgent.py b/flows/quietAgent.py index 6bc05ea..d723da7 100644 --- a/flows/quietAgent.py +++ b/flows/quietAgent.py @@ -51,20 +51,22 @@ def findQuietAgents(api: AirlockAPIWrapper): 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 if confirm: 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: - 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") return - # Convert 'datetime' column to timezone-aware datetime objects policy_exec_history["datetime"] = pd.to_datetime( 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() # 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 - most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates( - subset="hostname", keep="first" - ) + most_recent_exec = policy_exec_history.sort_values( + by="days_ago" + ).drop_duplicates(subset="hostname", keep="first") # Map most recent execution age to agents agents["days_since"] = agents["hostname"].map( @@ -101,7 +105,9 @@ def findQuietAgents(api: AirlockAPIWrapper): ) # 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 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 # Print results - - + message = ( f"Total agents: {total_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}%" ) logger.debug(message) - colorText(message,"green") + colorText(message, "green") get_sanitized_input("Press enter to continue") diff --git a/loading.png b/loading.png index 1dc81bd..a10998c 100644 Binary files a/loading.png and b/loading.png differ diff --git a/models/agent.py b/models/agent.py index 9a4c580..ee6678c 100644 --- a/models/agent.py +++ b/models/agent.py @@ -21,40 +21,36 @@ from models.policy import Policy @dataclass class Agent: + hostname: str agentid: str clientversion: str domain: str freespace: int - groupid: str # Changed to str to match UUID-style IDs - hostname: str + groupid: str ip: str localip: str lastcheckin: str os: str policyversion: str - status: int # raw status code + status: int username: str groupname: Optional[str] = field(default=None) status_text: Optional[str] = field(default=None) # Class-level status map - status_map: ClassVar[dict] = { - 0: "Offline", - 1: "Online", - 2: "Hidden", - 3: "Safemode" - } - + status_map: ClassVar[dict] = {0: "Offline", 1: "Online", 2: "Hidden", 3: "Safemode"} def enrich_with_policies(self, policies: List[Policy]): - """Enrich the agent with groupname and human-readable status.""" - self.status_text = self.status_map.get(self.status, "Unknown") - for policy in policies: - if policy.groupid == self.groupid: - self.groupname = policy.name - break - if not self.groupname: - self.groupname = "Unknown" + """Enrich the agent with groupname and human-readable status.""" + self.status_text = self.status_map.get(self.status, "Unknown") + for policy in policies: + if policy.groupid == self.groupid: + self.groupname = policy.name + break + if not self.groupname: + self.groupname = "Unknown" + + """ from models.agent import Agent diff --git a/models/execution.py b/models/execution.py index 13281d1..e1cf514 100644 --- a/models/execution.py +++ b/models/execution.py @@ -28,7 +28,6 @@ import pandas as pd import airlock_libs from services.API import AirlockAPIWrapper -from services.policyhandler import pullPolicyExechistories from utils.configmanager import get_protected_value, load_env_json from utils.utils import colorText, regulator @@ -36,11 +35,13 @@ logger = logging.getLogger(__name__) dotenv.load_dotenv() + @dataclass class Hash: """ Hash model representing Hash data """ + sha256: str applications: str baselines: str @@ -62,7 +63,7 @@ class Hash: sha384: str sha512: str at_decision: Optional[str] = None - + def to_dict(self): return asdict(self) @@ -100,11 +101,15 @@ class Hash: for hash_obj in hashes: publisher = hash_obj.publisher or "" description = hash_obj.description or "" - reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} + reputation = ( + hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} + ) scannermatch = reputation.get("scannermatch") logger.debug(f"Evaluating hash: {hash_obj}") - logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}") + logger.debug( + f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}" + ) # 1. Unapproved: bad publisher or PUP if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): @@ -128,9 +133,9 @@ class Hash: # 3. Approved or Unapproved based on threat level try: - score = int(scannermatch) # pyright: ignore[reportArgumentType] + score = int(scannermatch) # pyright: ignore[reportArgumentType] 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.") hash_obj.at_decision = "unapproved" unapproved_count += 1 @@ -139,15 +144,17 @@ class Hash: hash_obj.at_decision = "approved" approved_count += 1 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" needs_review_count += 1 - - logger.debug(f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}") + logger.debug( + f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}" + ) return hashes - @classmethod def export_to_csv(cls, hash_list, directory_path): """ @@ -204,11 +211,12 @@ class ExecutionHistoryRecord: localip: Optional[str] = None extid: 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 hash_obj: Optional[Hash] = None - @classmethod def from_dict(cls, data: dict): mandatory_fields = [ @@ -225,7 +233,9 @@ class ExecutionHistoryRecord: "datetime", ] 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: raise ValueError(f"Missing mandatory fields: {missing_fields}") @@ -255,7 +265,7 @@ class ExecutionHistoryRecord: extname=data.get("extname"), exttype=data.get("exttype"), extbrowser=data.get("extbrowser"), - hash_obj=data.get("hash_obj") + hash_obj=data.get("hash_obj"), ) @classmethod @@ -264,7 +274,9 @@ class ExecutionHistoryRecord: ) -> List["ExecutionHistoryRecord"]: executions = [] 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: data = json.loads(execs) exechistories = data.get("response", {}).get("exechistories", []) @@ -275,8 +287,12 @@ class ExecutionHistoryRecord: df = df.drop_duplicates(subset=["sha256", "filename", "hostname"]) df = df.sort_values(by=["sha256", "filename"]) - executions.extend([cls.from_dict(row.to_dict()) for _, row in df.iterrows()]) - logger.debug(f"Staging of Execution history for policy: {policy.name} is complete") + executions.extend( + [cls.from_dict(row.to_dict()) for _, row in df.iterrows()] + ) + logger.debug( + f"Staging of Execution history for policy: {policy.name} is complete" + ) print( colorText( f"Staging of Execution history for policy: {policy.name} is complete", @@ -285,20 +301,23 @@ class ExecutionHistoryRecord: ) return executions - + @staticmethod def enrich_with_hashes( - api: AirlockAPIWrapper, - executions: List["ExecutionHistoryRecord"] + api: AirlockAPIWrapper, executions: List["ExecutionHistoryRecord"] ) -> List["ExecutionHistoryRecord"]: """ 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}) - 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: - logger.warning("No sha256 values found in execution records. Skipping enrichment.") + logger.warning( + "No sha256 values found in execution records. Skipping enrichment." + ) return executions logger.debug("Querying hash data from API...") @@ -307,8 +326,10 @@ class ExecutionHistoryRecord: hash_objects = [] required_fields = { - f.name for f in dataclasses.fields(Hash) - if f.default == dataclasses.MISSING and f.default_factory == dataclasses.MISSING + f.name + 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()): @@ -346,11 +367,15 @@ class ExecutionHistoryRecord: exec_record.hash_obj = hash_obj 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 @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. @@ -374,11 +399,15 @@ class ExecutionHistoryRecord: publisher = hash_obj.publisher 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") 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 if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): @@ -402,7 +431,7 @@ class ExecutionHistoryRecord: # 3. Approved or Unapproved based on threat level try: - score = int(scannermatch) # pyright: ignore[reportArgumentType] + score = int(scannermatch) # pyright: ignore[reportArgumentType] logger.debug(f"Parsed scannermatch score: {score}") if threat_tolerance is not None and score >= threat_tolerance: logger.debug("Unapproved: Unsigned file with high threat score.") @@ -413,7 +442,9 @@ class ExecutionHistoryRecord: hash_obj.at_decision = "approved" approved_count += 1 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" needs_review_count += 1 @@ -423,12 +454,14 @@ class ExecutionHistoryRecord: ) return executions - - + @classmethod - def sort_by_hash_decision( - cls, executions: List["ExecutionHistoryRecord"] - ) -> Tuple[List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"]]: + def sort_by_hash_decision(cls, executions: List["ExecutionHistoryRecord"]) -> Tuple[ + List["ExecutionHistoryRecord"], + List["ExecutionHistoryRecord"], + List["ExecutionHistoryRecord"], + List["ExecutionHistoryRecord"], + ]: """ Sorts ExecutionHistoryRecord objects into approved, unapproved, needs_review, and unknown groups based on the value of hash_obj.at_decision. @@ -454,7 +487,9 @@ class ExecutionHistoryRecord: else: 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" Unapproved: {len(unapproved)}") logger.info(f" Needs Review: {len(needs_review)}") @@ -463,7 +498,6 @@ class ExecutionHistoryRecord: return approved, unapproved, needs_review, unknown - """ executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30) diff --git a/models/policy.py b/models/policy.py index 7040f32..dbfb357 100644 --- a/models/policy.py +++ b/models/policy.py @@ -29,7 +29,9 @@ class Policy: def __repr__(self): # 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"" def to_dict(self): @@ -53,7 +55,9 @@ class Allowlist: def __repr__(self): # 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"" def to_dict(self): diff --git a/requirements.txt b/requirements.txt index e932ee2..65ee4c7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,14 +1,14 @@ -cryptography==46.0.1 -keyring==25.6.0 -numpy==2.3.2 -pandas==2.3.1 -python-dotenv==1.1.1 -pymongo -requests==2.32.5 -schedule==1.2.2 -tqdm==4.67.1 -urllib3==2.5.0 -bson==0.5.10 - ---extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ -airlock_libs==1.0.3 \ No newline at end of file +cryptography==46.0.3 +keyring==25.6.0 +numpy==2.3.4 +pandas==2.3.3 +pymongo==4.15.3 +python-dotenv==1.2.1 +Requests==2.32.5 +textual==6.5.0 +tqdm==4.67.1 +urllib3==2.5.0 +pyperclip==1.11.0 + +--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ +airlock_libs==2.0.0 \ No newline at end of file diff --git a/screens/otpworkflowscreen.py b/screens/otpworkflowscreen.py new file mode 100644 index 0000000..fa5afc1 --- /dev/null +++ b/screens/otpworkflowscreen.py @@ -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 diff --git a/services/API.py b/services/API.py index e5a3c76..34ce788 100644 --- a/services/API.py +++ b/services/API.py @@ -23,7 +23,6 @@ import requests logger = logging.getLogger(__name__) - class AirlockAPIWrapper: """ A wrapper class for interacting with the Airlock API. @@ -141,25 +140,25 @@ class AirlockAPIWrapper: payload = {"status": "0"} result = self._post("/v1/otp/usage", payload) return pd.DataFrame(result["response"]["otpusage"]) - + def otp_find_enforced(self) -> pd.DataFrame: """Find OTPs that are awaiting activation.""" payload = {"status": "2"} result = self._post("/v1/otp/usage", payload) return pd.DataFrame(result["response"]["otpusage"]) - + def otp_find_revoked(self) -> pd.DataFrame: """Find OTPs that are awaiting activation.""" payload = {"status": "3"} result = self._post("/v1/otp/usage", payload) return pd.DataFrame(result["response"]["otpusage"]) - + def otp_find_by_agent(self, agentid) -> pd.DataFrame: """Find OTP by agent.""" payload = {"agentid": agentid} result = self._post("/v1/otp/usage", payload) return pd.DataFrame(result["response"]["otpusage"]) - + def otp_generate(self, agentid: str, duration: int, purpose: str) -> str: """Generate a new OTP for an agent.""" payload = { @@ -175,7 +174,7 @@ class AirlockAPIWrapper: payload = {"otpid": otpid} result = self._post("/v1/otp/activities", payload) return pd.DataFrame(result["response"]["otpactivities"]) - + def otp_revoke(self, otpid: str) -> dict: """ Revoke an active OTP. @@ -186,18 +185,17 @@ class AirlockAPIWrapper: """ payload = {"otpid": otpid} return self._post("/v1/otp/revoke", payload) - - def otp_validate(self, otpcode: str) -> dict: - """ - Validate an OTP code. - Parameters: - - otpcode (str): The OTP code to validate. - Returns: - - dict: JSON response indicating validity. - """ - payload = {"otpcode": otpcode} - return self._post("/v1/otp/validate", payload) + def otp_validate(self, otpcode: str) -> dict: + """ + Validate an OTP code. + Parameters: + - otpcode (str): The OTP code to validate. + Returns: + - dict: JSON response indicating validity. + """ + payload = {"otpcode": otpcode} + return self._post("/v1/otp/validate", payload) # Policy Management def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict: @@ -225,7 +223,7 @@ class AirlockAPIWrapper: payload = {"groupid": groupid} result = self._post("/v1/group/agents", payload) return pd.DataFrame(result["response"]["agents"]) - + def policy_list_allowlists(self, groupid: str) -> pd.DataFrame: """List allowlists assigned to a specific policy group.""" payload = {"groupid": groupid} @@ -236,31 +234,37 @@ class AirlockAPIWrapper: """Set audit mode for a policy group. 1=Audit, 0=Enforcement""" payload = {"groupid": groupid, "auditmode": auditmode} return self._post("/v1/group/settings/auditmode", payload) - - def policy_set_script_custom(self, - groupid: str, - script_custom: int, - scripts_audit: List[str], - scripts_disabled: List[str], - scripts_respect: List[str], - ) -> dict: + + def policy_set_script_custom( + self, + groupid: str, + script_custom: int, + scripts_audit: List[str], + scripts_disabled: List[str], + scripts_respect: List[str], + ) -> dict: """Set audit mode for a policy group. 1=Audit, 0=Enforcement""" - payload = {"groupid": groupid, - "script_custom": script_custom, - "scripts_audit": scripts_audit, - "scripts_disabled": scripts_disabled, - "scripts_respect": scripts_respect - } + payload = { + "groupid": groupid, + "script_custom": script_custom, + "scripts_audit": scripts_audit, + "scripts_disabled": scripts_disabled, + "scripts_respect": scripts_respect, + } return self._post("/v1/group/settings/script_custom", payload) # Execution History - def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> str: + def history_logging( + self, type: List[str], checkpoint: str, policy: List[str] + ) -> str: """Retrieve execution history logs.""" payload = {"type": type, "checkpoint": checkpoint, "policy": policy} result = self._post("/v1/logging/exechistories", payload) return result["response"]["exechistories"] - def history_execution(self, today: str, date_selected: str, agent_name: str) -> List[Dict]: + def history_execution( + self, today: str, date_selected: str, agent_name: str + ) -> List[Dict]: """ Retrieve execution history logs. diff --git a/services/agenthandler.py b/services/agenthandler.py index 531fc41..0ec3d3b 100644 --- a/services/agenthandler.py +++ b/services/agenthandler.py @@ -47,7 +47,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool): print(colorText("No agents selected or invalid history range.", "red")) return - historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d") + historical_date = (datetime.now() - timedelta(days=history_days)).strftime( + "%Y-%m-%d" + ) today = datetime.now().strftime("%Y-%m-%d") all_history = [] @@ -56,7 +58,11 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool): try: exechistory = api.history_execution(today, historical_date, agent.hostname) except Exception as e: - print(colorText(f"❌ Error retrieving history for {agent.hostname}: {e}", "red")) + print( + colorText( + f"❌ Error retrieving history for {agent.hostname}: {e}", "red" + ) + ) continue if isinstance(exechistory, list): @@ -76,7 +82,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool): print(colorText(f"{key}: {value}", "green")) print("\n") else: - print(colorText(f"No execution history found for {agent.hostname}.", "yellow")) + print( + colorText(f"No execution history found for {agent.hostname}.", "yellow") + ) if outputjson: print(json.dumps(all_history, indent=2)) @@ -92,6 +100,7 @@ def findAllAgents(api): return agents + def findAgents(api, return_dataframe): agents = selectAgents(api) working_dir = load_env("WORKING_DIR") @@ -113,8 +122,14 @@ def findAgents(api, return_dataframe): print(agent_df) logging.debug("Displayed DataFrame to console.") - user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower() - if user_input == 'y': + user_input = ( + get_sanitized_input( + "\nWould you like to export the results to a CSV file? (y/n): " + ) + .strip() + .lower() + ) + if user_input == "y": timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") filename = f"agentsearch_{timestamp}.csv" file_path = os.path.join(str(working_dir), filename) @@ -131,17 +146,27 @@ def findAgents(api, return_dataframe): else: logging.debug("User declined to export the DataFrame.") + def collect_device_names() -> List[str]: print(colorText("🔍 Device Search", "cyan")) - print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan")) - print(colorText("When you're done, press Enter twice (Three times if you have a single device).\n", "cyan")) + print( + colorText( + "Enter the device hostnames you'd like to search for, one per line.", "cyan" + ) + ) + print( + colorText( + "When you're done, press Enter twice (Three times if you have a single device).\n", + "cyan", + ) + ) print(colorText("Example:", "cyan")) print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan")) print(colorText("Paste or type your device names below:", "white")) device_input_lines = [] empty_line_count = 0 - valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$') + valid_line_pattern = re.compile(r"^[a-zA-Z0-9_\- ]+$") while True: line = get_sanitized_input("") @@ -158,7 +183,12 @@ def collect_device_names() -> List[str]: if valid_line_pattern.match(stripped_line): device_input_lines.append(stripped_line) else: - print(colorText(f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow")) + print( + colorText( + f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", + "yellow", + ) + ) return [name for name in device_input_lines if name] @@ -168,10 +198,13 @@ def choose_match_type() -> bool: return get_sanitized_input("").strip().lower() in ["y", "yes"] -def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool) -> List['Agent']: +def match_agents( + device_names: List[str], agents: List["Agent"], use_exact: bool +) -> List["Agent"]: if use_exact: return [ - agent for agent in agents + agent + for agent in agents if agent.hostname.lower() in [name.lower() for name in device_names] ] else: @@ -180,23 +213,38 @@ def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool return [agent for agent in agents if regex.search(agent.hostname)] -def show_unmatched(device_names: List[str], matched_agents: List['Agent'], use_exact: bool): +def show_unmatched( + device_names: List[str], matched_agents: List["Agent"], use_exact: bool +): if use_exact: - unmatched = [name for name in device_names if not any(agent.hostname.lower() == name.lower() for agent in matched_agents)] + unmatched = [ + name + for name in device_names + if not any( + agent.hostname.lower() == name.lower() for agent in matched_agents + ) + ] else: - unmatched = [name for name in device_names if not any(re.search(re.escape(name), agent.hostname, re.IGNORECASE) for agent in matched_agents)] + unmatched = [ + name + for name in device_names + if not any( + re.search(re.escape(name), agent.hostname, re.IGNORECASE) + for agent in matched_agents + ) + ] if unmatched: logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}") print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow")) -def enrich_agents(agents: List['Agent'], policies: List['Policy']): +def enrich_agents(agents: List["Agent"], policies: List["Policy"]): for agent in agents: agent.enrich_with_policies(policies) -def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']: +def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]: device_names = collect_device_names() if not device_names: logger.debug("No device names entered") @@ -227,11 +275,11 @@ def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']: if idx < len(matched_agents): line += f"{matched_agents[idx].hostname:<30}" logger.info(line) - + matched_agents = Selector.select_with_mode( matched_agents, label_func=lambda agent: agent.hostname, - header="Matched Devices:" + header="Matched Devices:", ) if not matched_agents: @@ -263,11 +311,17 @@ def moveAgentToRelatedPolicy( if agent.groupid in policy_relationship_map: target_policy = policy_relationship_map[agent.groupid] elif agent.groupid in policy_relationship_map.values(): - logger.debug(f"Agent {agent.hostname} is already in an audit group. No action needed.") - print(f"Agent {agent.hostname} is already in an audit group. No action needed.") + logger.debug( + f"Agent {agent.hostname} is already in an audit group. No action needed." + ) + print( + f"Agent {agent.hostname} is already in an audit group. No action needed." + ) return else: - logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.") + logger.warning( + f"Error: No corresponding audit policy found for groupid: {agent.groupid}." + ) return elif mode == "enforcement": @@ -275,10 +329,14 @@ def moveAgentToRelatedPolicy( if agent.groupid in inverse_map: target_policy = inverse_map[agent.groupid] elif agent.groupid in inverse_map.values(): - logger.info(f"Agent {agent.hostname} is already in an enforcement group. No action needed.") + logger.info( + f"Agent {agent.hostname} is already in an enforcement group. No action needed." + ) return else: - logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.") + logger.warning( + f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}." + ) return else: @@ -299,25 +357,32 @@ def toggleEnforcement(api: AirlockAPIWrapper): devices = selectAgents(api) for device in devices: print(device.hostname) - confirm = Selector.confirm("Would you like to continue with these devices? Y/N: ") + confirm = Selector.confirm( + "Would you like to continue with these devices? Y/N: " + ) if direction and devices and confirm: for device in devices: - result = moveAgentToRelatedPolicy(api,device, str(direction).lower()) + result = moveAgentToRelatedPolicy(api, device, str(direction).lower()) logger.info(f"{device.hostname}: result: {result}") get_sanitized_input("Press enter to continue") + def moveAgents(api: AirlockAPIWrapper): devices = selectAgents(api) for device in devices: print(device.hostname) - confirm_devices = Selector.confirm("Would you like to continue with these devices? Y/N: ") + confirm_devices = Selector.confirm( + "Would you like to continue with these devices? Y/N: " + ) if devices and confirm_devices: policies = selectPolicies(api, False) - confirm_move = Selector.confirm(f"Would you like to move these devices to {policies[0].name}?") + confirm_move = Selector.confirm( + f"Would you like to move these devices to {policies[0].name}?" + ) if confirm_move: for device in devices: result = api.agent_move(device.agentid, policies[0].groupid) logger.info(f"{device.hostname}: result: {result}") else: logger.info("Exiting without change") - get_sanitized_input("Press enter to continue") \ No newline at end of file + get_sanitized_input("Press enter to continue") diff --git a/services/policyhandler.py b/services/policyhandler.py index 5acc422..c32cdfc 100644 --- a/services/policyhandler.py +++ b/services/policyhandler.py @@ -34,7 +34,6 @@ from utils.utils import areYouSure, colorText, get_sanitized_input logger = logging.getLogger(__name__) - def pullPolicyExechistories( api: AirlockAPIWrapper, policy: Policy, @@ -72,7 +71,7 @@ def pullPolicyExechistories( ) as pbar: while True: histories = api.history_logging( - type=type, checkpoint=checkpoint, policy= [policy.name] + type=type, checkpoint=checkpoint, policy=[policy.name] ) # Ensure histories is a list of dictionaries @@ -98,13 +97,17 @@ def pullPolicyExechistories( # Update checkpoint on last item if index == len(histories) - 1: - checkpoint = history_item["checkpoint"] # pyright: ignore[reportArgumentType] + checkpoint = history_item[ + "checkpoint" + ] # pyright: ignore[reportArgumentType] filebar.desc = f"Checkpoint Progress: {checkpoint}" break try: history_date = datetime.datetime.strptime( - history_item["datetime"].replace(" +0000 UTC", ""), # pyright: ignore[reportArgumentType] + history_item["datetime"].replace( + " +0000 UTC", "" + ), # pyright: ignore[reportArgumentType] "%Y-%m-%dT%H:%M:%SZ", ).date() except ValueError: @@ -177,8 +180,10 @@ def pullPolicyExechistories( def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days): + import airlock_libs + 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: data = json.loads(exehist) executionhist_policy = pd.DataFrame(data["response"]["exechistories"]) @@ -203,7 +208,7 @@ def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days): executionhist_policy = executionhist_policy.sort_values( by=["sha256", "filename"] ) - logger.debug( f"Staging of Execution history for policy: {policy} is complete") + logger.debug(f"Staging of Execution history for policy: {policy} is complete") print( colorText( f"Staging of Execution history for policy: {policy} is complete", @@ -235,10 +240,10 @@ def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper): for enforcement_policy, audit_policy in policy_relationship_map.items(): api.policy_clone(enforcement_policy, audit_policy) api.policy_set_auditmode(audit_policy, "1") - + def confirmUpdateAfromE(api: AirlockAPIWrapper): areYouSure() confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") if confirmation.strip() == "I AGREE": - updateAuditPoliciesFromEnforcementPolices(api) \ No newline at end of file + updateAuditPoliciesFromEnforcementPolices(api) diff --git a/services/security.py b/services/security.py index aa95929..70229e2 100644 --- a/services/security.py +++ b/services/security.py @@ -28,8 +28,8 @@ import keyring # Constants KDF_ITERATIONS = 200_000 -SALT_SIZE = 16 # 128-bit Salt -NONCE_SIZE = 12 # AES-GCM +SALT_SIZE = 16 # 128-bit Salt +NONCE_SIZE = 12 # AES-GCM KEY_SIZE = 32 # AES-256 logger = logging.getLogger(__name__) @@ -49,9 +49,11 @@ def configure_keyring_backend(): system = platform.system() if system == "Windows": import keyring.backends.Windows + keyring.set_keyring(keyring.backends.Windows.WinVaultKeyring()) elif system == "Linux": import keyring.backends.kwallet + keyring.set_keyring(keyring.backends.kwallet.DBusKeyring()) else: raise EnvironmentError(f"Unsupported OS: {system}") @@ -68,8 +70,9 @@ def store_api_key(service: str, username: str, api_key: str, password: str): b64 = base64.b64encode(blob).decode() keyring.set_password(service, username, b64) - - logger.debug(f"API key for service '{service}' and user '{username}' stored successfully.") + logger.debug( + f"API key for service '{service}' and user '{username}' stored successfully." + ) print("\n✅ API key stored securely.") print("The program will now exit. Press Enter to continue...") @@ -90,8 +93,8 @@ def retrieve_api_key(service: str, username: str, password: str) -> str: raise ValueError("No stored secret for this service/username.") blob = base64.b64decode(b64) salt = blob[:SALT_SIZE] - nonce = blob[SALT_SIZE:SALT_SIZE + NONCE_SIZE] - ct = blob[SALT_SIZE + NONCE_SIZE:] + nonce = blob[SALT_SIZE : SALT_SIZE + NONCE_SIZE] + ct = blob[SALT_SIZE + NONCE_SIZE :] key = _derive_key(password.encode(), salt) aesgcm = AESGCM(key) pt = aesgcm.decrypt(nonce, ct, associated_data=None) @@ -124,7 +127,9 @@ def getAPI(USERNAME, SERVICE_NAME): if api_key_exists(SERVICE_NAME, USERNAME): for attempt in range(1, 4): - password = getpass(f"Attempt {attempt}/3 - Enter password to unlock your API key: ") + password = getpass( + f"Attempt {attempt}/3 - Enter password to unlock your API key: " + ) try: apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password) logging.debug("API key successfully retrieved.") @@ -134,9 +139,15 @@ def getAPI(USERNAME, SERVICE_NAME): logging.error("Failed to retrieve API key after 3 incorrect attempts.") raise ValueError("Failed to retrieve API key after 3 incorrect attempts.") else: - logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.") - api_key = getpass(f"No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip() - print("Please exit and relaunch program after saving your credential to avoid errors") + logging.warning( + f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'." + ) + api_key = getpass( + f"No API key found. Please enter your API key for '{SERVICE_NAME}': " + ).strip() + print( + "Please exit and relaunch program after saving your credential to avoid errors" + ) while True: password = getpass("Create a password to encrypt your API key: ") @@ -155,7 +166,9 @@ def getAPI(USERNAME, SERVICE_NAME): logging.error(f"Failed to store API key: {e}") break else: - logging.warning("Password does not meet complexity requirements. Try again.") + logging.warning( + "Password does not meet complexity requirements. Try again." + ) class APIKeyManager: @@ -169,4 +182,4 @@ class APIKeyManager: def get(cls) -> str: if cls._api_key is None: raise ValueError("API key not loaded. Call APIKeyManager.load() first.") - return cls._api_key \ No newline at end of file + return cls._api_key diff --git a/utils/configmanager.py b/utils/configmanager.py index 62c2bdf..b099fa3 100644 --- a/utils/configmanager.py +++ b/utils/configmanager.py @@ -29,14 +29,15 @@ PROTECTED_KEYS = [ "PATH_EXCLUSION_CONST", "MIN_FILES_FOR_PATH", "VT_THREAT_TOLERANCE", - "POLICY_MAP_ENF_AUD" + "POLICY_MAP_ENF_AUD", ] _protected_config = {} + def get_system_config_path() -> Path: # Check inside bundled EXE directory first - bundled_dir = Path(getattr(sys, '_MEIPASS', '')) + bundled_dir = Path(getattr(sys, "_MEIPASS", "")) bundled_path = bundled_dir / "system_config.json" if bundled_path.exists(): return bundled_path @@ -44,6 +45,7 @@ def get_system_config_path() -> Path: # Fallback to external location return Path(__file__).parent.parent / "system_config.json" + def load_protected_config() -> dict: global _protected_config try: @@ -52,19 +54,20 @@ def load_protected_config() -> dict: except FileNotFoundError: logging.warning("⚠️ system_config.json not found. Using built-in defaults.") system_config = { - "APPNAME": "AirlockTools", + "APPNAME": "Loxide", "PATH_EXCLUSION_CONST": 4, "MIN_FILES_FOR_PATH": 4, "VT_THREAT_TOLERANCE": 4, - "POLICY_MAP_ENF_AUD": { - "enforced_id": "audit_id" - } + "POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"}, } _protected_config = {key: system_config[key] for key in PROTECTED_KEYS} return _protected_config -def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]: + +def get_protected_value( + key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None +) -> Optional[T]: value = _protected_config.get(key) if value is None: logging.warning(f"Protected config key '{key}' not found.") @@ -74,9 +77,12 @@ def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: value = value.strip("'\"") return cast_type(value) except (ValueError, TypeError): - logging.warning(f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}.") + logging.warning( + f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}." + ) return default + def get_protected_json(key: str, default: str = "{}") -> dict: raw = _protected_config.get(key, default) if isinstance(raw, dict): @@ -85,13 +91,11 @@ def get_protected_json(key: str, default: str = "{}") -> dict: return json.loads(raw) except json.JSONDecodeError: try: - escaped = raw.encode('unicode_escape').decode('utf-8') + escaped = raw.encode("unicode_escape").decode("utf-8") return json.loads(escaped) except Exception as e: logging.error(f"Failed to parse protected JSON key '{key}': {e}") return json.loads(default) - - def load_env_json(key: str, default: str): @@ -100,13 +104,16 @@ def load_env_json(key: str, default: str): return json.loads(raw) except json.JSONDecodeError: try: - escaped = raw.encode('unicode_escape').decode('utf-8') + escaped = raw.encode("unicode_escape").decode("utf-8") return json.loads(escaped) except Exception as e: logging.error(f"Failed to parse {key}: {e}") return json.loads(default) -def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]: + +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. @@ -126,5 +133,7 @@ def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] value = value.strip("'\"") # Strip surrounding quotes return cast_type(value) except (ValueError, TypeError): - logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.") - return default \ No newline at end of file + logger.warning( + f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}." + ) + return default diff --git a/utils/selector.py b/utils/selector.py index 714e197..0a0f498 100644 --- a/utils/selector.py +++ b/utils/selector.py @@ -21,9 +21,12 @@ from utils.utils import colorText, get_sanitized_input logger = logging.getLogger(__name__) + class Selector: @staticmethod - def _get_sorted_items(items: List[Any], label_func: Callable[[Any], str]) -> List[Any]: + def _get_sorted_items( + items: List[Any], label_func: Callable[[Any], str] + ) -> List[Any]: return sorted(items, key=lambda item: label_func(item).lower()) @staticmethod @@ -31,10 +34,10 @@ class Selector: items: List[Any], label_func: Callable[[Any], str], num_columns: int = 4, - header: str = "Available Choices:" + header: str = "Available Choices:", ) -> None: # Force single column if items are DataFrame rows - + if items and isinstance(items[0], (pd.Series, dict)): num_columns = 1 @@ -51,9 +54,7 @@ class Selector: @staticmethod def _display_selected_items( - selected: List[Any], - label_func: Callable[[Any], str], - num_columns: int = 4 + selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 4 ) -> None: print(colorText("\nCurrent selections:", "cyan")) if not selected: @@ -92,7 +93,7 @@ class Selector: allow_multiple: bool = False, prompt_each: bool = False, header: str = "Available Choices:", - num_columns: int = 4 + num_columns: int = 4, ) -> Union[Optional[Any], List[Any]]: if not items: logger.warning("No items available for selection.") @@ -104,9 +105,19 @@ class Selector: if allow_multiple: while True: - Selector._display_choices(remaining_items, label_func, num_columns=num_columns, header=header) - Selector._display_selected_items(selected, label_func, num_columns=num_columns) - choice = get_sanitized_input("Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: ").strip().lower() + Selector._display_choices( + remaining_items, label_func, num_columns=num_columns, header=header + ) + Selector._display_selected_items( + selected, label_func, num_columns=num_columns + ) + choice = ( + get_sanitized_input( + "Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: " + ) + .strip() + .lower() + ) if choice == "q": break elif choice == "r": @@ -125,10 +136,14 @@ class Selector: logger.info(f"Selected: {label_func(item)}") else: 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 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: choice = int(get_sanitized_input("Select one item by number: ")) if 1 <= choice <= len(full_sorted_items): @@ -145,9 +160,14 @@ class Selector: def select_with_mode( items: List[Any], label_func: Callable[[Any], str], - header: str = "Available Choices:" + header: str = "Available Choices:", ) -> 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() if mode == "a": return items @@ -156,7 +176,7 @@ class Selector: label_func=label_func, allow_multiple=True, prompt_each=False, - header=header + header=header, ) if not selected: return items @@ -172,44 +192,38 @@ class Selector: @staticmethod def select_objects( - objects: List[Any], - allow_multiple: bool = False, - prompt_each: bool = False + objects: List[Any], allow_multiple: bool = False, prompt_each: bool = False ) -> Union[Optional[Any], List[Any]]: return Selector._select_from_list( objects, label_func=lambda obj: getattr(obj, "name", str(obj)), allow_multiple=allow_multiple, prompt_each=prompt_each, - header="Available Objects:" + header="Available Objects:", ) @staticmethod def select_string( - options: List[str], - allow_multiple: bool = False, - prompt_each: bool = False + options: List[str], allow_multiple: bool = False, prompt_each: bool = False ) -> Union[Optional[str], List[str]]: return Selector._select_from_list( options, label_func=str, allow_multiple=allow_multiple, prompt_each=prompt_each, - header="Available Options:" + header="Available Options:", ) @staticmethod def select_int( - options: List[int], - allow_multiple: bool = False, - prompt_each: bool = False + options: List[int], allow_multiple: bool = False, prompt_each: bool = False ) -> Union[Optional[int], List[int]]: return Selector._select_from_list( options, label_func=lambda x: str(x), allow_multiple=allow_multiple, prompt_each=prompt_each, - header="Available Integers:" + header="Available Integers:", ) @staticmethod @@ -217,7 +231,7 @@ class Selector: prompt: str, value_type: type = int, valid_range: Optional[tuple] = None, - allow_quit: bool = False + allow_quit: bool = False, ) -> Optional[Any]: while True: user_input = get_sanitized_input(prompt).strip().lower() @@ -255,7 +269,7 @@ class Selector: columns: Optional[List[str]] = None, allow_multiple: bool = False, prompt_each: bool = False, - header: str = "Available Rows:" + header: str = "Available Rows:", ) -> List[pd.Series]: if df.empty: print("DataFrame is empty.") @@ -272,7 +286,7 @@ class Selector: label_func=label_func, allow_multiple=allow_multiple, prompt_each=prompt_each, - header=header + header=header, ) if isinstance(result, pd.Series): @@ -286,7 +300,7 @@ class Selector: def select_dataframe_with_mode( df: pd.DataFrame, columns: Optional[List[str]] = None, - header: str = "Available Rows:" + header: str = "Available Rows:", ) -> List[pd.Series]: if df.empty: print("⚠️ DataFrame is empty.") @@ -305,7 +319,12 @@ class Selector: print(f"{i}: {label_func(row)}") # Prompt for mode once - print(colorText("\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white")) + print( + colorText( + "\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", + "white", + ) + ) mode = get_sanitized_input("").strip().lower() if mode == "a": @@ -317,7 +336,7 @@ class Selector: label_func=label_func, allow_multiple=True, prompt_each=False, - header=header + header=header, ) if not selected: @@ -331,4 +350,4 @@ class Selector: return [pd.Series(row) for row in items if row not in selected] else: print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow")) - return [] \ No newline at end of file + return [] diff --git a/utils/setup.py b/utils/setup.py index f7ba7f2..4277389 100644 --- a/utils/setup.py +++ b/utils/setup.py @@ -30,16 +30,16 @@ from utils.configmanager import PROTECTED_KEYS, load_protected_config def get_base_directory() -> Path: system = platform.system() home = Path.home() - if system == 'Windows': - return Path(os.getenv('APPDATA', home / 'AppData' / 'Roaming')) / "AirlockTools" - elif system == 'Darwin': - return home / 'Library' / 'Application Support' / "AirlockTools" + if system == "Windows": + return Path(os.getenv("APPDATA", home / "AppData" / "Roaming")) / "Loxide" + elif system == "Darwin": + return home / "Library" / "Application Support" / "Loxide" else: - return home / '.local' / 'share' / "AirlockTools" + return home / ".local" / "share" / "Loxide" def configure_logging(log_dir: Path, log_level: str = "DEBUG"): - log_file = log_dir / "airlocktools.log" + log_file = log_dir / "Loxide.log" config = { "version": 1, # Required key for dictConfig format version @@ -58,17 +58,17 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"): "file": { "class": "logging.handlers.TimedRotatingFileHandler", "filename": str(log_file), - "when": "midnight", # Rotate logs at midnight - "interval": 1, # Every 1 day - "backupCount": 7, # Keep 7 days of logs - "encoding": "utf-8", # Ensure UTF-8 encoding - "level": "DEBUG", # Always log DEBUG and above - "formatter": "detailed", # Use detailed format + "when": "midnight", # Rotate logs at midnight + "interval": 1, # Every 1 day + "backupCount": 7, # Keep 7 days of logs + "encoding": "utf-8", # Ensure UTF-8 encoding + "level": "DEBUG", # Always log DEBUG and above + "formatter": "detailed", # Use detailed format }, "console": { "class": "logging.StreamHandler", "level": log_level.upper(), # Configurable log level - "formatter": "simple", # Use simple format + "formatter": "simple", # Use simple format }, }, "root": { @@ -82,9 +82,9 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"): try: config["handlers"]["eventlog"] = { "class": "logging.handlers.NTEventLogHandler", - "appname": "AirlockTools", # Event log source name - "level": "CRITICAL", # Only log critical errors - "formatter": "simple", # Use simple format + "appname": "Loxide", # Event log source name + "level": "CRITICAL", # Only log critical errors + "formatter": "simple", # Use simple format } config["root"]["handlers"].append("eventlog") except Exception as e: @@ -94,9 +94,11 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"): logging.config.dictConfig(config) logging.getLogger().debug("✅ Logging configured.") - + 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" @@ -108,45 +110,45 @@ def load_system_config() -> dict: except FileNotFoundError: logging.warning("⚠️ system_config.json not found. Using built-in defaults.") return { - "APPNAME": "AirlockTools", + "APPNAME": "Loxide", "LOG_LEVEL": "DEBUG", "PATH_EXCLUSION_CONST": 4, "MIN_FILES_FOR_PATH": 4, "VT_THREAT_TOLERANCE": 4, - "POLICY_MAP_ENF_AUD": { - "enforced_id": "audit_id" - } + "POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"}, } + def load_user_config(config_dir: Path) -> dict: user_config_path = config_dir / "user_config.json" if not user_config_path.exists(): - default_user_config = { - "URL": "", - "LOG_LEVEL": "INFO" - } + default_user_config = {"URL": "", "LOG_LEVEL": "INFO"} with open(user_config_path, "w") as f: json.dump(default_user_config, f, indent=4) logging.debug(f"Created user config at {user_config_path}") with open(user_config_path, "r") as f: return json.load(f) + def write_config_to_env(config: dict, env_path: Path): for key, value in config.items(): if key in PROTECTED_KEYS: continue # Skip protected keys try: - serialized = json.dumps(value) if isinstance(value, (list, dict)) else str(value) + serialized = ( + json.dumps(value) if isinstance(value, (list, dict)) else str(value) + ) set_key(env_path, key, serialized) except Exception as e: logging.warning(f"Failed to write {key} to .env: {e}") + def setup(): base_dir = get_base_directory() dirs = { - 'config': base_dir / 'config', - 'cache': base_dir / 'cache', - 'logs': base_dir / 'logs', + "config": base_dir / "config", + "cache": base_dir / "cache", + "logs": base_dir / "logs", } for name, path in dirs.items(): @@ -154,7 +156,7 @@ def setup(): logging.debug(f"{name.capitalize()} directory ensured at: {path}") 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" if not env_path.exists(): @@ -171,7 +173,7 @@ def setup(): "Approved": [], "Needs_Review": ["Review_First", "Review_Second", "HTML"], "Preflight": ["HTML"], - "Archived": [] + "Archived": [], } for folder_name, subfolders in folders_structure.items(): @@ -183,7 +185,7 @@ def setup(): subfolder_path.mkdir(parents=True, exist_ok=True) 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} protected_config = load_protected_config() @@ -194,10 +196,12 @@ def setup(): if not url: url = os.getenv("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 set_key(env_path, "URL", url) os.environ["URL"] = url logging.debug(f"Service URL set to: {url}") - write_config_to_env(merged_config, env_path) \ No newline at end of file + write_config_to_env(merged_config, env_path) diff --git a/utils/test.py b/utils/test.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/tui.py b/utils/tui.py index 36deeda..6eb070d 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -5,7 +5,7 @@ import sys import dotenv from dotenv import set_key from textual.app import App, ComposeResult -from textual.containers import Horizontal, Vertical +from textual.containers import Vertical from textual.reactive import reactive from textual.screen import Screen from textual.widgets import ( @@ -16,18 +16,24 @@ from textual.widgets import ( Static, Tab, 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.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.API import AirlockAPIWrapper from services.policyhandler import confirmUpdateAfromE from utils.configmanager import load_env from utils.setup import get_base_directory, load_user_config 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() @@ -58,7 +64,9 @@ def _persist_user_theme(theme_name: str) -> None: config_dir.mkdir(parents=True, exist_ok=True) if not user_config_path.exists(): # 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 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) - - # --------------------------------------------------------------------------- # 1) SCREEN # --------------------------------------------------------------------------- @@ -105,7 +111,7 @@ class MainMenuScreen(Screen): ("🔀 - Move - Other", "move_other_button"), ], "otp": [ - ("🔐 - Generate OTPs", "otp_generate_button"), + ("🎫 - Generate OTPs", "otp_generate_button"), ("📊 - OTP Activities By Agent", "otp_activities_button"), ("❌ - Revoke OTPs", "otp_revoke_button"), ], @@ -115,42 +121,24 @@ class MainMenuScreen(Screen): ], } - # textual themes to expose - 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: + def __init__(self, api: AirlockAPIWrapper) -> None: super().__init__() + self.api = api self.extras = load_env("EXTRAS") wd = load_env("WORKING_DIR") or os.getcwd() if not os.path.isdir(wd): wd = os.getcwd() self.working_dir = wd - def _make_buttons_for(self, tab_id: str) -> Vertical: defs = self.BUTTON_DEFS.get(tab_id, []) buttons = [] for label, btn_id in defs: btn = Button(label, id=btn_id) - btn.styles.width = "100%" # Make button span full width of parent + btn.styles.width = "100%" buttons.append(btn) return Vertical(*buttons) - - - def compose(self) -> ComposeResult: yield Header(show_clock=True, icon="⚙") @@ -208,8 +196,7 @@ class MainMenuScreen(Screen): new_index = current + direction if 0 <= new_index < len(buttons): buttons[new_index].focus() - - + def switch_tab(self, tab_id: str) -> None: self.current_tab = tab_id content = self.query_one("#content", Vertical) @@ -221,88 +208,62 @@ class MainMenuScreen(Screen): elif tab_id == "dir": content.mount(DirectoryTree(self.working_dir, id="dir_tree")) elif tab_id == "p_tree": - layout = Horizontal() - 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()) - - + content.mount(PolicyTreeWidget(self.app.policies, self.app.devices)) elif tab_id == "settings": - # Create and mount the horizontal container - 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 - + content.mount(ThemeSelector()) else: content.mount(Static(f"Unknown tab: {tab_id}")) def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: self.switch_tab(event.tab.id) - def on_tree_node_selected(self, message: Tree.NodeSelected) -> None: - node = message.node - data = node.data + def on_multi_agent_selector_agents_selected( + self, message: MultiAgentSelector.AgentsSelected + ) -> 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: - details = "\n".join(f"{key}: {value}" for key, value in data.items()) - else: - details = f"Selected: {node.label}" + def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: + """Handle OTP generation request from the workflow.""" + global _PENDING_JOB - 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 logger.debug("Directory file selected: %s", path) try: @@ -316,14 +277,6 @@ class MainMenuScreen(Screen): button_id = event.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: case "find_device_button": _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {}) @@ -341,7 +294,10 @@ class MainMenuScreen(Screen): case "move_other_button": _PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {}) 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": _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {}) case "otp_revoke_button": @@ -359,14 +315,10 @@ class MainMenuScreen(Screen): self.app.exit() - - - - # --------------------------------------------------------------------------- # 2) APP # --------------------------------------------------------------------------- -class AirlockTools(App): +class Loxide(App): CSS = """ #logo { width: 100%; @@ -388,12 +340,23 @@ class AirlockTools(App): if not os.path.isdir(wd): wd = os.getcwd() 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.push_screen(MainMenuScreen()) + self.push_screen(MainMenuScreen(api)) def action_quit(self) -> None: global _PENDING_JOB @@ -407,7 +370,6 @@ class AirlockTools(App): screen.switch_tab("dir") - # --------------------------------------------------------------------------- # 3) TERMINAL + LEGACY # --------------------------------------------------------------------------- @@ -422,6 +384,7 @@ def _restore_terminal_for_legacy() -> None: if os.name == "nt": try: import ctypes + kernel32 = ctypes.windll.kernel32 handle = kernel32.GetStdHandle(-11) mode = ctypes.c_ulong() @@ -447,7 +410,7 @@ def _run_legacy_job(func, args, kwargs) -> None: # --------------------------------------------------------------------------- # 4) PUBLIC ENTRYPOINT # --------------------------------------------------------------------------- -def run_AirlockTools(api: AirlockAPIWrapper) -> None: +def run_Loxide(api: AirlockAPIWrapper) -> None: global _PENDING_JOB while True: @@ -456,7 +419,7 @@ def run_AirlockTools(api: AirlockAPIWrapper) -> None: dotenv.load_dotenv(dotenv_path=env_path, override=True) _PENDING_JOB = None - app = AirlockTools(api) + app = Loxide(api) try: app.run() @@ -478,6 +441,46 @@ def run_AirlockTools(api: AirlockAPIWrapper) -> None: # just loop again; fresh .env was already loaded at the top 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 @@ -486,4 +489,4 @@ def run_AirlockTools(api: AirlockAPIWrapper) -> None: # --------------------------------------------------------------------------- if __name__ == "__main__": api = AirlockAPIWrapper() - run_AirlockTools(api) + run_Loxide(api) diff --git a/utils/utils.py b/utils/utils.py index 539a588..26b2023 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -28,8 +28,6 @@ import pandas as pd logger = logging.getLogger(__name__) - - def import_to_dataframe(file_path: str) -> pd.DataFrame: df = pd.DataFrame() @@ -96,17 +94,17 @@ def choose_file(initial_directory=None, required_substring=None): return file_path - - def get_sanitized_input(prompt: str) -> str: while True: user_input = input(prompt) if user_input.strip() == "": return user_input # Allow blank lines - if re.match(r'^[a-zA-Z0-9_\- .]+$', user_input.strip()): + if re.match(r"^[a-zA-Z0-9_\- .]+$", user_input.strip()): return user_input else: - print("Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed.") + print( + "Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed." + ) def regulator(paths, case_insensitive=True): @@ -119,8 +117,10 @@ def regulator(paths, case_insensitive=True): pattern = "(?i)" + pattern # Add inline case-insensitive flag print(f"Regulator is providing: {pattern}") return pattern + + def irtang(): - print( + print( colorText( r""" ███ @@ -149,6 +149,8 @@ def irtang(): "yellow", ) ) + + def displayIntro(): print( @@ -164,6 +166,8 @@ def displayIntro(): "cyan", ) ) + + def welcome(): print( 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(): @@ -348,7 +362,11 @@ def printDeviceEnforceChecklist(): "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( colorText( " 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("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("7. Liftoff ------------------------------------------------------", "cyan")) + print( + colorText( + "7. Liftoff ------------------------------------------------------", "cyan" + ) + ) print( colorText( " Apply path exclusions according to allowed and approved paths", "cyan", ) ) - print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan")) - print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan")) + print( + colorText(" Apply signed or attested hashes to Parent Allow List", "cyan") + ) + print( + colorText( + " Apply approved, but unsigned hashes to the Child Allow List", "cyan" + ) + ) print( colorText( @@ -523,10 +559,9 @@ def formatHTML(df, output_html_path=None, overwrite=True): return styled_html - def open_directory(path): system = platform.system() - + if system == "Windows": os.startfile(path) elif system == "Linux": @@ -537,9 +572,9 @@ def open_directory(path): def print_x_wide(items: list, width: int): for i in range(0, len(items), width): - row = items[i:i+width] + row = items[i : i + width] print(" | ".join(row)) - + def clear_screen(): - os.system('cls' if os.name == 'nt' else 'clear') + os.system("cls" if os.name == "nt" else "clear") diff --git a/widgets/OTP_generate.py b/widgets/OTP_generate.py new file mode 100644 index 0000000..b35b39d --- /dev/null +++ b/widgets/OTP_generate.py @@ -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 diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py new file mode 100644 index 0000000..e46c636 --- /dev/null +++ b/widgets/multiagentselector.py @@ -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 diff --git a/widgets/policytreewidget.py b/widgets/policytreewidget.py new file mode 100644 index 0000000..3a0efff --- /dev/null +++ b/widgets/policytreewidget.py @@ -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 diff --git a/widgets/themeselector.py b/widgets/themeselector.py new file mode 100644 index 0000000..fda7cd1 --- /dev/null +++ b/widgets/themeselector.py @@ -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))