Files
AirlockTools/Development/WIP/Sync Code/utils.py
T
2025-10-17 09:26:56 -04:00

709 lines
26 KiB
Python

# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import logging
import os
import platform
import re
import subprocess
import tempfile
import tkinter as tk
from tkinter import filedialog
import pandas as pd
from utils.configmanager import load_env
logger = logging.getLogger(__name__)
def import_to_dataframe(file_path: str) -> pd.DataFrame:
df = pd.DataFrame()
try:
if not os.path.exists(file_path):
print(colorText(f"Error: File '{file_path}' does not exist.", "red"))
return df
ext = os.path.splitext(file_path)[1].lower()
if ext == ".csv":
df = pd.read_csv(file_path)
elif ext == ".parquet":
df = pd.read_parquet(file_path)
else:
print(colorText(f"Error: Unsupported file extension '{ext}'.", "red"))
return df
if df.empty:
print(colorText("Error: File has headers but no data rows.", "red"))
else:
print(colorText(f"Data loaded successfully from {file_path}", "green"))
return df
except pd.errors.EmptyDataError:
print(
colorText(
"Notice: CSV file is completely empty, falling back to empty frame",
"white",
)
)
return pd.DataFrame()
except Exception as e:
print(colorText(f"Error reading file: {e}", "red"))
return pd.DataFrame()
def choose_directory():
root = tk.Tk()
root.withdraw() # Hide the main window
directory = filedialog.askdirectory(title="Select a Directory")
print("Selected directory:", directory)
return directory
def choose_file(initial_directory=None, required_substring=None):
"""Open a file dialog and ensure the selected file contains a required substring."""
while True:
root = tk.Tk()
root.withdraw() # Hide the main window
file_path = filedialog.askopenfilename(initialdir=initial_directory)
if not file_path:
print("No file selected.")
return None
if required_substring and required_substring not in file_path:
print(
f"The selected file must contain '{required_substring}' in its path or name. Please try again."
)
else:
return file_path
def get_sanitized_input(prompt: str) -> str:
while True:
user_input = input(prompt)
if user_input.strip() == "":
return user_input # Allow blank lines
if re.match(r'^[a-zA-Z0-9_\- .]+$', user_input.strip()):
return user_input
else:
print("Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed.")
def regulator(paths, case_insensitive=True):
"""
Build a regex pattern that matches any of the given Windows path fragments.
"""
escaped = [re.escape(p) for p in paths]
pattern = "(?:" + "|".join(escaped) + ")"
if case_insensitive:
pattern = "(?i)" + pattern # Add inline case-insensitive flag
print(f"Regulator is providing: {pattern}")
return pattern
def displayIntro():
print(
colorText(
r"""
███
████ ░████████
█████████████ ███████████████
█████████████████████ █████████████████████
███████████████████ ██████████████████████▓
███████████████████ ██████████████████████
█████████████████████ ███████████████████████
████████████████████████████████████████████████████████
█████████ ██ ██ █████████
█████████ ██ ███ █ █████████
█████████ ██ ████ █████ █████████████
█████████ ██ ██████ █████████████
████████ ██ ███████ ████████████░
███████ ██ ██▓ ██████ ████████████
██████ ██ ████ █████ ███████████
█████████████████████████████████████████████████
▒████████████████████ ██████████████████
███████████████████ ███████████████▒
███████████████ █████████████
██████████ ███████████
████████
████
""",
"yellow",
)
)
print(
colorText(
r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""",
"cyan",
)
)
print(
colorText(
"=================================================================================",
"cyan",
)
)
print(
colorText(
"======================== Welcome to the Airlock API Tool ========================",
"cyan",
)
)
print(
colorText(
"=================================================================================",
"cyan",
)
)
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
"cyan",
)
)
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
"\nSequentually follow these steps to prepare a policy for enforcement:",
"white",
)
)
print(
colorText(
"\n1. Choose which originating policy or policies to move to enforcement",
"cyan",
)
)
if not selected_policies:
print(colorText(" [✗] No policies have been chosen", "red"))
else:
print(colorText("The following policies have been choosen:", "green"))
for policy in selected_policies:
print(colorText(f" [✓] {policy.name}", "green"))
print(colorText("2. Choose the destination policy and allowlist", "cyan"))
if not destination_policy:
print(colorText(" [✗] No destination policy has been chosen", "red"))
elif destination_policy:
print(colorText(f" [✓] {destination_policy[0].name} has been selected as the destination policy", "green"))
if not destination_allowlist:
print(colorText(" [✗] No allowlist has been chosen", "red"))
elif destination_allowlist:
print(
colorText(
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
"green",
)
)
print(
colorText(
"3. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
"cyan",
)
)
if not selected_policies:
print(colorText(" [✗] No policies have been chosen", "red"))
else:
if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
print(colorText(" [✓] Data has been fetched", "green"))
else:
print(colorText(" [✗] Data has not been fetched", "red"))
print(colorText("4. Manually review the files:", "cyan"))
print(
colorText(
" 'prepare_policy\\needs_approved\\good_hashes.csv' and 'prepare_policy\\needs_approved\\unknown_hashes.csv'\n",
"cyan",
)
)
print(
colorText(
" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
"cyan",
)
)
print(
colorText(
" If metarules need to be created, please make note of them, and remove the row from the csv.",
"cyan",
)
)
print(
colorText(
" When complete, save both csv files to the directory 'approved' and choose this option.",
"cyan",
)
)
print(
colorText(
" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
"cyan",
)
)
if os.path.exists(f"{working_dir}\\Approved\\approved_executions.csv"):
print(colorText(" [✓] Reviewed hashes have been loaded", "green"))
else:
print(colorText(" [✗] Reviewed hashes have not been loaded", "red"))
if os.path.exists(
f"{working_dir}\\Needs_Review\\Review_Second\\primary_Paths.csv",
):
print(colorText(" [✓] Path review list created", "green"))
else:
print(colorText(" [✗] Path review list has not been created", "red"))
print(
colorText(
"5. Manually review the files \n 'prepare_policy\\needs_approved\\primary_Paths.csv'\n 'prepare_policy\\needs_approved\\secondary_Paths.csv'",
"cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions you do not approve of. The secondary list can be not added at all if nothing is useful",
"cyan",
)
)
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
print(
colorText(
" Do the same process with the list of publishers forthe same directories",
"cyan",
)
)
print(colorText(" Preflight Lists will be generated", "cyan"))
if os.path.exists(
f"{working_dir}\\Approved\\primary_Paths.csv",
):
print(colorText(" [✓] Reviewed path list detected", "green"))
else:
print(colorText(" [✗] Path review list has not been detected", "red"))
if os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv") and os.path.exists(
f"{working_dir}\\Preflight\\approved_hashes.csv"
):
print(colorText(" [✓] Preflight Path Exclusion List has been generated", "green"))
else:
print(colorText(" [✗] Preflight Path Exclusion List has not been generated", "red"))
print(colorText("6. Test ------------------------------------------------------", "cyan"))
print(colorText(" Print rather than apply selected data.", "cyan"))
print(colorText("7. Liftoff ------------------------------------------------------", "cyan"))
print(
colorText(
" Apply path exclusions according to allowed and approved paths",
"cyan",
)
)
print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan"))
print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
print(
colorText(
"R. Remove/Reset Generated data - will prompt to allow keeping execution history",
"cyan",
)
)
print(colorText("F. 📂 - Open Working Directory", "cyan"))
print(colorText("Q. 🔚 - Quit", "cyan"))
def areYouSure():
print(
colorText(
"🛑****************************************************************************************************************************************🛑",
"red",
)
)
print(
colorText(
"⚠️=========================================================================================================================================⚠️",
"yellow",
)
)
print(
colorText(
"🛑========================================================================================================================================🛑",
"red",
)
)
print(
colorText(
"⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️",
"yellow",
)
)
print(
colorText(
"🛑========================================================================================================================================🛑",
"red",
)
)
print(
colorText(
"⚠️=========================================================================================================================================⚠️",
"yellow",
)
)
print(
colorText(
"🛑****************************************************************************************************************************************🛑",
"red",
)
)
def locked():
print(
colorText(
r"""
████████████████████████████████████████████████████████████████
███ ██
██ ██████ ███
██ ████████████ ███
██ ████ ███ ███
██ ███ ███ ███
██ ███ ███ ███
██ ▒████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ███
███ ███
████████████████████████████████████████████████████████████████████
▒██████████████████████████████████████████████████████████████████▒
▒████
▒████
▓██████████████████████████████████████████
█████████████████████████████████████████████░
""",
"yellow",
)
)
def printDeviceEnforceChecklist():
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
"cyan",
)
)
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
"\nSequentually follow these steps to prepare a policy for enforcement:",
"white",
)
)
print(
colorText(
"\n1. Choose which originating policy or policies to move to enforcement",
"cyan",
)
)
print(
colorText(
"2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
"cyan",
)
)
print(colorText("3. Manually review the files:", "cyan"))
print(
colorText(
" 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'",
"cyan",
)
)
print(
colorText(
" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
"cyan",
)
)
print(
colorText(
" If metarules need to be created, please make note of them, and remove the row from the csv.",
"cyan",
)
)
print(
colorText(
" When complete, save both csv files to the directory 'approved' and choose this option.",
"cyan",
)
)
print(
colorText(
" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
"cyan",
)
)
print(
colorText(
"4. Manually review the file 'needs_approved\\paths_needing_review.csv'",
"cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions you do not approve of",
"cyan",
)
)
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
print(
colorText(
" Do the same process with the list of publishers forthe same directories",
"cyan",
)
)
print(colorText(" Preflight Lists will be generated", "cyan"))
print(colorText("5. Choose the destination policy and parent and child allow list", "cyan"))
print(colorText("6. Test ------------------------------------------------------", "cyan"))
print(colorText(" Print rather than apply selected data.", "cyan"))
print(colorText("7. Liftoff ------------------------------------------------------", "cyan"))
print(
colorText(
" Apply path exclusions according to allowed and approved paths",
"cyan",
)
)
print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan"))
print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
print(
colorText(
"R. Remove/Reset Generated data - will prompt to allow keeping execution history",
"cyan",
)
)
print(colorText("B. Back", "cyan"))
def colorText(text, color):
colors = {
"red": "\033[91m",
"green": "\033[92m",
"yellow": "\033[93m",
"blue": "\033[94m",
"magenta": "\033[95m",
"cyan": "\033[96m",
"white": "\033[97m",
"reset": "\033[0m",
}
return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"
def formatHTML(df, output_html_path=None, overwrite=True):
from datetime import datetime
# Get current date and filename for subtitle
today = datetime.now().strftime("%d %B %Y") # Changed to "Day Month Year"
filename = output_html_path.replace(".html", "") if output_html_path else "Report"
dark_css = """
<style>
body {
background-color: #000000;
margin: 0;
padding: 0;
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
color: #f8f8f2;
}
.header {
text-align: center;
margin: 20px auto;
padding: 10px;
border-bottom: 2px solid #ffd700;
max-width: 95%;
}
.header h1 {
color: #ffd700;
margin: 0;
font-size: 32px;
}
.header p {
color: #00bfff;
margin: 5px 0 0 0;
font-size: 18px;
}
.table-container {
overflow-y: scroll;
margin: 0 auto;
width: 95%;
max-height: calc(80vh - 100px);
display: block;
border: 1px solid #3a3a4d;
margin-bottom: 0;
}
table {
border-collapse: collapse;
font-size: 14px;
background-color: #1e1e2f;
color: #f8f8f2;
width: max-content;
}
th, td {
border: 1px solid #3a3a4d;
text-align: left;
padding: 10px;
max-width: 300px;
word-wrap: break-word;
overflow-wrap: break-word;
}
/* First column: no wrap */
td:nth-child(1), th:nth-child(1) {
white-space: nowrap;
max-width: none !important;
word-wrap: normal !important;
}
th {
background-color: #2e2e40;
color: #ffd700;
position: sticky;
top: 0;
z-index: 10;
}
tr:nth-child(even) {
background-color: #262638;
}
tr:hover {
background-color: #33334d;
color: #00bfff;
}
/* Custom scrollbar styling */
.table-container::-webkit-scrollbar {
width: 12px;
}
.table-container::-webkit-scrollbar-track {
background: #1e1e2f;
}
.table-container::-webkit-scrollbar-thumb {
background-color: #3a3a4d;
border-radius: 6px;
}
</style>
"""
header = f"""
<div class="header">
<h1>Airlock Tools</h1>
<p>{filename} - {today}</p>
</div>
"""
html_table = df.to_html(index=False, escape=False)
styled_html = (
f"<html>\n"
f"<head><title>Airlock Tools Report</title></head>\n"
f"<body>\n"
f"{dark_css}\n"
f"{header}\n"
f"<div class='table-container'>\n"
f" {html_table}\n"
f"</div>\n"
f"</body>\n"
f"</html>"
)
if output_html_path:
with open(output_html_path, "w", encoding="utf-8") as f:
f.write(styled_html)
print(f"✅ Styled table saved to '{output_html_path}'")
elif overwrite:
with tempfile.NamedTemporaryFile(
suffix=".html", delete=False, mode="w", encoding="utf-8"
) as f:
f.write(styled_html)
temp_path = f.name
print(f"✅ Styled table saved to temporary file: {temp_path}")
else:
return styled_html
def open_directory(path):
system = platform.system()
if system == "Windows":
os.startfile(path)
elif system == "Linux":
subprocess.run(["xdg-open", path])
else:
raise OSError(f"Unsupported operating system: {system}")