mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Memory Optimization Draft one complete

Browse Source
This commit is contained in:
=
2025-09-01 20:19:11 -04:00
parent 7ef5a83ff2
commit 67615f364c
14 changed files with 440 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/hashfunctions.py
+33 -14
View File
@@ -41,23 +41,28 @@ def aggregateHashes(executions_json) -> pd.DataFrame:
return agg_df
def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame:
"""
Takes output of aggregatedHashes, queries API for those hashes, flattens response while keeping one row per hash,
aggregate applications and baselines into lists, then merges results back into agg_df to create a
"""
if 'sha256' not in agg_df.columns or agg_df.empty:
print("⚠️ 'sha256' column missing or DataFrame is empty. Skipping API query.")
return agg_df.copy() # Return as-is to avoid breaking downstream logic
endpoint = url + '/v1/hash/query'
payload = {
"hashes": agg_df['sha256'].tolist()
"hashes": agg_df['sha256'].tolist()
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
data = response.json()
data = response.json()
results = data.get("response", {}).get("results", [])
rows = []
for res in results:
row = {"sha256": res.get("sha256"), "result": res.get("result")}
@@ -80,10 +85,24 @@ def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame:
df_api = pd.DataFrame(rows)
if 'sha256' not in df_api.columns:
print("⚠️ API response missing 'sha256'. Skipping merge.")
return agg_df.copy()
df = agg_df.merge(df_api, on="sha256", how="left")
aug_df = df[['sha256', 'filename_x', 'description', 'productname', 'productversion', 'publisher_y', 'publisher_x', 'netdomain', 'hostname', 'username', 'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount', 'reputation_status', 'reputation_threatlevel', 'reputation_threatname', 'reputation_timestamp', 'pprocess', 'gprocess', 'commandline', ]]
# Only include columns that exist to avoid KeyErrors
expected_columns = ['sha256', 'filename_x', 'description', 'productname', 'productversion',
'publisher_y', 'publisher_x', 'netdomain', 'hostname', 'username',
'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
'reputation_status', 'reputation_threatlevel', 'reputation_threatname',
'reputation_timestamp', 'pprocess', 'gprocess', 'commandline']
available_columns = [col for col in expected_columns if col in df.columns]
aug_df = df[available_columns]
return aug_df
def categorizeHashes(aug_df: pd.DataFrame, threat_tolerance: int, untrusted_publishers: list):
if untrusted_publishers is None:
@@ -94,29 +113,29 @@ def categorizeHashes(aug_df: pd.DataFrame, threat_tolerance: int, untrusted_publ
def reputationtool(row):
val = row["reputation_scannermatch"]
if pd.isna(val) or val == "N/A":
return row["publisher_y"] == "Not Signed"
return row["publisher"] == "Not Signed"
try:
return int(val) > threat_tolerance
except (ValueError, TypeError):
return row["publisher_y"] == "Not Signed"
return row["publisher"] == "Not Signed"
df["reputation_flag"] = df.apply(reputationtool, axis=1)
mask_needsreview = (
((df["publisher_y"] == "Not Signed") & df["reputation_flag"]) |
((df["publisher"] == "Not Signed") & df["reputation_flag"]) |
(df["reputation_status"] == "UNKNOWN")
)
mask_approved = (
(
(df["publisher_y"] != "Not Signed") &
~df["publisher_y"].isin(untrusted_publishers) &
(df["publisher"] != "Not Signed") &
~df["publisher"].isin(untrusted_publishers) &
~df["reputation_status"].isna()
) |
(
(df["publisher_y"] == "Not Signed") &
(df["publisher"] == "Not Signed") &
~df["reputation_flag"] &
~df["publisher_y"].isin(untrusted_publishers) &
~df["publisher"].isin(untrusted_publishers) &
~df["reputation_status"].isna()
)
)
utils/pathfunctions.py
+26
View File
@@ -17,6 +17,7 @@ import pandas as pd
import os
from itertools import chain
import ast
import re
@@ -134,3 +135,28 @@ def filter_and_drop(approved, eligiblepaths, min_hashes):
filtered = filtered[filtered['sha256'].apply(len) >= min_hashes]
return filtered
def inspect_parquet(path):
try:
df = pd.read_parquet(path)
print(f"✅ Successfully read: {path}")
print(f"📄 Columns: {df.columns.tolist()}")
print(f"🔢 Rows: {len(df)}")
return df
except Exception as e:
print(f"❌ Error reading {path}: {e}")
return pd.DataFrame()
def regulator(paths, case_insensitive=True):
"""
Build a Python raw string regex that matches any of the given Windows path fragments.
"""
escaped = [re.escape(p) for p in paths]
pattern = "(?:" + "|".join(escaped) + ")"
if case_insensitive:
pattern = pattern
print(f"Regulator is providing {pattern}")
return f'r"{pattern}"'
utils/pretty.py
+75
View File
@@ -136,3 +136,78 @@ def style_dataframe_dark(df, output_html_path=None, overwrite=True):
print(f"✅ Styled table saved to temporary file: {temp_path}")
else:
return styled_html
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 areYouSure():
print(colorText(f"*******************************************************************************************************************************************","red"))
print(colorText(f"*=========================================================================================================================================*","yellow"))
print(colorText(f"*=========================================================================================================================================*","red"))
print(colorText(f"*-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------*", "yellow"))
print(colorText(f"*=========================================================================================================================================*","red"))
print(colorText(f"*=========================================================================================================================================*","yellow"))
print(colorText(f"*******************************************************************************************************************************************","red"))
def locked():
print(colorText(r"""
████████████████████████████████████████████████████████████████
███ ██
██ ██████ ███
██ ████████████ ███
██ ████ ███ ███
██ ███ ███ ███
██ ███ ███ ███
██ ▒████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ███
███ ███
████████████████████████████████████████████████████████████████████
▒██████████████████████████████████████████████████████████████████▒
▒████
▒████
▓██████████████████████████████████████████
█████████████████████████████████████████████░
""", "yellow"))