diff --git a/AirlockTools.py b/AirlockTools.py index 54e679d..bb7c7b6 100644 --- a/AirlockTools.py +++ b/AirlockTools.py @@ -65,13 +65,14 @@ def main(): ct.apivalidation() register_function("monitorOTP", utils.otpfunctions.monitorOTP) - register_function("updateAudit", utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices) - # register_function("monitorLA", la.monitorLA) + register_function("monitorLA", la.scheduleAddingLAHashes) + register_function("updateAuditPolicies", utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices) + if not os.path.exists("scheduling\\jobs.json"): recurring_job("monitorOTP", "monitorOTP", interval=60, unit="seconds", args=[url, pups]) - recurring_job("monitorLA", "monitorLA", interval=45, unit="seconds", args=[url, pups]) - recurring_job("updateAudit", "updateAudit", interval=10, unit="minutes", args=[url, policy_relationship_map]) + recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[url, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant]) + recurring_job("updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[url, policy_relationship_map]) else: reload_jobs() start_scheduler() @@ -84,12 +85,13 @@ def main(): def menu_main(): while True: ct.displayIntro(); - print(ct.colorText("1. 🖥️ - Get All Events for Single Device", "yellow")) + print(ct.colorText("1. 🖥️ - Get All Events for Single Device", "yellow")) print(ct.colorText("2. 🎫 - OTP", "yellow")) print(ct.colorText("3. 🔇 - Find Quiet Hosts", "yellow")) print(ct.colorText("4. 🔒 - Prepare Policy For Enforcement", "yellow")) print(ct.colorText("5. 🔄 - Update Audit Policies from Enforcement Policies", "yellow")) - print(ct.colorText("6 🔍 - Device Search", "yellow")) + print(ct.colorText("6. 🔍 - Device Search", "yellow")) + print(ct.colorText("7. ➡️ - Move Devices to Local Approval", "yellow")) print(ct.colorText("Q. 🔚 - Quit", "yellow")) choice = input(ct.colorText("\nEnter Menu Item: ", "white")) @@ -110,8 +112,7 @@ def menu_main(): utils.clientfunctions.findAgents(url, devicelist, False) elif choice == "7": - pass - + la.moveToLocalApproval(url, policy_relationship_map) elif choice == "8": las = la.getLocalApprovals(url) diff --git a/utils/clientfunctions.py b/utils/clientfunctions.py index af4b9b7..08ae71f 100644 --- a/utils/clientfunctions.py +++ b/utils/clientfunctions.py @@ -67,7 +67,7 @@ def getDestAllowlistFromClientID(url, clientid): result = json.loads(response.text) data = pd.DataFrame(result["response"]["agents"]) - allowlist = getDestAllowlist(url,data.loc[0, "groupid"]) + allowlist = policyf.getDestAllowlist(url,data.loc[0, "groupid"]) return allowlist diff --git a/utils/hashfunctions.py b/utils/hashfunctions.py index 52ef616..831330f 100644 --- a/utils/hashfunctions.py +++ b/utils/hashfunctions.py @@ -287,7 +287,7 @@ def generatePublist(all_hashes, bad_publisher_list): all_approved_hashes = ct.tryToReadParquet(all_hashes) #Drop all not signed, only keep unique values - publist = all_approved_hashes[all_approved_hashes['publisher'] != "Not Signed"].drop_duplicates(subset='publisher') + publist = all_approved_hashes[all_approved_hashes['publisher'] != "Not Signed"].drop_duplicates(subset=['publisher']) #Remove Bad publisher if somehow they made it this far pattern = pathf.regulator(bad_publisher_list) publist = publist[~publist["publisher"].str.contains(pattern, na=False)] diff --git a/utils/localapproval.py b/utils/localapproval.py index 5771cc1..da4ce75 100644 --- a/utils/localapproval.py +++ b/utils/localapproval.py @@ -22,13 +22,14 @@ from utils.perstscheduler import register_function, run_once_job, recurring_job, #Standard Libary Imports: import json -import math import os import re import time +import datetime #3rd Party Imports: import pandas as pd +import numpy as np import requests def getLocalApprovals(url): @@ -60,23 +61,65 @@ def getLocalApprovals(url): return local_approval -def getLocalApprovalActivities(url, policy_relationship_map): +def scheduleAddingLAHashes(url, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant): + try: + register_function("add_hash", returnFromLocalApproval) + register_function("move_device", clientf.moveAgentToEnforcement) + except Exception as e: + print(f"[ERROR] Failed to register functions: {e}") + return - localapprovals = getLocalApprovals(url) - - newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move = monitorAuditStatus(url, policy_relationship_map) + try: + approvals_df = getNewLocalApprovals(url) + if approvals_df.empty: + print("[INFO] No new local approvals found. Nothing to schedule.") + return + batches = approvals_df.groupby('batchid') + except Exception as e: + print(f"[ERROR] Failed to retrieve or group local approvals: {e}") + return - register_function("add_hash_and_return_enforcement", returnFromLocalApproval) + for batchid, batch_df in batches: + try: + duration_minutes = int(batch_df['duration'].iloc[0]) + start_time = datetime.datetime.now() + run_time = start_time + datetime.timedelta(minutes=duration_minutes) + early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95)) - + early_timestamp = early_time.timestamp() + run_timestamp = run_time.timestamp() + # Schedule add_hash job + try: + run_once_job( + f"add_hash_{batchid}", + "add_hash", + early_timestamp, + [url, batch_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant], + None + ) + print(f"[INFO] Scheduled add_hash for batch {batchid} at {early_time}") + except Exception as e: + print(f"[ERROR] Failed to schedule add_hash for batch {batchid}: {e}") + # Schedule move_device jobs + devices = batch_df['agentid'].drop_duplicates().tolist() + for device in devices: + try: + run_once_job( + f"move_device_{device}_{batchid}", + "move_device", + run_timestamp, + [url, device, policy_relationship_map], + None + ) + print(f"[INFO] Scheduled move_device for device {device} in batch {batchid} at {run_time}") + except Exception as e: + print(f"[ERROR] Failed to schedule move_device for device {device} in batch {batchid}: {e}") - - - - #run_once_job(f" ", "addhash", time.time() + early, [url, clientid, pid, pups], None) - + except Exception as e: + print(f"[ERROR] Failed to process batch {batchid}: {e}") + def returnFromLocalApproval(url, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant): # Get unique policy names from device list policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist()) @@ -91,8 +134,8 @@ def returnFromLocalApproval(url, device_df, policy_relationship_map, bad_publish policy_types = [1, 2, 6, 7] # Define output directories - parq_base_dir = "Local_Approval\\parquet\\" - needappr_base_dir = "Local_Approval\\needs_approved" + parq_base_dir = "Local_Approval\\PARQ\\" + needappr_base_dir = "Local_Approval\\PARQ\\" # Ensure output directories exist os.makedirs(parq_base_dir, exist_ok=True) @@ -136,10 +179,6 @@ def returnFromLocalApproval(url, device_df, policy_relationship_map, bad_publish allowlist = policyf.getDestAllowlist(url, destination_id) policyf.addHash(url, allowlist, hashes[hashes['group'] == policy_id]) - devices_to_move = device_df[device_df['groupid'] == policy_id] - devices = devices_to_move['agentid'].drop_duplicates().tolist() - for device in devices: - clientf.moveAgentToEnforcement(url, device, policy_relationship_map) def moveToLocalApproval(url, policy_relationship_map): @@ -275,8 +314,13 @@ def getNewLocalApprovals(url): # Find new entries new_entries = current_la[~current_la['key'].isin(old_la['key'])] + # Convert 'granted' to datetime and filter by last 10 minutes + new_entries['granted'] = pd.to_datetime(new_entries['granted'], utc=True, errors='coerce') + ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10) + recent_entries = new_entries[new_entries['granted'] > ten_minutes_ago] + # Save current approvals for next run current_la.drop(columns=['key'], inplace=True) current_la.to_parquet(old_la_path, index=False) - return new_entries + return recent_entries \ No newline at end of file