diff --git a/utils/policyfunctions.py b/utils/policyfunctions.py new file mode 100644 index 0000000..b996238 --- /dev/null +++ b/utils/policyfunctions.py @@ -0,0 +1,657 @@ +# 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 . + +#Local Imports +import utils.utils as ct +import utils.hashfunctions as hashf +import utils.pathfunctions as pathf +import utils.policyfunctions as policyf + +#Standard Libary Imports: +import datetime +import gc +import json +import os +import re +import sys +#3rd Party Imports: +import ijson +import pandas as pd +import requests +import tqdm +from bson import ObjectId + +def addHash(url, policy, hash): + print(f"Adding the following: {hash} \n to {policy}:") + for p in hash: + pass + # print(p) + +def addPath(url, policy, hash): + print(f"Adding the following Path Exclusions to {policy}:") + for p in hash: + print(p) + +def addPub(url, policy, publist): + print(f"Adding the following Publishers to {policy}:") + for p in publist: + print(p) + +def addHashReal(url, allowlistID, hashlist): + endpoint = url + '/v1/hash/application/add' + + payload = { + "applicationid" : allowlistID, + "hashes" : hashlist + } + headers = { + "X-APIKey": os.getenv('APIKEY') + } + payload = json.dumps(payload) + response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) + response.raise_for_status() # Raise an error for bad status codes + parse_text = json.loads(response.text) + print(parse_text) + +def addPathReal(url, grouplistID, pathlist): + endpoint = url + '/v1/group/path/add' + payload = { + "groupid" : grouplistID, + "path" : pathlist + } + headers = { + "X-APIKey": os.getenv('APIKEY') + } + print(payload) + payload = json.dumps(payload) + response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) + print(response.text) + +def addPubReal(url, grouplistID, publist): + endpoint = url + '/v1/group/publisher/add' + payload = { + "groupid" : grouplistID, + "publisher" : publist + } + headers = { + "X-APIKey": os.getenv('APIKEY') + } + print(payload) + payload = json.dumps(payload) + response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) + print(response.text) + +def getPolicyInfo(url, policy, type, days, parquet=True): + executionhist_policy = pd.DataFrame() + exehist = pullPolicyExechistories(url, policy, type, days, True) + if exehist is not None: + data = json.loads(exehist) + executionhist_policy = pd.DataFrame(data["response"]["exechistories"]) + if not executionhist_policy.empty: + executionhist_policy = executionhist_policy[['datetime','sha256', 'publisher', 'filename', 'hostname', 'username', 'pprocess', 'gprocess', 'commandline']] + executionhist_policy['policy'] = policy # Add policy column here + executionhist_policy = executionhist_policy.drop_duplicates(subset=['sha256', 'filename', 'hostname']) + executionhist_policy = executionhist_policy.sort_values(by=['sha256', 'filename']) + if parquet: + executionhist_policy.to_parquet(f"prepare_policy\\parquet\\execution_history_{policy}.parquet", index=False) + print(ct.colorText(f"Staging of Execution history for policy: {policy} is complete", "green")) + del data + del exehist + gc.collect() + return executionhist_policy + +def sendToPolicy(url, paths, hashes, publishers, destination_name, destination_id, allowlist_name, allowlist_id): + pathexclusions = pd.read_parquet(paths) + allowbyhash = pd.read_parquet(hashes) + publishers = ct.tryToReadCSV(publishers) + + ct.areYouSure() + confirmation = input(ct.colorText("Type 'I AGREE' to continue: ","white")) + + + if confirmation.strip().upper() == "I AGREE": + print(ct.colorText("Proceeding with the code...", "yellow")) + print(ct.colorText(f"Adding path exclusions to {destination_name}", "yellow")) + + # Get unique combinations of longestcfp and file_extension + unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates() + + # Regex to match a Windows drive letter at the start (e.g., C:\) + drive_letter_pattern = re.compile(r'^[a-zA-Z]:\\') + + # Build processed paths like \\path\\**.exe or C:\path\**.jar + 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) + ] + + addPathReal(url, destination_id,processed_paths) + + print(ct.colorText(f"Adding publishers to {destination_name}", "yellow")) + + if publishers.empty: + print(ct.colorText("The publishers list is empty.", "red")) + else: + publisher_list = publishers['publisher'].tolist() + addPubReal(url, destination_id, publisher_list) + + + print(ct.colorText(f"These hashes would be added to {allowlist_name}", "yellow")) + + allowlist = allowbyhash['sha256'].unique().tolist() + addHash(url, allowlist_id,allowlist) + + ct.locked() + + exit() + + else: + print(ct.colorText("Operation aborted. You MUST EXPLICITLY AGREE to proceed.", "red")) + +def pullPolicyExechistories(url, policiesnames, type, days, outputjson: bool): + file_path = 'chunkinator.json' + if not os.path.exists(file_path): + with open(file_path, 'w') as file: + json.dump({'error': 'Success', 'response': {'exechistories': []}}, file) + print(f"File '{file_path}' has been created.") + else: + print(f"File '{file_path}' already exists.") + headers = {"X-APIKey": os.getenv('APIKEY')} + checkpoint = str(skipback(days)) + json_output = {'error': 'Success', 'response': {'exechistories': []}} + with tqdm.tqdm(file=sys.stdout, leave=True, total=10000, desc=f"Checkpoint Progess: {checkpoint}", colour="blue", initial=1) as filebar: + with tqdm.tqdm(file=sys.stdout, leave=True, total=100, desc=f"Total of {policiesnames} Complete: ") as pbar: + while True: + item = {} + json_response_data = checkpoint_stomper(checkpoint, url, type, policiesnames, headers) + histories = json_response_data['response']['exechistories'] + filebar.total=len(histories) + if not histories: + break + match_found = True + if match_found == True: + for index, item in enumerate(histories): + if index == len(histories) - 1: + checkpoint = item['checkpoint'] + filebar.desc = f"Checkpoint Progress: {checkpoint}" + break + else: + if (datetime.date.today() - datetime.timedelta(days=days) > datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()): + pass + else: json_output['response']['exechistories'].append(item) + filebar.update(1) + filebar.refresh() + seen = {} + if os.path.exists(file_path): + with open(file_path, 'r') as file: + existing_data = json.load(file) + combined = existing_data['response']['exechistories'] + json_output['response']['exechistories'] + else: + combined = json_output['response']['exechistories'] + for item in combined: + key = (item.get('sha256'), item.get('filename'), item.get('hostname')) + seen[key] = item + deduplicated = list(seen.values()) + with open(file_path, 'w') as file: + json.dump({'error': 'Success', 'response': {'exechistories': deduplicated}}, file) + json_output['response']['exechistories'].clear() + date_diff = datetime.date.today() - datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date() + percentage_diff = (((days + 10) - date_diff.days) / (days + 10)) * 100 + pbar.n = round(percentage_diff) + pbar.set_description_str(f"Total of {policiesnames} Complete: ") + pbar.refresh() + filebar.n = 1 + with open(file_path, 'r') as file: + final_output = json.load(file) + os.remove(file_path) + return json.dumps(final_output) if outputjson else None + +def checkpoint_stomper(checkpoint, url, type, policy, headers): + json_output = {'error': 'Success', 'response': {'exechistories': []}} + endpoint = url + '/v1/logging/exechistories' + payload_dict = { + "type":[type], + "checkpoint": checkpoint, + "policy": [policy] + } + payload = json.dumps(payload_dict) + with requests.request("POST", endpoint, headers=headers, data=payload, verify=False, stream=True) as response: + parser = ijson.items(response.raw, 'response.exechistories.item') + for item in parser: + key = (item.get('sha256'), item.get('hostname')) + if key not in json_output: + json_output['response']['exechistories'].append(item) + parse_text = json.loads(json.dumps(json_output)) + return parse_text + +def listPolicies(url): + endpoint = url + '/v1/group' + print(ct.colorText("[+] Grabbing All Policies", "cyan")) + payload = {} + headers = { + "X-APIKey": os.getenv('APIKEY') + } + response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) + parse_text = json.loads(response.text) + return parse_text + +def choosePolicies(url): + allpolicies = listPolicies(url) + policiesnames = [] + policyids = [] + for index, list in enumerate(allpolicies['response']['groups'], start=1): + print(ct.colorText(f"{index}. {list['name']}", "yellow")) + policiesnames.append(list['name']) + policyids.append(list['groupid']) + choice = input(ct.colorText("Select Policy Group: ", "white")) + choice = int(choice) - 1 + return choice, policiesnames, policyids + +def getPolicyDataframe(url) -> pd.DataFrame: + return pd.DataFrame(listPolicies(url)['response']['groups']) + +def listATPolicies(url): + endpoint = url + '/v1/group' + print(ct.colorText("[+] Grabbing All Policies", "cyan")) + + headers = { + "X-APIKey": os.getenv('APIKEY') + } + + try: + response = requests.post(endpoint, headers=headers, json={}, verify=False) + response.raise_for_status() + parse_text = response.json() + + at_policies = {} + + for index, group in enumerate(parse_text.get('response', {}).get('groups', []), start=1): + name = group.get('name', '') + if "AT" in name: + print(ct.colorText(f"{index}. {name}", "yellow")) + at_policies[name] = group.get('groupid') + + return at_policies + + except requests.exceptions.RequestException as e: + print(ct.colorText(f"[!] Request failed: {e}", "red")) + return {} + except (KeyError, json.JSONDecodeError) as e: + print(ct.colorText(f"[!] Failed to parse response: {e}", "red")) + return {} + +def listAllowlists(url: str) -> tuple[int, list, list]: + endpoint = url + '/v1/application' + print(ct.colorText("[+] Grabbing All Allowlists", "cyan")) + payload = {} + headers = { + "X-APIKey": os.getenv('APIKEY') + } + response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) + parse_text = json.loads(response.text) + policiesnames = [] + policyids = [] + + for index, item in enumerate(parse_text['response']['applications'], start=1): + if index >= 38: + print(ct.colorText(f"{index}. {item['name']}", "yellow")) + policiesnames.append(item['name']) + policyids.append(item['applicationid']) + + while True: + try: + choice = int(input(ct.colorText("Select allowlist: ", "white"))) + if choice < 38 or choice > len(parse_text['response']['applications']): + print(ct.colorText("Please only choose an allowlist designed for this use - '38+'", "red")) + else: + adjusted_choice = choice - 38 + return adjusted_choice, policiesnames, policyids + except ValueError: + print(ct.colorText("Invalid input. Please enter a number.", "red")) + +def skipback(days): + """ + Generate a MongoDB ObjectId for a given number of days ago from today. + """ + adjusted_days = days + date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=adjusted_days) + timestamp = int(date_days_ago.timestamp()) + hex_timestamp = format(timestamp, '08x') + objectid_hex = hex_timestamp + '0000000000000000' + return ObjectId(objectid_hex) + +def sendToPolicyTest(url, paths, hashes, publishers, destination_name, destination_id, allowlist_id, allowlist_name): + pathexclusions = pd.read_parquet(paths) + allowbyhash = pd.read_parquet(hashes) + publishers = ct.tryToReadCSV(publishers) + + + print(ct.colorText(f"These path exclusions would be added to {destination_name}", "yellow")) + + # Get unique combinations of longestcfp and file_extension + unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates() + + # Regex to match a Windows drive letter at the start (e.g., C:\) + drive_letter_pattern = re.compile(r'^[a-zA-Z]:\\') + + # Build processed paths like \\path\\**.exe or C:\path\**.jar + 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) + ] + + addPath(url, destination_id,processed_paths) + + print(ct.colorText(f"These publishers would added to {destination_name}", "yellow")) + + if publishers.empty: + print(ct.colorText("The publishers list is empty.", "red")) + else: + publisher_list = publishers['publisher'].tolist() + addPub(url, destination_id, publisher_list) + + print(ct.colorText(f"These hashes would be added to {allowlist_name}", "yellow")) + + allowlist = allowbyhash['sha256'].unique().tolist() + addHash(url, allowlist_id,allowlist) + +def updateAuditPoliciesFromEnforcementPolices(url, policy_relationship_map): + for enforcement_policy, audit_policy in policy_relationship_map.items(): + assignPoliciesfromGroup(url, enforcement_policy, audit_policy) + turnOnAudit(url, audit_policy) + +def assignPoliciesfromGroup(url, source_policy_id, target_policy_id): + + endpoint = url + '/v1/group/assign' + + payload = { + "groupid" : {source_policy_id}, + "targetgroupid" : {target_policy_id} + } + headers = { + "X-APIKey": os.getenv('APIKEY') + } + response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) + parse_text = json.loads(response.text) + print(parse_text) + +def turnOnAudit(url, policyid): + + endpoint = url + '/v1/group/settings/auditmode' + + payload = { + "groupid" : {policyid}, + "auditmode" : "1" + } + headers = { + "X-APIKey": os.getenv('APIKEY') + } + response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) + parse_text = json.loads(response.text) + print(parse_text) + +def agentsInPolicy(url, policyid): + endpoint = url + '/v1/group/agents' + + payload = { + "groupid" : {policyid} + } + headers = { + "X-APIKey": os.getenv('APIKEY') + } + response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) + parse_text = json.loads(response.text) + print(parse_text) + +def prepare_to_enforce(url, bad_publisher_list, pups, badpathparts, threat_tolerance_constant = 4, path_exclusion_constant = 4, min_files_for_path = 4): + + destination_name = " " + destination_id = " " + allowlist_name = " " + allowlist_id = " " + policylist = [] + type = [1, 2, 6, 7] + + parq_base_dir = "prepare_policy\\parquet\\" + appr_base_dir = "prepare_policy\\approved\\" + needappr_base_dir = "prepare_policy\\needs_approved\\" + pflight_base_dir = "prepare_policy\\preflight\\" + + #If the directorys where we're going to store our output dont exist, make them. + os.makedirs(parq_base_dir, exist_ok=True) + os.makedirs(needappr_base_dir, exist_ok=True) + os.makedirs(appr_base_dir, exist_ok=True) + os.makedirs(pflight_base_dir, exist_ok=True) + + while True: + + ct.printEnforceChecklist(parq_base_dir,appr_base_dir, needappr_base_dir, pflight_base_dir, policylist, allowlist_name, destination_name) + + choice = input(ct.colorText("\nEnter your choice: ", "white")) + + if choice == "1": + + policylist, policyids = getMultiplePolicySelections(url) + + elif choice == "2": + + print(ct.colorText(f"Please choose destination_name Policy for Path Exclusions","white")) + choice, policynames, policyid = choosePolicies(url) + #print(allowlist_parent_tuple) + destination_name = policynames[choice] + destination_id = policyid[choice] + + print(ct.colorText(f"Please choose Allowlist for Hashes","white")) + choice, allowlists,allowid = listAllowlists(url) + #print(allowlist_parent_tuple) + allowlist_name = allowlists[choice] + allowlist_id = allowid[choice] + + print(destination_name, allowlist_name) + + + elif choice == "3": + + policyf.buildExecHistory(url, + policylist, + parq_base_dir, + needappr_base_dir, + type, + threat_tolerance_constant, + bad_publisher_list, + pups, + ) + csvs = [f"{needappr_base_dir}unknown_hashes.csv", f"{needappr_base_dir}good_hashes.csv"] + #Since we want to build paths as if they were all in the same policy to begin with, lets group them that way + for csv in csvs: + df = ct.tryToReadCSV(csv) + df['policy'] = destination_name + df.to_csv(csv) + + + elif choice == "4": + pathf.generateApprovables(needappr_base_dir, appr_base_dir, parq_base_dir, badpathparts, bad_publisher_list, path_exclusion_constant, min_files_for_path) + + elif choice == "5": + savePreFlights(appr_base_dir, parq_base_dir, pflight_base_dir) + + + elif choice == "6": + if os.path.exists(f"{pflight_base_dir}final_path_exclusions.html") and os.path.exists(f"{pflight_base_dir}\\final_hash_approvals.html") and allowlist_name != " " and destination_name != " ": + sendToPolicyTest( + url, + f"{parq_base_dir}final_path_exclusions.parquet", + f"{parq_base_dir}final_hash_approvals.parquet", + f"{appr_base_dir}publishers.parquet", + destination_name, + destination_id, + allowlist_name, + allowlist_id + ) + + elif choice == "7": + if os.path.exists(f"{pflight_base_dir}final_path_exclusions.html") and os.path.exists(f"{pflight_base_dir}\\final_hash_approvals.html") and allowlist_name != " " and destination_name != " ": + sendToPolicy( + url, + f"{parq_base_dir}final_path_exclusions.parquet", + f"{parq_base_dir}final_hash_approvals.parquet", + f"{appr_base_dir}publishers.parquet", + destination_name, + destination_id, + allowlist_name, + allowlist_id, + ) + elif choice == "R": + + pathf.clean_folders(parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir) + + elif choice == "Q": + break + else: + print(ct.colorText("Invalid choice. Please try again.", "red")) + +def buildExecHistory(url, + policylist, + parq_base_dir, + needappr_base_dir, + type, + threat_tolerance_constant, + bad_publisher_list, + pups + ): + exe_hist_parq_list = [] + while True: + try: + history_days = int(input("Enter the how many days in of history do you want to pull - select a number between 1 and 150: ")) + if 1 <= history_days <= 150: + break + else: + print("Invalid input. Please enter a number between 1 and 150.") + except ValueError: + print("Invalid input. Please enter a valid integer.") + for policy in policylist: + policy_exec_history = policyf.getPolicyInfo(url, policy, type, history_days) + policy_exec_history['policy'] = policy + policy_exec_history.to_parquet(f"{parq_base_dir}Exec_Hist_{policy}.parquet") + exe_hist_parq_list.append(f"{parq_base_dir}Exec_Hist_{policy}.parquet") + + augmented_hashlist = hashf.combineHashes(url, exe_hist_parq_list) + augmented_hashlist.to_parquet(f"{parq_base_dir}augmentedHashlist.parquet",index=False) + + needsreview_df, approved_df, unapproved_df = hashf.categorizeHashes(augmented_hashlist, threat_tolerance_constant, bad_publisher_list, pups) + needsreview_df.to_parquet(f"{parq_base_dir}needsreview.parquet",index=False) + approved_df.to_parquet(f"{parq_base_dir}approved.parquet",index=False) + unapproved_df.to_parquet(f"{parq_base_dir}unapproved.parquet",index=False) + + + condensed_executions = hashf.condenseExecutions(exe_hist_parq_list) + condensed_executions.to_parquet(f"{parq_base_dir}condensed_executions.parquet", index=False) + + unknown, good, bad = hashf.divideSortedHashExecutions( + f"{parq_base_dir}needsreview.parquet", + f"{parq_base_dir}approved.parquet", + f"{parq_base_dir}unapproved.parquet", + f"{parq_base_dir}condensed_executions.parquet", + pups + ) + + dataframes = { + "unknown_hashes" : unknown, + "good_hashes": good, + "bad_hashes": bad + } + + for name, df in dataframes.items(): + df.to_csv(f"{needappr_base_dir}{name}.csv", index=False) + df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False) + ct.style_dataframe_dark(df, f"{needappr_base_dir}{name}.html") + +def savePreFlights(appr_base_dir, parq_base_dir, pflight_base_dir): + if os.path.exists(f"{appr_base_dir}primary_Paths.csv"): + if not os.path.exists(f"{parq_base_dir}final_hash_approvals.parquet") and not os.path.exists(f"{parq_base_dir}final_path_exclusions.parquet"): + + pathexclusions, allowbyhash = hashf.generatePreflights( + f"{parq_base_dir}all_hashes.parquet", + f"{appr_base_dir}primary_Paths.csv", + f"{appr_base_dir}secondary_Paths.csv") + + dataframes = { + "final_path_exclusions" : pathexclusions, + "final_hash_approvals": allowbyhash + } + + for name, df in dataframes.items(): + df.to_csv(f"{pflight_base_dir}{name}.csv", index=False) + df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False) + ct.style_dataframe_dark(df, f"{pflight_base_dir}{name}.html") + +def getMultiplePolicySelections(url): + + policynameslist = [] + policyidlist = [] + while True: + choice, policynames, policyid = choosePolicies(url) + selected_policy = policynames[choice] + selected_policyid = policyid[choice] + + if selected_policy not in policynameslist: + policynameslist.append(selected_policy) + + if selected_policyid not in policyidlist: + policyidlist.append(selected_policyid) + + + while True: + answer = input(ct.colorText("Do you want to load another policy? (yes/no): ", "white")).strip().lower() + if answer in ("no", "n"): + break # Exit the inner loop and then the outer loop + elif answer in ("yes", "y"): + break # Exit the inner loop and continue the outer loop + else: + print(ct.colorText("Please answer with 'yes' or 'no'.", "red")) + + if answer in ("no", "n"): + break + + + return policynameslist, policyidlist + +def getDestAllowlist(url, groupid): + + allowlists = getPolicyAllowlists(url, groupid) + + matches = allowlists[ + allowlists['name'].str.contains('local', case=False, na=False) & + allowlists['name'].str.contains('approval', case=False, na=False), + 'applicationid' + ].values + app_id = matches[0] if len(matches) > 0 else None + + return app_id + +def getPolicyAllowlists(url, groupid): + + endpoint = url + '/v1/group/policies' + payload = { + "groupid" : f"{groupid}" + } + + headers = {"X-APIKey": os.getenv('APIKEY')} + payload = json.dumps(payload) + + response = requests.post(endpoint, headers=headers, data=payload, verify=False) + result = json.loads(response.text) + data = pd.DataFrame(result["response"]["applications"]) + return data \ No newline at end of file