diff --git a/utils/hashfunctions.py b/utils/hashfunctions.py new file mode 100644 index 0000000..ff57f43 --- /dev/null +++ b/utils/hashfunctions.py @@ -0,0 +1,84 @@ +import pandas as pd +import requests +import os + +def aggregateHashes(executions_json: dict) -> pd.DataFrame: + """ + Takes the executions, aggregates all the data with sha256 as primary, then returns aggregated dataframe + """ + + exechistories = executions_json.get("response", {}).get("exechistories", []) + df = pd.DataFrame(exechistories) + + if df.empty: + return df + + # Aggregate by sha256 - keep all entries in lists + agg_df = df.groupby("sha256").agg(lambda x: list(x)).reset_index() + + 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 + """ + endpoint = url + 'v1/hash/query' + payload = agg_df['sha256'].tolist() + headers = {"X-APIKey": os.getenv('APIKEY')} + + response = requests.post(endpoint, headers=headers, json=payload, verify=False) + data = response.json() + results = data.get("response", {}).get("results", []) + + rows = [] + for res in results: + row = {"sha256": res.get("sha256"), "result": res.get("result")} + + if "data" in res: + d = res["data"] + for key in ["filename", "filepath", "description", "filesize", "md5", + "productname", "productversion", "publisher", "createtime", "modtime", + "sha128", "sha384", "sha512", "datetime"]: + row[key] = d.get(key) + + row["applications"] = d.get("applications", []) + row["baselines"] = d.get("baselines", []) + + reputation = d.get("reputation", {}) + for k, v in reputation.items(): + row[f"reputation_{k}"] = v + + rows.append(row) + + df_api = pd.DataFrame(rows) + + aug_df = agg_df.merge(df_api, on="sha256", how="left") + + return aug_df + +def categorize_hashes(aug_df: pd.DataFrame, threat_tolerance: int, untrusted_publishers: list): + """ + Categorize hashes into needsreview, approved, and remaining based on publisher and threat level. + + """ + if untrusted_publishers is None: + untrusted_publishers = [] + + # Flatten threatlevel from nested reputation dict + df = df.copy() + df["threatlevel"] = df["reputation"].apply(lambda x: x.get("threatlevel") if pd.notnull(x) else None) + + # Masks for each category + mask_needsreview = (df["publisher"] == "Not Signed") & (df["threatlevel"] > threat_tolerance) + mask_approved = (df["publisher"] != "Not Signed") & (~df["publisher"].isin(untrusted_publishers)) + + # Create DataFrames for each category + needsreview_df = df[mask_needsreview].drop(columns=["threatlevel"]) + approved_df = df[mask_approved].drop(columns=["threatlevel"]) + remaining_df = df[~(mask_needsreview | mask_approved)].drop(columns=["threatlevel"]) + + return needsreview_df, approved_df, remaining_df + +def approve_hashes(approved_df: pd.DataFrame): + pass \ No newline at end of file