Files
AirlockTools/utils/otpfunctions.py
T

205 lines
7.9 KiB
Python

import json
import requests
import os
import pandas as pd
import utils.pretty as ct
import math
import time
import utils.clientfunctions as clientf
import utils.pathfunctions as pathf
import utils.policyfunctions as policyf
from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler, find_and_prioritize_jobs_by_pid
def getActiveOTP(url):
endpoint = url + f'/v1/otp/usage'
payload = {
"status" : "1"
}
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)
otp = pd.DataFrame(result["response"]["otpusage"])
if os.path.exists("OTP\\PARQ\\newest_active_OTP.parquet"):
previous_run = pd.read_parquet("OTP\\PARQ\\newest_active_OTP.parquet")
previous_run.to_parquet("OTP\\PARQ\\old_active_OTP.parquet", index=False)
os.remove("OTP\\PARQ\\newest_active_OTP.parquet")
otp.to_parquet("OTP\\PARQ\\newest_active_OTP.parquet", index=False)
if not otp.empty:
ct.style_dataframe_dark(otp, f"OTP\\HTML\\newest_active_OTP.html")
def getOTPActivities(url, otpid):
endpoint = url + f'/v1/otp/activities'
payload = {"otpid": f"{otpid}"}
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)
new_data = pd.DataFrame(result["response"]["otpactivities"])
# Define file path
parquet_path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
# Check if file exists and read it
if os.path.exists(parquet_path):
existing_data = pd.read_parquet(parquet_path)
combined_data = pd.concat([existing_data, new_data], ignore_index=True)
combined_data.drop_duplicates(inplace=True)
else:
combined_data = new_data
# Save combined data
combined_data.to_parquet(parquet_path, index=False)
# Optional: generate styled HTML if there's data
if not combined_data.empty:
ct.style_dataframe_dark(combined_data, f"OTP/HTML/OTP_activities_{otpid}.html")
def monitorOTP(url, pups):
getActiveOTP(url)
old_otp_path = "OTP\\PARQ\\old_active_OTP.parquet"
new_otp_path = "OTP\\PARQ\\newest_active_OTP.parquet"
if os.path.exists(old_otp_path):
old_active_OTP = pd.read_parquet(old_otp_path)
else:
old_active_OTP = pd.DataFrame(columns=['otpid']) # Ensure expected column exists
current_active_OTP = pd.read_parquet(new_otp_path)
if 'otpid' not in current_active_OTP.columns: current_active_OTP = pd.DataFrame(columns=['otpid'])
if 'otpid' not in old_active_OTP.columns: old_active_OTP = pd.DataFrame(columns=['otpid'])
newly_added = current_active_OTP[~current_active_OTP['otpid'].isin(old_active_OTP['otpid'])]
still_in_OTP = old_active_OTP[old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
no_longer_OTP = old_active_OTP[~old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
register_function("addhash", addOTPHashes)
for _, row in newly_added.iterrows():
clientid = row['clientid']
duration = (int(row['duration']) * 60)
hostname = row['hostname']
purpose = row ['purpose']
pid = row['otpid']
early = math.floor(duration * .95)
#If newly added to the list - schedule adding the majority of the executions prior to the expiration of OTP period.
run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}", "addhash", time.time() + early, [url, clientid, pid, pups], None)
print(f"Processing: {pid} with other data: {row}")
for _, row in still_in_OTP.iterrows():
pid = row['otpid']
#While still in OTP, continue to update activities list
getOTPActivities(url,pid)
for _, row in no_longer_OTP.iterrows():
clientid = row['clientid']
hostname = row['hostname']
purpose = row ['purpose']
pid = row['otpid']
allowlist = clientf.getDestAllowlistFromClientID(url,clientid)
policy = clientf.getPolicyFromClientID(url,clientid)
"""
Devices can come out of OTP either by timeout, or by early move out of OTP. If they are manually moved out prior to the job to add hashes can run, we want to accelerate the job.
But first, we want to update the OTP activities one final time for the pid, then move up any jobs if they exist, then add the hashes to the local approval allowlist
"""
getOTPActivities(url,pid)
find_and_prioritize_jobs_by_pid(pid, 1)
addOTPHashes(url, clientid,pid, pups)
finalhashesadded = pd.read_parquet(f"OTP\\PARQ\\otp_activities_{pid}.parquet")
finalhashesadded['policy'] = policy
finalhashesadded['allowlist'] = allowlist
finalhashesadded['added_at'] = time.localtime()
if not os.path.exists(f"OTP\\PARQ\\localapprovalhistory.parquet"):
df = pd.DataFrame()
df.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
history = pd.read_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
history = pd.concat([history, finalhashesadded], ignore_index=True)
ct.style_dataframe_dark(history, f"localapproval_history.html")
history.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
os.remove(f"OTP\\PARQ\\otp_activities_{pid}.parquet")
del finalhashesadded
def addOTPHashes(url, clientid, otpid, pups):
path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
activities = pd.read_parquet(path)
pattern = pathf.regulator(pups)
allowlist = clientf.getDestAllowlistFromClientID(url, clientid)
# Initialize or preserve 'hash_added' column
if "hash_added" not in activities.columns: activities["hash_added"] = None
# Identify rows that should be added (not matching pattern and not already added)
approve_by_hash = activities[
~activities["filename"].str.contains(pattern, na=False) & (activities["hash_added"] != "added")
]
hashes_to_add = approve_by_hash["sha256"].tolist()
# Add hashes to policy
if hashes_to_add:
policyf.addHash(url, allowlist, hashes_to_add)
# Update 'hash_added' column
activities["hash_added"] = activities.apply(
lambda row: "do not add" if pd.notna(row["filename"]) and pattern in row["filename"]
else ("added" if row["sha256"] in hashes_to_add else row["hash_added"]),
axis=1
)
# Save the updated DataFrame
activities.to_parquet(path)
def generateOTP(url, agentid):
purpose = input(ct.colorText(" Please enter the purpose for the OTP: ", "white"))
possible_durations = [15, 60, 360, 1440, 10080]
duration_selected = " "
print(ct.colorText("Please select a duration:", "white"))
for i, option in enumerate(possible_durations, start=1):
print(f"{i}. {option}")
try:
choice = int(input("Enter the number of your choice: "))
if 1 <= choice <= len(possible_durations):
duration_selected = possible_durations[choice - 1]
print(ct.colorText(f"You selected: {duration_selected}", "yellow"))
else:
print(ct.colorText("Invalid choice.", "red"))
except ValueError:
print(ct.colorText("Invalid input. Please enter a number.", "red"))
endpoint = url + '/v1/otp/retrieve'
payload = {
"duration" : f"{duration_selected}",
"agentid" : f"{agentid}",
"purpose" : f"{purpose}"
}
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)
otpcode = result["response"]["otpcode"]
print(ct.colorText(f"The OPT code is: {otpcode}", "yellow"))