RustImplementation #23

Merged
mysticmomba merged 118 commits from RustImplementation into master 2025-11-04 18:13:24 -05:00
11 changed files with 328 additions and 28 deletions
Showing only changes of commit b0f4d84e78 - Show all commits
+1
View File
@@ -4,3 +4,4 @@
*__pycache__* *__pycache__*
*.parquet *.parquet
chunkinator.json chunkinator.json
jobs.json
+11 -5
View File
@@ -16,7 +16,7 @@ import argparse
import dotenv import dotenv
import os import os
import pandas as pd import pandas as pd
import time
import urllib3 import urllib3
import utils.allowlist import utils.allowlist
import utils.clientfunctions import utils.clientfunctions
@@ -27,6 +27,7 @@ import utils.pathfunctions
import utils.policyfunctions import utils.policyfunctions
import utils.pretty as ct 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) urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
@@ -51,15 +52,20 @@ def main():
args = parser.parse_args() args = parser.parse_args()
if args.monitorOTP: if args.monitorOTP:
# Non-interactive logic # Non-interactive logic
if not os.path.exists("parquet"): os.makedirs("parquet") 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() apivalidation()
print(f"Running non-interactively to start monitoring OTP") 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: else:
# Interactive logic # Interactive logic
apivalidation() apivalidation()
BIN
View File
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
+13
View File
@@ -0,0 +1,13 @@
[
{
"id": "monitor",
"type": "recurring",
"interval": 10,
"unit": "seconds",
"function": "monitorOTP",
"args": [
"https://172.17.22.240:3129"
],
"kwargs": {}
}
]
+2
View File
@@ -373,3 +373,5 @@ def generatePreflights(first_policy, second_policy):
del allowbyhash del allowbyhash
del pathexclusions del pathexclusions
gc.collect() gc.collect()
+86 -17
View File
@@ -4,6 +4,15 @@ import os
import pandas as pd import pandas as pd
import utils.pretty as ct import utils.pretty as ct
import shutil 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): def getActiveOTP(url):
@@ -19,45 +28,105 @@ def getActiveOTP(url):
response = requests.post(endpoint, headers=headers, data=payload, verify=False) response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text) result = json.loads(response.text)
otp = pd.DataFrame(result["response"]["otpusage"]) 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: if not otp.empty:
ct.style_dataframe_dark(otp, f"newest_active_OTP.html") ct.style_dataframe_dark(otp, f"newest_active_OTP.html")
def getOTPActivities(url, otpid): def getOTPActivities(url, otpid):
endpoint = url + f'/v1/otp/activities' endpoint = url + f'/v1/otp/activities'
payload = { payload = {"otpid": f"{otpid}"}
"otpid" : f"{otpid}"
}
headers = {"X-APIKey": os.getenv('APIKEY')} headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload) payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False) response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text) result = json.loads(response.text)
otp = pd.DataFrame(result["response"]["otpactivities"]) new_data = 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")
# 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"): # Save combined data
shutil.copy2("parquet\\newest_active_OTP.parquet", "parquet\\old_active_OTP.parquet") combined_data.to_parquet(parquet_path, index=False)
old_active_OTP = pd.read_parquet("parquet\\old_active_OTP.parquet") # Optional: generate styled HTML if there's data
current_active_OTP = pd.read_parquet("parquet\\newest_active_OTP.parquet") 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'])] 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'])] 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'])] 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(): 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'] pid = row['otpid']
# Access other columns via row['column_name'] # Access other columns via row['column_name']
print(f"Processing PID: {pid} with other data: {row}") 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)
+209
View File
@@ -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()
"""
+1 -1
View File
@@ -25,7 +25,7 @@ import utils.allowlist
def addHash(url, policy, hash): 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: for p in hash:
print(p) print(p)