Files
AirlockTools/utils/policyfunctions.py
T
2025-09-11 10:25:16 -04:00

259 lines
11 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 gc
import json
import os
import pandas as pd
import re
import requests
import utils.pretty as ct
import datetime
import requests
import ijson
from bson import ObjectId
import datetime
import tqdm
import sys
def addHash(url, policy, hash):
print(f"Adding the following: {hash} \n to {policy}:")
for p in hash:
print(p)
def addPath(url, policy, hash):
print(f"Adding the following Path Exclusions to {policy}:")
for p in hash:
print(p)
def addHashReal(url, allowlistID, hashlist):
endpoint = url + '/v1/hash/application/add'
print(ct.colorText("[+] Grabbing All Categories", "cyan"))
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'
print(ct.colorText("[+] Grabbing All Categories", "cyan"))
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 getPolicyInfo(url, policy, days):
executionhist_policy = pd.DataFrame()
exehist = pullPolicyExechistories(url, policy, days, True)
data = json.loads(exehist)
executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
if not executionhist_policy.empty:
executionhist_policyxecutionhist_policy = executionhist_policy[['sha256', 'publisher', 'filename', 'hostname', 'username', 'pprocess', 'gprocess', 'commandline']]
executionhist_policy = executionhist_policy.drop_duplicates(subset=['sha256', 'filename', 'hostname'])
executionhist_policy = executionhist_policy.sort_values(by=['sha256', 'filename'])
executionhist_policy.to_parquet(f"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, first_policy, second_policy, destination_name, destination_id, allowlist_parent_name, allowlist_parent_id, allowlist_child_name, allowlist_child_id):
pathexclusions = pd.read_parquet(f"parquet\\final_path_exclusions_{first_policy}_{second_policy}.parquet")
allowbyhash = pd.read_parquet(f"parquet\\final_hash_approvals_{first_policy}_{second_policy}.parquet")
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"))
pathexcludelist = pathexclusions['longestcfp'].unique().tolist()
# Regex to match a Windows drive letter at the start (e.g., C:\)
drive_letter_pattern = re.compile(r'^[a-zA-Z]:\\')
# Processed list
processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + "**"
for path in pathexcludelist
]
addPath(url, destination_id,processed_paths)
print(ct.colorText(f"Adding hashes to {allowlist_parent_name}", "yellow"))
allowlist_parenthashlist = allowbyhash[allowbyhash['reputation_status'] == 'KNOWN']['sha256'].unique().tolist()
addHash(url, allowlist_parent_id,allowlist_parenthashlist)
print(ct.colorText(f"Adding hashes to {allowlist_child_name}", "yellow"))
allowlist_childhashlist = allowbyhash[allowbyhash['reputation_status'] == 'UNKNOWN']['sha256'].unique().tolist()
addHash(url, allowlist_child_id, allowlist_childhashlist)
ct.locked()
exit()
else:
print(ct.colorText("Operation aborted. You MUST EXPLICITLY AGREE to proceed.", "red"))
def pullPolicyExechistories(url, policiesnames, 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 crated.")
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:
json_response_data = checkpoint_stomper(checkpoint, url, 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, policy, headers):
json_output = {'error': 'Success', 'response': {'exechistories': []}}
endpoint = url + '/v1/logging/exechistories'
payload_dict = {
"type":[1,2,6,7],
"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)
policiesnames = []
policyids = []
for index, list in enumerate(parse_text['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 listAllowlists(url):
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, list in enumerate(parse_text['response']['applications'], start=1):
if index >= 38:
print(ct.colorText(f"{index}. {list['name']}", "yellow"))
policiesnames.append(list['name'])
policyids.append(list['applicationid'])
choice = int(input(ct.colorText("Select allowlist: ", "white")))
if choice < 38:
print(ct.colorText("Please only choose an allowlist designed for this use - '38+'","red"))
elif choice >= 38:
choice = choice - 38
return choice, policiesnames, policyids
#Need else and catch for upper bound
def skipback(days):
"""
Generate a MongoDB ObjectId for a given number of days ago from today.
Adds 1 extra day to the input to look further back.
"""
adjusted_days = days + 10
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)