diff --git a/.gitignore b/.gitignore index 5145354..09db371 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ *.csv *__pycache__* *.parquet -chunkinator.json \ No newline at end of file +chunkinator.json +jobs.json \ No newline at end of file diff --git a/AirlockTools.py b/AirlockTools.py index f1bbb3d..14dee05 100644 --- a/AirlockTools.py +++ b/AirlockTools.py @@ -16,7 +16,7 @@ import argparse import dotenv import os import pandas as pd - +import time import urllib3 import utils.allowlist import utils.clientfunctions @@ -27,6 +27,7 @@ import utils.pathfunctions import utils.policyfunctions import utils.pretty as ct +from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) @@ -51,15 +52,20 @@ def main(): args = parser.parse_args() if args.monitorOTP: - - - # Non-interactive logic - if not os.path.exists("parquet"): os.makedirs("parquet") + # Non-interactive logic + if not os.path.exists("scheduling"): os.makedirs("scheduling") + if not os.path.exists("OTP"): os.makedirs("OTP") + if not os.path.exists("OTP\\HTML"): os.mkdir("OTP\\HTML") + if not os.path.exists("OTP\\PARQ"): os.mkdir("OTP\\PARQ") apivalidation() print(f"Running non-interactively to start monitoring OTP") + register_function("monitorOTP", utils.otpfunctions.monitorOTP) + + if not os.path.exists("scheduling\\jobs.json"): recurring_job("monitor", "monitorOTP", interval=60, unit="seconds", args=[url]) + else: + reload_jobs() + start_scheduler() - - # Process input here else: # Interactive logic apivalidation() diff --git a/parquet/newOTP b/parquet/newOTP deleted file mode 100644 index f80e99c..0000000 Binary files a/parquet/newOTP and /dev/null differ diff --git a/parquet/newest_active_OTP b/parquet/newest_active_OTP deleted file mode 100644 index f80e99c..0000000 Binary files a/parquet/newest_active_OTP and /dev/null differ diff --git a/parquet/old_active_OTP b/parquet/old_active_OTP deleted file mode 100644 index f80e99c..0000000 Binary files a/parquet/old_active_OTP and /dev/null differ diff --git a/parquet/otp_activities b/parquet/otp_activities deleted file mode 100644 index 6b2160c..0000000 Binary files a/parquet/otp_activities and /dev/null differ diff --git a/scheduling/jobs.json b/scheduling/jobs.json new file mode 100644 index 0000000..e75d055 --- /dev/null +++ b/scheduling/jobs.json @@ -0,0 +1,13 @@ +[ + { + "id": "monitor", + "type": "recurring", + "interval": 10, + "unit": "seconds", + "function": "monitorOTP", + "args": [ + "https://172.17.22.240:3129" + ], + "kwargs": {} + } +] \ No newline at end of file diff --git a/utils/hashfunctions.py b/utils/hashfunctions.py index d617118..3a3f819 100644 --- a/utils/hashfunctions.py +++ b/utils/hashfunctions.py @@ -372,4 +372,6 @@ def generatePreflights(first_policy, second_policy): del allowbyhash del pathexclusions - gc.collect() \ No newline at end of file + gc.collect() + + \ No newline at end of file diff --git a/utils/otpfunctions.py b/utils/otpfunctions.py index ab1c05a..d4776ee 100644 --- a/utils/otpfunctions.py +++ b/utils/otpfunctions.py @@ -4,6 +4,15 @@ import os import pandas as pd import utils.pretty as ct import shutil +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 + + + def getActiveOTP(url): @@ -19,45 +28,105 @@ def getActiveOTP(url): response = requests.post(endpoint, headers=headers, data=payload, verify=False) result = json.loads(response.text) otp = pd.DataFrame(result["response"]["otpusage"]) - otp.to_parquet("parquet\\newest_active_OTP.parquet", index=False) + otp.to_parquet("OTP\\PARQ\\newest_active_OTP.parquet", index=False) if not otp.empty: ct.style_dataframe_dark(otp, f"newest_active_OTP.html") - def getOTPActivities(url, otpid): - endpoint = url + f'/v1/otp/activities' - payload = { - "otpid" : f"{otpid}" - } - + 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) - otp = pd.DataFrame(result["response"]["otpactivities"]) - otp.to_parquet("parquet/otp_activities.parquet", index=False) - if not otp.empty: - ct.style_dataframe_dark(otp, f"OTP_activities_{otpid}.html") + new_data = pd.DataFrame(result["response"]["otpactivities"]) + # Define file path + parquet_path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet" -def compareOTPHistory(): + # 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 - if not os.path.exists("parquet\\old_active_OTP"): - shutil.copy2("parquet\\newest_active_OTP.parquet", "parquet\\old_active_OTP.parquet") + # Save combined data + combined_data.to_parquet(parquet_path, index=False) - old_active_OTP = pd.read_parquet("parquet\\old_active_OTP.parquet") - current_active_OTP = pd.read_parquet("parquet\\newest_active_OTP.parquet") + # 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) + + if not os.path.exists("OTP\\old_active_OTP"): + shutil.copy2("OTP\\PARQ\\newest_active_OTP.parquet", "OTP\\PARQ\\old_active_OTP.parquet") + + old_active_OTP = pd.read_parquet("OTP\\PARQ\\old_active_OTP.parquet") + current_active_OTP = pd.read_parquet("OTP\\PARQ\\newest_active_OTP.parquet") 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 = (row['duration'] * 60) + hostname = row['hostname'] + purpose = row ['purpose'] + pid = row['otpid'] + early = math.floor(duration * .9) + + run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}", "addhash", time.time() + early, args=[url, clientid, pid, pups]) + run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}", "addhash", time.time() + duration, args=[url, clientid, pid, pups]) + print(f"Processing: {pid} with other data: {row}") + + + for _, row in still_in_OTP.iterrows(): + pid = row['otpid'] + getOTPActivities(url,pid) + + for _, row in no_longer_OTP.iterrows(): pid = row['otpid'] # Access other columns via row['column_name'] print(f"Processing PID: {pid} with other data: {row}") + + +def addOTPHashes(url, clientid, otpid, pups): + path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet" + activities = pd.read_parquet(path) + + pattern = pathf.regulator(pups) + policy = clientf.getDestAllowlistFromClientID(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, policy, 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) diff --git a/utils/perstscheduler.py b/utils/perstscheduler.py new file mode 100644 index 0000000..d8bedb0 --- /dev/null +++ b/utils/perstscheduler.py @@ -0,0 +1,209 @@ +import schedule +import time +import json +import os +from typing import Callable, Any, List, Dict + +# File where all jobs are persisted +JOBS_FILE = "scheduling\\jobs.json" + +# Registry of functions that can be scheduled +FUNCTION_MAP: Dict[str, Callable] = {} + +# ------------------------------- +# Function Registration +# ------------------------------- + +def register_function(name: str, func: Callable): + """ + Register a function so it can be called by name later. + + Example: + register_function("say_hello", say_hello) + """ + FUNCTION_MAP[name] = func + +# ------------------------------- +# Persistence Helpers +# ------------------------------- + +def load_jobs() -> List[Dict[str, Any]]: + """Load jobs from the JSON file, or return [] if none exist.""" + if not os.path.exists(JOBS_FILE): + return [] + with open(JOBS_FILE, "r") as f: + return json.load(f) + +def save_jobs(jobs: List[Dict[str, Any]]): + """Save jobs to the JSON file (overwrite).""" + with open(JOBS_FILE, "w") as f: + json.dump(jobs, f, indent=4) + +# ------------------------------- +# Run Once Jobs +# ------------------------------- + +def run_once_job(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None): + """ + Schedule a job to run once at a specific timestamp. + """ + args = args or [] + kwargs = kwargs or {} + + def job_wrapper(): + """Executes the job once, then removes it.""" + if func_name not in FUNCTION_MAP: + print(f"[ERROR] Function '{func_name}' is not registered.") + return + FUNCTION_MAP[func_name](*args, **kwargs) + + # Remove from persistence + jobs = load_jobs() + jobs = [j for j in jobs if j["id"] != job_id] + save_jobs(jobs) + + # Clear from in-memory schedule + schedule.clear(job_id) + + delay_seconds = run_at_timestamp - time.time() + if delay_seconds <= 0: + print(f"[WARN] Job {job_id} scheduled in the past. Skipping.") + return + + # Schedule via schedule library + schedule.every(delay_seconds).seconds.do(job_wrapper).tag(job_id) + + # Save to JSON + jobs = load_jobs() + jobs.append({ + "id": job_id, + "type": "once", + "run_at": run_at_timestamp, + "function": func_name, + "args": args, + "kwargs": kwargs + }) + save_jobs(jobs) + +# ------------------------------- +# Recurring Jobs +# ------------------------------- + +def recurring_job(job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None): + """ + Schedule a recurring job. + + Args: + job_id: Unique job name + func_name: Function to call (must be registered) + interval: Number of units between runs + unit: "seconds", "minutes", "hours", "days" + args: Positional arguments for the function + kwargs: Keyword arguments for the function + """ + args = args or [] + kwargs = kwargs or {} + + def job_wrapper(): + if func_name not in FUNCTION_MAP: + print(f"[ERROR] Function '{func_name}' is not registered.") + return + FUNCTION_MAP[func_name](*args, **kwargs) + + # Pick correct scheduling unit + if unit == "seconds": + schedule.every(interval).seconds.do(job_wrapper).tag(job_id) + elif unit == "minutes": + schedule.every(interval).minutes.do(job_wrapper).tag(job_id) + elif unit == "hours": + schedule.every(interval).hours.do(job_wrapper).tag(job_id) + elif unit == "days": + schedule.every(interval).days.do(job_wrapper).tag(job_id) + else: + raise ValueError(f"Unsupported unit: {unit}") + + # Save to JSON + jobs = load_jobs() + # Ensure no duplicate job id in file + jobs = [j for j in jobs if j["id"] != job_id] + jobs.append({ + "id": job_id, + "type": "recurring", + "interval": interval, + "unit": unit, + "function": func_name, + "args": args, + "kwargs": kwargs + }) + save_jobs(jobs) + +# ------------------------------- +# Reload Saved Jobs +# ------------------------------- + +def reload_jobs(): + """Reload jobs from JSON and reschedule them.""" + jobs = load_jobs() + for job in jobs: + if job["type"] == "once": + # Only reschedule if still in the future + if job["run_at"] > time.time(): + run_once_job(job["id"], job["function"], job["run_at"], job.get("args"), job.get("kwargs")) + elif job["type"] == "recurring": + recurring_job( + job["id"], + job["function"], + job["interval"], + job["unit"], + job.get("args"), + job.get("kwargs") + ) + +# ------------------------------- +# Scheduler Loop +# ------------------------------- + +def start_scheduler(): + """ + Start the scheduler loop (blocking). + Call this once in your main program to begin. + """ + try: + while True: + schedule.run_pending() + time.sleep(0.5) + except KeyboardInterrupt: + print("[INFO] Scheduler stopped.") + + +""" +import time +from persistent_schedule import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler + +# Example functions +def greet(name, loud=False): + if loud: + print(f"HELLO, {name}!") + else: + print(f"Hello, {name}.") + +def add(a, b): + print(f"{a} + {b} = {a + b}") + +# Register functions +register_function("greet", greet) +register_function("add", add) + +# Reload saved jobs +reload_jobs() + +# Schedule a run-once job in 5 seconds +run_once_job("job1", "greet", time.time() + 5, args=["Alice"], kwargs={"loud": True}) + +# Schedule a recurring job every 10 seconds +recurring_job("job2", "add", interval=10, unit="seconds", args=[2, 3]) + +# Start scheduler loop +start_scheduler() + +""" \ No newline at end of file diff --git a/utils/policyfunctions.py b/utils/policyfunctions.py index 90a7c59..f6cfb71 100644 --- a/utils/policyfunctions.py +++ b/utils/policyfunctions.py @@ -25,7 +25,7 @@ import utils.allowlist def addHash(url, policy, hash): - print(f"Adding the following hashes to {policy}:") + print(f"Adding the following: {hash} \n to {policy}:") for p in hash: print(p)