Fixed issue with generating multiple OTP - now prints a nice list for Copy/Paste. Began splitting client / server functionality. Demoted selector and setup to util from service. Fixed some typos.
This commit is contained in:
@@ -23,11 +23,11 @@ import time
|
||||
import dotenv
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from utils.setup import get_base_directory
|
||||
from models.agent import Agent
|
||||
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.scheduler import (
|
||||
from Server.scheduler import (
|
||||
register_function,
|
||||
run_once_job,
|
||||
)
|
||||
|
||||
+26
-188
@@ -16,205 +16,43 @@
|
||||
|
||||
|
||||
import logging
|
||||
from services.agenthandler import selectAgents
|
||||
from utils.utils import colorText
|
||||
from utils.selector import Selector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
"""
|
||||
def getActiveOTP(url):
|
||||
|
||||
endpoint = url + f'/v1/otp/usage'
|
||||
payload = {
|
||||
"status" : "1"
|
||||
}
|
||||
|
||||
headers = {"X-APIKey": load_env('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:
|
||||
formatHTML(otp, f"OTP\\HTML\\newest_active_OTP.html")
|
||||
|
||||
def getOTPActivities(url, otpid):
|
||||
endpoint = url + f'/v1/otp/activities'
|
||||
payload = {"otpid": f"{otpid}"}
|
||||
headers = {"X-APIKey": load_env('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:
|
||||
formatHTML(combined_data, f"OTP/HTML/OTP_activities_{otpid}.html")
|
||||
|
||||
def monitorOTP(url, pups):
|
||||
|
||||
getActiveOTP(url)
|
||||
from services.API import AirlockAPIWrapper
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
if 'otpid' not in current_active_OTP.columns: current_active_OTP = pd.DataFrame(columns=['otpid'])
|
||||
if 'otpid' not in old_active_OTP.columns: old_active_OTP = pd.DataFrame(columns=['otpid'])
|
||||
|
||||
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 * .95)
|
||||
#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}", "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, policyid = clientf.getPolicyFromClientID(url,clientid)
|
||||
|
||||
"""
|
||||
|
||||
# Devices can come out of OTP either by timeout, or by early move out of OTP. If they are manually moved out prior to the job to add hashes can run, we want to accelerate the job.
|
||||
# But first, we want to update the OTP activities one final time for the pid, then move up any jobs if they exist, then add the hashes to the local approval allowlist
|
||||
"""
|
||||
|
||||
getOTPActivities(url,pid)
|
||||
find_and_prioritize_jobs_by_pid(pid, 1)
|
||||
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.localtime()
|
||||
|
||||
if not os.path.exists(f"OTP\\PARQ\\localapprovalhistory.parquet"):
|
||||
df = pd.DataFrame()
|
||||
df.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
|
||||
|
||||
history = pd.read_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
|
||||
history = pd.concat([history, finalhashesadded], ignore_index=True)
|
||||
formatHTML(history, f"localapproval_history.html")
|
||||
|
||||
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 = 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:
|
||||
#TODO add the api call
|
||||
pass
|
||||
|
||||
# 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)
|
||||
|
||||
def generateOTP(url, agentid):
|
||||
def generate(api: AirlockAPIWrapper):
|
||||
agents = selectAgents(api)
|
||||
purpose = input(colorText(" Please enter the purpose for the OTP: ", "white"))
|
||||
|
||||
possible_durations = [15, 60, 360, 1440, 10080]
|
||||
duration_selected = " "
|
||||
|
||||
print(colorText("Please select a duration:", "white"))
|
||||
for i, option in enumerate(possible_durations, start=1):
|
||||
print(f"{i}. {option}")
|
||||
duration_selected = Selector.select_int(possible_durations)
|
||||
otp_dict = {}
|
||||
if duration_selected:
|
||||
for agent in agents:
|
||||
otp_code = api.otp_generate(agent.agentid, duration_selected, purpose)
|
||||
logger.info(f"Generated OTP for {agent.hostname}: {otp_code}")
|
||||
otp_dict[agent.hostname] = otp_code
|
||||
|
||||
try:
|
||||
choice = int(input("Enter the number of your choice: "))
|
||||
if 1 <= choice <= len(possible_durations):
|
||||
duration_selected = possible_durations[choice - 1]
|
||||
print(colorText(f"You selected: {duration_selected}", "yellow"))
|
||||
else:
|
||||
print(colorText("Invalid choice.", "red"))
|
||||
except ValueError:
|
||||
print(colorText("Invalid input. Please enter a number.", "red"))
|
||||
return otp_dict
|
||||
|
||||
endpoint = url + '/v1/otp/retrieve'
|
||||
payload = {
|
||||
"duration" : f"{duration_selected}",
|
||||
"agentid" : f"{agentid}",
|
||||
"purpose" : f"{purpose}"
|
||||
}
|
||||
def otp_activities_by_agent(api: AirlockAPIWrapper):
|
||||
agents = selectAgents(api)
|
||||
otp_dict = {}
|
||||
for agent in agents:
|
||||
otp_info = api.otp_find_by_agent(agent.agentid)
|
||||
otp_dict[agent.hostname] = otp_info
|
||||
|
||||
headers = {"X-APIKey": load_env('APIKEY')}
|
||||
payload = json.dumps(payload)
|
||||
return otp_dict
|
||||
|
||||
|
||||
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
|
||||
result = json.loads(response.text)
|
||||
otpcode = result["response"]["otpcode"]
|
||||
print(colorText(f"The OPT code is: {otpcode}", "yellow"))
|
||||
|
||||
"""
|
||||
def revoke(api: AirlockAPIWrapper):
|
||||
otp_dict = otp_activities_by_agent(api)
|
||||
list_to_revoke = [entry["otpid"] for entry in otp_dict]
|
||||
for revokee in list_to_revoke:
|
||||
api.otp_revoke(revokee)
|
||||
+3
-3
@@ -24,7 +24,7 @@ import pandas as pd
|
||||
from models.execution import ExecutionHistoryRecord, Hash
|
||||
from models.policy import Allowlist, Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.selector import Selector
|
||||
from utils.selector import Selector
|
||||
from utils.utils import (
|
||||
colorText,
|
||||
formatHTML,
|
||||
@@ -42,10 +42,10 @@ dotenv.load_dotenv()
|
||||
|
||||
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
|
||||
|
||||
allowlists = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||
logger.debug("Prompting for Policies")
|
||||
print(colorText("Please select policy/policies", "white"))
|
||||
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
|
||||
selected = Selector.select_objects(policies, allow_multiple, prompt_each=True)
|
||||
|
||||
if selected is None:
|
||||
return []
|
||||
|
||||
+1
-1
@@ -21,7 +21,7 @@ import pandas as pd
|
||||
from flows.prepPolicy import selectPolicies
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import getPolicyInfo
|
||||
from services.selector import Selector
|
||||
from utils.selector import Selector
|
||||
from utils.utils import colorText, load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
Reference in New Issue
Block a user