Files
AirlockTools/AirlockTools.py
T
2025-09-09 17:06:38 -04:00

283 lines
10 KiB
Python

# 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 argparse
import dotenv
import os
import pandas as pd
import urllib3
import utils.allowlist
import utils.clientfunctions
import utils.getdeviceevents
import utils.hashfunctions
import utils.otpfunctions
import utils.pathfunctions
import utils.policyfunctions
import utils.pretty as ct
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
dotenv.load_dotenv()
#Constants
url = os.getenv('url')
bad_publisher_list = ["Brave","Zoom", "GlavSoft", "VNC"]
pups = ["logmein", "invalid"]
badpathparts = ["users", "wwwroot", "windows\\temp", "windows\\task", "windows\\system32", "startup", "windows\\fonts", "Recycle.Bin", "AppData", "programdata"]
path_exclusion_constant = 4
min_files_for_path = 4
threat_tolerance_constant = 4
def main():
parser = argparse.ArgumentParser(description="Your script description")
parser.add_argument('--monitorOTP', action='store_true', help='Run in non-interactive mode')
# Add other arguments as needed
args = parser.parse_args()
if args.monitorOTP:
# Non-interactive logic
if not os.path.exists("parquet"): os.makedirs("parquet")
apivalidation()
print(f"Running non-interactively to start monitoring OTP")
# Process input here
else:
# Interactive logic
apivalidation()
menu_main()
def apivalidation():
match os.getenv('APIKEY'):
case '':
print(ct.colorText("Please add your API Key to the .env file", "red"))
def tryToReadCSV(csv):
try:
df =pd.read_csv(csv)
if df.empty:
print(ct.colorText("Error: CSV file has headers but no data rows.", "red"))
else:
print(ct.colorText(f"Data loaded successfully from {csv}", "green"))
except pd.errors.EmptyDataError:
print(ct.colorText("Notice : CSV file is completely empty (no headers, no data), falling back to empty frame", "white"))
df = pd.DataFrame() # Create an empty DataFrame as fallback
return df
def tryToReadParquet(parquet):
try:
df = pd.read_parquet(parquet)
if df.empty:
print(ct.colorText("Error: Parquet file has headers but no data rows.", "red"))
else:
print(ct.colorText(f"Data loaded successfully from {parquet}", "green"))
except pd.errors.EmptyDataError:
print(ct.colorText("Notice : Parquet file is completely empty (no headers, no data), falling back to empty frame", "white"))
df = pd.DataFrame() # Create an empty DataFrame as fallback
return df
def deduplicate_list(lst):
seen = set()
return [x for x in lst if not (x in seen or seen.add(x))]
def menu_main():
while True:
ct.displayIntro();
print(ct.colorText("1. Get All Events for Single Device", "yellow"))
print(ct.colorText("2. Placeholder for Local Approval", "yellow"))
print(ct.colorText("3. Placeholder for Another Tool", "yellow"))
print(ct.colorText("4. Prepare Policy For Enforcement", "yellow"))
print(ct.colorText("Q. Quit", "yellow"))
choice = input(ct.colorText("\nEnter Menu Item: ", "white"))
if choice == '1':
utils.getdeviceevents.devicehistory(url,False)
elif choice == "2":
menu_local_approve()
elif choice == "3":
menu_feature2()
elif choice == "4":
menu_prepare_to_enforce()
elif choice == "Q":
break
else:
print(ct.colorText("Invalid choice. Please try again.","red"))
def menu_local_approve():
while True:
print("\n--- Submenu ---")
print("1. Sub-option A")
print("2. Sub-option B")
print("3. Return to Main Menu")
choice = input("Enter your choice: ")
if choice == "1":
print("You selected Sub-option A")
elif choice == "2":
print("You selected Sub-option B")
elif choice == "3":
print("Returning to Main Menu...")
break
else:
print("Invalid choice. Please try again.")
def menu_feature2():
while True:
print("\n--- Submenu ---")
print("1. Sub-option A")
print("2. Sub-option B")
print("3. Return to Main Menu")
choice = input("Enter your choice: ")
if choice == "1":
print("You selected Sub-option A")
elif choice == "2":
print("You selected Sub-option B")
elif choice == "3":
print("Returning to Main Menu...")
break
else:
print("Invalid choice. Please try again.")
def menu_prepare_to_enforce():
first_policy = " "
second_policy = " "
destination_name = " "
destination_id = " "
allowlist_parent_name = " "
allowlist_parent_id = " "
allowlist_child_name = " "
allowlist_child_id = " "
#If the directorys where we're going to store our output dont exist, make them.
if not os.path.exists("parquet"): os.makedirs("parquet")
if not os.path.exists("needs_approved"): os.makedirs("needs_approved")
if not os.path.exists("approved"): os.makedirs("approved")
if not os.path.exists("preflight"): os.makedirs("preflight")
while True:
ct.printEnforceChecklist(first_policy, second_policy, allowlist_child_name, allowlist_parent_name, destination_name)
choice = input(ct.colorText("\nEnter your choice: ", "white"))
if choice == "1":
choice, policynames, policyid = utils.allowlist.listPolicies(url)
first_policy = policynames[choice]
while True:
answer = input(ct.colorText(f"{"Do you want to load a second policy?"} (yes/no): ", "white").strip().lower())
if answer in ("yes", "y"):
choice, policynames, policyid = utils.allowlist.listPolicies(url)
second_policy = policynames[choice]
break
elif answer in ("no", "n"):
second_policy = first_policy
break
else:
print(ct.colorText("Please answer with 'yes' or 'no'.", "red"))
elif choice == "2":
if not os.path.exists(f"parquet\\execution_history_{first_policy}.parquet"):
utils.policyfunctions.getPolicyInfo(url, first_policy, 60)
if not os.path.exists(f"parquet\\execution_history_{second_policy}.parquet"):
utils.policyfunctions.getPolicyInfo(url, second_policy, 60)
if not os.path.exists(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet"):
utils.hashfunctions.combineHashes(url, first_policy, second_policy)
if not os.path.exists(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet"):
utils.hashfunctions.categorizeHashes(
first_policy,
second_policy,
pd.read_parquet(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet"),
threat_tolerance_constant,
bad_publisher_list,
pups
)
if not os.path.exists(f"parquet\\condensed_executions_{first_policy}_{second_policy}.parquet"):
utils.hashfunctions.condenseExecutions(first_policy,second_policy)
if os.path.exists(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet") & os.path.exists(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet") & os.path.exists(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet"):
utils.hashfunctions.divideSortedHashExecutions(first_policy,second_policy,pups)
elif choice == "3":
if os.path.exists(f"approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv") and os.path.exists(f"approved\\hashes_rep_good_{first_policy}_{second_policy}.csv"):
utils.pathfunctions.generatePathReview(first_policy, second_policy, badpathparts, min_files_for_path)
else:
print(ct.colorText(f"Please manually approve hashes prior to this step","red"))
elif choice == "4":
if os.path.exists(f"approved\\path_needs_approved_{first_policy}_{second_policy}.csv"):
if not os.path.exists(f"parquet\\final_hash_approvals_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\final_path_exclusions_{first_policy}_{second_policy}.parquet"):
utils.hashfunctions.generatePreflights(first_policy, second_policy)
elif choice == "5":
print(ct.colorText(f"Please choose destination_name Policy for Path Exclusions","white"))
choice, policynames, policyid = utils.allowlist.listPolicies(url)
#print(allowlist_parent_tuple)
destination_name = policynames[choice]
destination_id = policyid[choice]
print(ct.colorText(f"Please choose Parent Allowlist for Known Hashes","white"))
choice, allowlists,allowid = utils.allowlist.listAllowlists(url)
#print(allowlist_parent_tuple)
allowlist_parent_name = allowlists[choice]
allowlist_parent_id = allowid[choice]
print(ct.colorText(f"Please choose Child Allowlist for Less-Known Hashes","white"))
choice, allowlists, allowid = utils.allowlist.listAllowlists(url)
#print(allowlist_child_tuple)
allowlist_child_name = allowlists[choice]
allowlist_child_id = allowid[choice]
elif choice == "6":
if os.path.exists(f"preflight\\final_path_exclusions_{first_policy}_{second_policy}.html") and os.path.exists(f"preflight\\final_hash_approvals_{first_policy}_{second_policy}.html") and allowlist_parent_name != " " and allowlist_child_name != " " and destination_name != " ":
utils.policyfunctions.sendToPolicy(
url,
first_policy,
second_policy,
destination_name,
destination_id,
allowlist_parent_name,
allowlist_parent_id,
allowlist_child_name,
allowlist_child_id
)
elif choice == "Q":
break
else:
print(ct.colorText("Invalid choice. Please try again.", "red"))
if __name__ == "__main__":
main()