AirlockTools/utils/otpfunctions.py

import json
import requests
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):

    endpoint = url + f'/v1/otp/usage'
    payload = {
            "status" : "0"
    }

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


    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 * .9)
        #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}, early run", "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)
        #When duration completes - 
        
        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.time()
        
        if not os.path.exists(f"OTP\\PARQ\\localapprovalhistory.parquet"): 
            df = pd.DataFrame()
            df.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
        else:
            history = pd.read_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
            history = pd.concat([history, finalhashesadded], ignore_index=True)
            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)