AirlockTools/utils/localapproval.py

#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

#Local Imports
import utils.clientfunctions as clientf
import utils.pathfunctions as pathf
import utils.policyfunctions as policyf
import utils.utils as ct
from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler, find_and_prioritize_jobs_by_pid

#Standard Libary Imports:
import json
import math
import os
import re
import time

#3rd Party Imports:
import pandas as pd
import requests

def getLocalApprovals(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)
    local_approval = pd.DataFrame(result["response"]["otpusage"])
    if os.path.exists("Local_Approval\\PARQ\\newest_local_approval.parquet"):
        previous_run = pd.read_parquet("Local_Approval\\PARQ\\newest_local_approval.parquet")
        previous_run.to_parquet("Local_Approval\\PARQ\\last_local_approval.parquet", index=False)
        os.remove("Local_Approval\\PARQ\\newest_local_approval.parquet")

    #Only keep rows presumably created by the generate local approval function
    local_approval = local_approval[local_approval['purpose'].str.startswith('🎫 Local Approval 🎫')]

    local_approval['batchid'] = local_approval['purpose'].apply(lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1))
    
    if not local_approval.empty:
        ct.style_dataframe_dark(local_approval, f"Local_Approval\\HTML\\newest_local_approval.html")
        local_approval.to_parquet("Local_Approval\\PARQ\\newest_local_approval.parquet", index=False)
    
    return local_approval
   
def getLocalApprovalActivities(url, policy_relationship_map):

    localapprovals = getLocalApprovals(url)
    
    newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move = monitorAuditStatus(url, policy_relationship_map)

    register_function("add_hash_and_return_enforcement", returnFromLocalApproval)

    



 
  
     
       #run_once_job(f" ", "addhash", time.time() + early, [url, clientid, pid, pups], None)
        
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())

    # Create inverse map to go from Audit to Enforcement
    inverse_map = {v: k for k, v in policy_relationship_map.items()}

    # Fetch all policies 
    all_policies = policyf.getPolicyDataframe(url)

    # Define policy types
    policy_types = [1, 2, 6, 7]

    # Define output directories
    parq_base_dir = "Local_Approval\\parquet\\"
    needappr_base_dir = "Local_Approval\\needs_approved"

    # Ensure output directories exist
    os.makedirs(parq_base_dir, exist_ok=True)
    os.makedirs(needappr_base_dir, exist_ok=True)

    # Build execution history
    policyf.buildExecHistory(
        url,
        policies_in_devicelist,
        parq_base_dir,
        needappr_base_dir,
        policy_types,
        threat_tolerance_constant,
        bad_publisher_list,
        pups,
    )

    # Read hash data
    unknown_hashes = ct.tryToReadCSV(f"{needappr_base_dir}unknown_hashes.csv")
    good_hashes = ct.tryToReadCSV(f"{needappr_base_dir}good_hashes.csv")
    hashes = pd.concat([unknown_hashes, good_hashes], ignore_index=True)

    # Process each policy
    for policy in policies_in_devicelist:
        # Filter for matching policy
        matching_rows = all_policies[all_policies['name'] == policy]
        
        if matching_rows.empty:
            print(f"Warning: No group ID found for policy '{policy}'. Skipping.")
            continue

        # Extract group ID
        policy_id = matching_rows['groupid'].values[0]

        # Map to destination ID
        destination_id = inverse_map.get(policy_id)
        if destination_id is None:
            print(f"Warning: No corresponding enforcement policy found for group ID '{policy_id}'. Skipping.")
            continue

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

    possible_durations = [15, 60, 360, 1440, 10080]
    duration_selected = None

    print(ct.colorText("Please select a duration:", "white"))
    for i, option in enumerate(possible_durations, start=1):
        print(f"{i}. {option}")

    try:
        choice = int(input("Enter the number of your choice: "))
        if 1 <= choice <= len(possible_durations):
            duration_selected = possible_durations[choice - 1]
            print(ct.colorText(f"You selected: {duration_selected}", "yellow"))
        else:
            print(ct.colorText("❌ Invalid choice.", "red"))
            return
    except ValueError:
        print(ct.colorText("❌ Invalid input. Please enter a number.", "red"))
        return

    devicelist = clientf.promptForDevices()
    device_df = clientf.findAgents(url, devicelist, True)

    batch = int(time.time())

    if device_df is None or device_df.empty:
        print(ct.colorText("❌ No agents found or error retrieving agents.", "red"))
        return

    for row in device_df.itertuples(index=False):
        try:
            addLocalApproval(url, batch, duration_selected, row.agentid)
            clientf.moveAgentToAudit(url, row.agentid, policy_relationship_map)
        except Exception as e:
            print(ct.colorText(f"❌ Error processing agent {row.agentid}: {e}", "red"))

def addLocalApproval(url, batchid, duration_selected, agentid):
    purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
    endpoint = url + '/v1/otp/retrieve'
    payload = {
        "duration": str(duration_selected),
        "agentid": str(agentid),
        "purpose": purpose
    }

    headers = {"X-APIKey": os.getenv('APIKEY')}
    response = requests.post(endpoint, headers=headers, data=json.dumps(payload), verify=False)

    try:
        result = response.json()
        otpcode = result["response"]["otpcode"]
        print(ct.colorText(f"The OTP code is: {otpcode}", "yellow"))
    except Exception as e: 
        print(ct.colorText(f"An unexpected error occurred for agent {agentid}: {str(e)}", "red"))
        
def monitorAuditStatus(url: str, policy_relationship_map: dict):
    
    # Simulated current agent list
    current_agent_list = clientf.findAllAgents(url)

    # Load old agent list
    old_agent_path = "Local_Approval\\PARQ\\last_agent_list.parquet"
    if os.path.exists(old_agent_path):
        old_agent_list = pd.read_parquet(old_agent_path)
    else:
        old_agent_list = pd.DataFrame(columns=current_agent_list.columns)

    # Merge on hostname
    merged = pd.merge(
        old_agent_list[['hostname', 'groupid']],
        current_agent_list[['hostname', 'groupid']],
        on='hostname',
        how='outer',
        suffixes=('_old', '_current'),
        indicator=True
    )

    # Reverse map for enforcement
    reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
    known_transitions = set(policy_relationship_map.items()) | set(reverse_policy_map.items())

    # 1. Newly added
    newly_added = merged[merged['_merge'] == 'right_only']

    # 2. Same policy
    same_policy = merged[
        (merged['_merge'] == 'both') &
        (merged['groupid_old'] == merged['groupid_current'])
    ]

    # 3. Moved to audit
    moved_to_audit = merged[
        (merged['_merge'] == 'both') &
        (merged['groupid_old'].isin(policy_relationship_map)) &
        (merged['groupid_current'] == merged['groupid_old'].map(policy_relationship_map))
    ]

    # 4. Moved to enforcement
    moved_to_enforcement = merged[
        (merged['_merge'] == 'both') &
        (merged['groupid_old'].isin(reverse_policy_map)) &
        (merged['groupid_current'] == merged['groupid_old'].map(reverse_policy_map))
    ]

    # 5. Unusual moves
    unusual_move = merged[
        (merged['_merge'] == 'both') &
        (merged['groupid_old'] != merged['groupid_current']) &
        merged.apply(lambda row: (row['groupid_old'], row['groupid_current']) not in known_transitions, axis=1)
    ]

    current_agent_list.to_parquet("Local_Approval\\PARQ\\last_agent_list.parquet", index=False)

    # Return all five DataFrames
    return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move

def getNewLocalApprovals(url):
    current_la = getLocalApprovals(url)

    # Load old approval list
    old_la_path = "Local_Approval\\PARQ\\last_la.parquet"
    if os.path.exists(old_la_path):
        old_la = pd.read_parquet(old_la_path)
    else:
        old_la = pd.DataFrame(columns=current_la.columns)

    # Create composite keys
    current_la['key'] = current_la['clientid'].astype(str) + "_" + current_la['granted'].astype(str)
    old_la['key'] = old_la['clientid'].astype(str) + "_" + old_la['granted'].astype(str)

    # Find new entries
    new_entries = current_la[~current_la['key'].isin(old_la['key'])]

    # 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