Scheduler added

.gitignore

@@ -4,3 +4,4 @@
 *__pycache__*
 *.parquet
 chunkinator.json
+jobs.json

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()

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": {}
+    }
+]

utils/hashfunctions.py

@@ -373,3 +373,5 @@ def generatePreflights(first_policy, second_policy):
     del allowbyhash
     del pathexclusions
     gc.collect()
+
+    

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)

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()
+
+"""

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)