AirlockTools/utils/allowlist.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/>.
import datetime
import requests
import json
import os
import utils.pretty as ct
import ijson
import os
from bson import ObjectId
import datetime

def pullPolicyExechistories(url, policiesnames, days, outputjson: bool):
   
    file_path = 'chunkinator.json'

    # If chunkintor exists - kill and make new
    if os.path.exists(file_path):
       os.remove(file_path)

    if not os.path.exists(file_path):
        with open(file_path, 'w') as file:
            json.dump({'error': 'Success', 'response': {'exechistories': []}}, file)
        print(f"File '{file_path}' has been created.")
    else:
        print(f"File '{file_path}' already exists.")

    headers = {"X-APIKey": os.getenv('APIKEY')}

    checkpoint = str(skipback(days))
    json_output = {'error': 'Success', 'response': {'exechistories': []}}

    while True:
        json_response_data = checkpoint_stomper(checkpoint, url, policiesnames, headers)
        histories = json_response_data['response']['exechistories']
        if not histories:
            break

        array_dividend = max(round(len(histories) / 20), 1)
        match_found = False

        for index, item in enumerate(histories[::array_dividend]):
            item_date = datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()
            if item_date >= datetime.date.today() - datetime.timedelta(days):
                match_found = True
                break

        checkpoints_processed = round(len(histories) / array_dividend)
        print(ct.colorText(
            f"{index + 1}/{checkpoints_processed} checkpoint(s) processed. " # type: ignore
            f"{'Found with Date Match.' if match_found else ''} Last Checkpoint: {item['checkpoint']}.", "blue")) # type: ignore

        checkpoint = item['checkpoint'] # type: ignore

        if match_found:
            for item in histories:
                item_date = datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()
                if item_date >= datetime.date.today() - datetime.timedelta(days):
                    json_output['response']['exechistories'].append(item)

            # Deduplicate and write to file
            seen = {}
            if os.path.exists(file_path):
                with open(file_path, 'r') as file:
                    existing_data = json.load(file)
                combined = existing_data['response']['exechistories'] + json_output['response']['exechistories']
            else:
                combined = json_output['response']['exechistories']

            for item in combined:
                key = (item.get('sha256'), item.get('filename'), item.get('hostname'))
                seen[key] = item

            deduplicated = list(seen.values())
            with open(file_path, 'w') as file:
                json.dump({'error': 'Success', 'response': {'exechistories': deduplicated}}, file)

            # Reset output to free memory
            json_output['response']['exechistories'].clear()

    # Final output
    with open(file_path, 'r') as file:
        final_output = json.load(file)
    os.remove(file_path)

    return json.dumps(final_output) if outputjson else None


def checkpoint_stomper(checkpoint, url, policy, headers):
    json_output = {'error': 'Success', 'response': {'exechistories': []}}
    endpoint = url + '/v1/logging/exechistories'
    payload_dict = {
        "type":[1,2,6,7],
        "checkpoint": checkpoint,
        "policy": [policy]
    }
    payload = json.dumps(payload_dict)
    with requests.request("POST", endpoint, headers=headers, data=payload, verify=False, stream=True) as response:
        parser = ijson.items(response.raw, 'response.exechistories.item')
        for item in parser:
            key = (item.get('sha256'), item.get('hostname'))
            if key not in json_output:
                json_output['response']['exechistories'].append(item)   
    parse_text = json.loads(json.dumps(json_output))
    return parse_text

def listPolicies(url): 
    endpoint = url + '/v1/group'
    print(ct.colorText("[+] Grabbing All Policies", "cyan"))
    payload = {}
    headers = {
        "X-APIKey": os.getenv('APIKEY')
    }
    response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
    parse_text = json.loads(response.text)
    policiesnames = []
    policyids = []
    for index, list in enumerate(parse_text['response']['groups'], start=1):
        print(ct.colorText(f"{index}. {list['name']}", "yellow"))
        policiesnames.append(list['name'])
        policyids.append(list['groupid'])
    choice = input(ct.colorText("Select Policy Group: ", "white"))
    choice = int(choice) - 1
    return choice, policiesnames, policyids

def listAllowlists(url): 
    endpoint = url + '/v1/application'
    print(ct.colorText("[+] Grabbing All Allowlists", "cyan"))
    payload = {}
    headers = {
        "X-APIKey": os.getenv('APIKEY')
    }
    response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
    parse_text = json.loads(response.text)
    policiesnames = []
    policyids = []
    for index, list in enumerate(parse_text['response']['applications'], start=1):
        if index >= 38:
            print(ct.colorText(f"{index}. {list['name']}", "yellow"))
            policiesnames.append(list['name'])
            policyids.append(list['applicationid'])
    choice = int(input(ct.colorText("Select allowlist: ", "white")))
    if choice < 38:
        print(ct.colorText("Please only choose an allowlist designed for this use - '38+'","red"))
    elif choice >= 38:
        choice = choice - 38
        return choice, policiesnames, policyids
    #Need else and catch for upper bound

def skipback(days):
    """
    Generate a MongoDB ObjectId for a given number of days ago from today.
    Adds 1 extra day to the input to look further back.
    """
    adjusted_days = days + 1
    date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=adjusted_days)
    timestamp = int(date_days_ago.timestamp())
    hex_timestamp = format(timestamp, '08x')
    objectid_hex = hex_timestamp + '0000000000000000'
    return ObjectId(objectid_hex)