#    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"])

    allowlist = policyf.getDestAllowlist(url,data.loc[0, "groupid"]) 
 
    return allowlist

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_name = getPolicyName(url,data.loc[0, "groupid"]) 
    policy_id   = data.loc[0, "groupid"]
    return policy_name, policy_id
    
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 findQuietAgents(url):
    # Get policy selection and agent list
    choice, policynames, policyid = policyf.choosePolicies(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 statistics
    total_agents = len(agents)
    ready_agents = agents['enforce_ready'].sum()
    not_ready_agents = total_agents - ready_agents
    ready_percentage = (ready_agents / total_agents) * 100

    # Print results
    print(f"Total agents: {total_agents}")
    print(f"Agents marked as 'enforce_ready': {ready_agents}")
    print(f"Agents not ready: {not_ready_agents}")
    print(f"Percentage ready for enforcement: {ready_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)

def moveAgentToAudit(url, agentid, policy_relationship_map):
    policy_name, policy_id = getPolicyFromClientID(url, agentid)

    print(policy_id)


    if policy_id in policy_relationship_map:
        target_policy = policy_relationship_map[policy_id]
    elif policy_id in policy_relationship_map.values():
        print(f"Agent {agentid} is already in an audit group. No action needed.")
        return
    else:
        print(f"Error: No corresponding audit policy found for policy: {policy_name} - {policy_id}.")
        return

    moveAgent(url, agentid, target_policy, "audit")

def moveAgentToEnforcement(url, agentid, policy_relationship_map):
    policy_name, policy_id = getPolicyFromClientID(url, agentid)

    # Invert the map for audit → enforcement
    inverse_map = {v: k for k, v in policy_relationship_map.items()}

    if policy_id in inverse_map:
        target_policy = inverse_map[policy_id]
    elif policy_id in inverse_map.values():
        print(f"Agent {agentid} is already in an enforcement group. No action needed.")
        return
    else:
        print(f"Error: No corresponding enforcement policy found for policy: {policy_name} - {policy_id}.")
        return

    moveAgent(url, agentid, target_policy, "enforcement")

def moveAgent(url, agentid, target_policy, direction):
    endpoint = f"{url}/v1/agent/move"
    payload = {
        "groupid": target_policy,
        "agentid": agentid
    }

    headers = {"X-APIKey": os.getenv('APIKEY')}
    response = None  # Initialize to avoid unbound errors

    try:
        response = requests.post(endpoint, headers=headers, data=json.dumps(payload), verify=False)
        response.raise_for_status()  # Raises HTTPError for bad status codes

        result = response.json()

        # Check if 'error' key exists and if it's not a success message
        if "error" in result and result["error"].lower() != "success":
            print(f"API returned an error: {result['error']}")
        else:
            print(f"✅ Agent {agentid} successfully moved to {direction} group {target_policy}.")

    except requests.exceptions.HTTPError as http_err:
        print(f"HTTP error occurred: {http_err}")
        if response is not None:
            print("Raw response:", response.text)
    except requests.exceptions.RequestException as req_err:
        print(f"Request error occurred: {req_err}")
    except ValueError:
        print("Failed to parse JSON response.")
        if response is not None:
            print("Raw response:", response.text)
    except Exception as e:
        print(f"Unexpected error: {e}")
        if response is not None:
            print("Raw response:", response.text)