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()
)
)