AirlockTools/utils/clientfunctions.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.hashfunctions as hashf
import utils.utils as ct
import utils.pathfunctions as pathf
import utils.policyfunctions as policyf

#Standard Libary Imports:
import datetime
import json
import os
import re

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

def findAgentID(url):
       
    print(ct.colorText("WARNING: Device Name is Case Sensitive", "red"))
    hostname = input(ct.colorText("Enter Device Name: ", "white"))
    
    endpoint = url + '/v1/agent/find'
    payload = {
            "hostname" : f"{hostname}"
    }

    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)
    data = pd.DataFrame(result["response"]["agents"])
    return data.loc[0, "agentid"]

def getDestAllowlistFromClientID(url, clientid):

    #This is dependant on their being an allowlist in the policy containing the name "localapproval"

    endpoint = url + '/v1/agent/find'
    payload = {
            "agentid" : f"{clientid}"
    }

    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)
    data = pd.DataFrame(result["response"]["agents"])

    allowlists = getPolicyAllowlists(url,data.loc[0, "groupid"]) 
 
    matches = allowlists.loc[
        allowlists['name'].str.contains('local', case=False, na=False) &
        allowlists['name'].str.contains('approval', case=False, na=False),
        'applicationid'
    ].values
    app_id = matches[0] if len(matches) > 0 else None
    
    return app_id

def getPolicyFromClientID(url, clientid):

    endpoint = url + '/v1/agent/find'
    payload = {
            "agentid" : f"{clientid}"
    }

    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)
    data = pd.DataFrame(result["response"]["agents"])

    policy = getPolicyName(url,data.loc[0, "groupid"]) 
  
    return policy
    
def getPolicyAllowlists(url, groupid):

    endpoint = url + '/v1/group/policies'
    payload = {
            "groupid" : f"{groupid}"
    }

    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)
    data = pd.DataFrame(result["response"]["applications"])
    return data

def getPolicyName(url, groupid):
    endpoint = url + '/v1/group/'
    payload = {}

    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)

    data = pd.DataFrame(result["response"]["groups"])

    # Safely attempt to get the policy name
    filtered = data.loc[data['groupid'] == str(groupid), 'name']
    if not filtered.empty:
        name = filtered.values[0]
    else:
        name = None  # or "Unknown", depending on your preference

    return name


def devicehistory(url, outputjson: bool):
    endpoint = url + '/v1/getexechistory'
    print("\n")
    print(ct.colorText("1. Today", "yellow"))
    print(ct.colorText("2. Last 24 Hours", "yellow"))
    print(ct.colorText("3. Past 7 Days", "yellow"))
    print(ct.colorText("4. Past 30 Days", "yellow"))
    print(ct.colorText("5. Custom Date Range","yellow"))
    choice = input(ct.colorText("\nSelect Date Range: ", "white"))
    today = datetime.date.today()
    today = today.strftime("%Y-%m-%d")
    date_selected = " "
    if choice == '1':
        date_selected = today
    elif choice == '2':
        date_selected = datetime.date.today() - datetime.timedelta(days=1)
        date_selected = date_selected.strftime('%Y-%m-%d')
    elif choice == '3':
        date_selected = datetime.date.today() - datetime.timedelta(days=7)
        date_selected = date_selected.strftime('%Y-%m-%d')
    elif choice == '4':
        date_selected = datetime.date.today() - datetime.timedelta(days=30)
        date_selected = date_selected.strftime('%Y-%m-%d')
    elif choice == "5":
        print(ct.colorText("Please Input Dates as YYYY-MM-DD", "cyan"))
        date_selected = input(ct.colorText("From: ", "white"))
        today = input(ct.colorText("Date To: ", "white"))
    print(ct.colorText("WARNING: Device Name is Case Sensitive", "red"))
    device = input(ct.colorText("Enter Device Name: ", "white"))
    payload_dict = {
        "datefrom": date_selected,
        "dateto": today,
        "hostname": device
    }
    payload = json.dumps(payload_dict)
    print(ct.colorText(payload, "green"))
    headers = {
        "X-APIKey": os.getenv('APIKEY')
    }

    response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)

    if outputjson:
        return response

    parse_text = json.loads(response.text)

    # Safely get exechistory
    exechistory = parse_text.get('response', {}).get('exechistory')

    if isinstance(exechistory, list):
        for block in exechistory:
            print(ct.colorText(f"Command: {block.get('commandline', 'N/A')}", "green"))
            print(ct.colorText(f"Date: {block.get('datetime', 'N/A')}", "green"))
            print(ct.colorText(f"Filename: {block.get('filename', 'N/A')}", "green"))
            print(ct.colorText(f"Policy Name: {block.get('policyname', 'N/A')}", "green"))
            print(ct.colorText(f"Hostname: {block.get('hostname', 'N/A')}", "green"))
            print(ct.colorText(f"Hash: {block.get('sha256', 'N/A')}", "green"))
            print("\n")
    else:
        print(ct.colorText("No execution history found or data is not in expected format.", "red"))
        
def findAllAgents(url):        
    endpoint = url + '/v1/agent/find'
    payload = {}
    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)
    data = pd.DataFrame(result["response"]["agents"])
    
    group_ids = sorted(data['groupid'].unique().tolist())
    group_policy_map = {}

    for groupid in group_ids:
        policy_name = getPolicyName(url, groupid)
        group_policy_map[groupid] = policy_name

    data['policy_name'] = data['groupid'].map(group_policy_map)

    status_map = {
        0: 'Offline',
        1: 'Online',
        2: 'Hidden',
        3: 'Safemode'
    }

    data['status'] = data['status'].map(status_map)
    return(data)

def findAgents(url, device_input_str, return_dataframe):

    os.makedirs("device_search", exist_ok=True)
    df = findAllAgents(url)

    #Parse the input string into device names (newline-separated only)
    device_names = device_input_str.strip().split('\n')
    device_names = [name.strip() for name in device_names if name.strip()]

    #Build a regex pattern for case-insensitive matching
    pattern = '|'.join([re.escape(name) for name in device_names])
    regex = re.compile(pattern, re.IGNORECASE)

    #Filter the DataFrame using regex
    matched_df = df[df['hostname'].apply(lambda x: bool(regex.search(str(x))))]
    
    if return_dataframe : return matched_df
    else:
        #Export to CSV with timestamp
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        filename = f"agentsearch_{timestamp}.csv"
        matched_df.to_csv(f"device_search\\{filename}", index=False)

        print(ct.colorText(f"\n✅ Matched devices exported to: device_search\\{filename}","green"))
    
def promptForDevices():
    
    print(ct.colorText("🔍 Device Search", "cyan"))
    print(ct.colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
    print(ct.colorText("When you're done, press Enter twice.\n", "cyan"))
    print(ct.colorText("Example:", "cyan"))
    print(ct.colorText("H00000", "cyan"))
    print(ct.colorText("UTN00000", "cyan"))
    print(ct.colorText("i-hSuperSecretServer", "cyan"))
    print(ct.colorText("u-hVenderBroke\n", "cyan"))


    # Collect multiline input from user
    print(ct.colorText("Paste or type your device names below:", "white"))
    device_input_lines = []
    while True:
        line = input()
        if line == "":
            break
        device_input_lines.append(line)

    device_input_str = "\n".join(device_input_lines)

    return device_input_str


def returnToEnforcement(url, device_df, policy_relationship_map, bad_publisher_list, pups, badpathparts, threat_tolerance_constant, path_exclusion_constant, min_files_for_path):
    policylist = sorted(device_df['policy_name'].unique().tolist())
    type = [1, 2, 6, 7]
  
    parq_base_dir = "prepare_policy\\parquet\\"
    appr_base_dir = "prepare_policy\\approved\\"
    needappr_base_dir = "prepare_policy\\needs_approved"
    pflight_base_dir = "prepare_policy\\preflight"
    
    #If the directorys where we're going to store our output dont exist, make them. 
    os.makedirs(parq_base_dir, exist_ok=True)
    os.makedirs(needappr_base_dir, exist_ok=True)
    os.makedirs(appr_base_dir, exist_ok=True)
    os.makedirs(pflight_base_dir, exist_ok=True)

    while True:
        
        #ct.printDeviceEnforceChecklist()
        
        choice = input(ct.colorText("\nEnter your choice: ", "white"))
        
        if choice == "1":
            
            policyf.buildExecHistory(url,
                        policylist, 
                        parq_base_dir, 
                        needappr_base_dir, 
                        type, 
                        threat_tolerance_constant,
                        bad_publisher_list,
                        pups,
                        )
            
        elif choice == "2":
            
            pathf.generateApprovables(needappr_base_dir, appr_base_dir, parq_base_dir, badpathparts, bad_publisher_list, path_exclusion_constant, min_files_for_path, True)
                
        elif choice == "3":
            policyf.savePreFlights(appr_base_dir, parq_base_dir, pflight_base_dir)
                     
        """

        elif choice == "4":
            if os.path.exists(f"preflight\\final_path_exclusions.html") and os.path.exists(f"preflight\\final_hash_approvals.html")    and allowlist_parent_name != " " and allowlist_child_name != " " and destination_name != " ":
                sendToPolicyTest(
                    url,
                    first_policy,
                    second_policy,
                    destination_name,
                    destination_id,
                    allowlist_parent_name,
                    allowlist_parent_id, 
                    allowlist_child_name,
                    allowlist_child_id
                    )

        elif choice == "5":
            if os.path.exists(f"preflight\\final_path_exclusions.html") and os.path.exists(f"preflight\\final_hash_approvals.html")    and allowlist_parent_name != " " and allowlist_child_name != " " and destination_name != " ":
                sendToPolicy(
                    url,
                    first_policy,
                    second_policy,
                    destination_name,
                    destination_id,
                    allowlist_parent_name,
                    allowlist_parent_id, 
                    allowlist_child_name,
                    allowlist_child_id
                    )
        elif choice == "R":
            
            pathf.clean_folders(enforcement_prep)

        elif choice == "Q":
            break
        else:
            print(ct.colorText("Invalid choice. Please try again.", "red"))    
    """
        

def findQuietAgents(url):
    # Get policy selection and agent list
    choice, policynames, policyid = policyf.listPolicies(url)
    policy = policynames[choice]
    groupid = policyid[choice]
    agents = findGroupAgents(url, groupid)

    # Prompt user for history range
    while True:
        try:
            history_days = int(input("Enter how many days of history to pull (1–150): "))
            if 1 <= history_days <= 150:
                break
            else:
                print("Invalid input. Please enter a number between 1 and 150.")
        except ValueError:
            print("Invalid input. Please enter a valid integer.")

    while True:
        try:
            required_quiet = int(input("Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): "))
            if 1 <= required_quiet <= 365:
                break
            else:
                print("Invalid input. Please enter a number between 1 and 365.")
        except ValueError:
            print("Invalid input. Please enter a valid integer.")

    # Get policy execution history
    policy_exec_history = policyf.getPolicyInfo(url, policy, [1, 2, 6, 7], history_days, False)

    # Convert 'datetime' column to timezone-aware datetime objects
    policy_exec_history['datetime'] = pd.to_datetime(policy_exec_history['datetime'], format='%Y-%m-%dT%H:%M:%SZ', utc=True)

    # Get current UTC time
    now = datetime.datetime.now(datetime.timezone.utc)

    # Calculate days ago
    policy_exec_history['days_ago'] = policy_exec_history['datetime'].apply(lambda dt: (now - dt).days)

    # Count total executions per hostname
    hostname_counts = policy_exec_history['hostname'].value_counts()

    # Map execution counts to agents
    agents['execution_count'] = agents['hostname'].map(hostname_counts).fillna(0).astype(int)

    # Find most recent execution per hostname
    most_recent_exec = policy_exec_history.sort_values(by='days_ago').drop_duplicates(subset='hostname', keep='first')

    # Map most recent execution age to agents
    agents['days_since'] = agents['hostname'].map(most_recent_exec.set_index('hostname')['days_ago'])

    #Check for enforcement readyness
    agents['required_quiet'] = required_quiet
    agents['enforce_ready'] = agents['days_since'].apply(
        lambda x: True if pd.isna(x) or x > required_quiet else False
    )

    # Sort agents by execution count and hostname
    agents = agents.sort_values(by=['execution_count', 'hostname'], ascending=[True, True])

    # Save to CSV
    print(ct.colorText(f"Saving CSV to {policy}_agents_last_{history_days}_days.csv", "green"))
    agents.to_csv(f"{policy}_agents_last_{history_days}_days.csv", index=False)

    # Summary stats
    zero_count = (agents['execution_count'] == 0).sum()
    total_hosts = len(agents)
    zero_percentage = (zero_count / total_hosts) * 100

    print(f"Number of hosts in policy: {total_hosts}")
    print(f"Number of hosts with execution_count = 0: {zero_count}")
    print(f"Percentage of hosts with execution_count = 0: {zero_percentage:.2f}%")


def findGroupAgents(url, groupid):        
    endpoint = url + '/v1/agent/find'
    payload = {
        "groupid" : f"{groupid}"
    }
    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)
    data = pd.DataFrame(result["response"]["agents"])
    group_ids = sorted(data['groupid'].unique().tolist())
    group_policy_map = {}

    for groupid in group_ids:
        policy_name = getPolicyName(url, groupid)
        group_policy_map[groupid] = policy_name

    data['policy_name'] = data['groupid'].map(group_policy_map)

    status_map = {
        0: 'Offline',
        1: 'Online',
        2: 'Hidden',
        3: 'Safemode'
    }

    data['status'] = data['status'].map(status_map)
    return(data)
    return(data)