Added Move Agent, UI Tweaks

Rust/services.rs

@@ -0,0 +1,65 @@
+use std::str::FromStr;
+
+use pyo3::prelude::*;
+use reqwest::{Client, header::{HeaderMap, HeaderName, HeaderValue}};
+use serde::{Deserialize, Serialize};
+use serde_json::{Value, json};
+
+#[pyfunction]
+pub fn pull_policy_exec_histories() {
+    println!("Hello World from Rust!");
+}
+
+#[pyfunction]
+pub fn api(py: Python<'_>, x: Py<PyAny>) {
+    let base_url: String = x.getattr(py, "base_url").unwrap().to_string();
+    println!("{}", base_url);
+}
+
+#[tokio::main]
+#[pyfunction]
+pub async fn history_logging(py: Python<'_>, py_self: Py<PyAny>, exec_types: String, checkpoint_number: String, policy_names: String ) -> String {
+    #[derive(Serialize, Debug)]
+    struct Payload {
+        checkpoint: String,
+    }
+    let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
+    let headers = py_self.getattr(py, "headers").unwrap().to_string();
+    let headers_replace = headers.replace('\'', "\"");
+    let parsed: Value = serde_json::from_str(&headers_replace.as_str()).unwrap();
+    let mut header_map = HeaderMap::new();
+    if let Some(obj) = parsed.as_object() {
+        for (key, value) in obj {
+            if let Some(v) = value.as_str() {
+                let val = HeaderValue::from_str(v).unwrap();
+                header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val);
+            }
+        }
+    }
+    let exec_types = exec_types.replace(" ", "");
+    let payload_dict: Payload = Payload {
+        checkpoint:exec_types.to_string(),
+    };
+    println!("{:?}", payload_dict);
+    let client = Client::builder()
+        .danger_accept_invalid_certs(true)
+        .default_headers(header_map)
+        .timeout(std::time::Duration::from_secs(30))
+        .build()
+        .unwrap();
+    let res = client
+        .post(format!("{}/v1/logging/exechistories", base_url))
+        .json(&payload_dict)
+        .send()
+        .await;
+    match res {
+        Ok(res) => {
+            println!("{:?}", res.text().await.unwrap());
+        }
+        Err(e) => {
+            eprintln!("{}", e);
+        }
+    } 
+    let testingstring: String = "Testing".to_string();
+    return testingstring
+}

flows/quietAgent.py

@@ -23,7 +23,7 @@ from flows.prepPolicy import selectPolicies
 from services.API import AirlockAPIWrapper
 from services.policyhandler import getPolicyInfo
 from utils.selector import Selector
-from utils.utils import colorText, load_env
+from utils.utils import colorText, get_sanitized_input, load_env
 
 logger = logging.getLogger(__name__)
 
@@ -50,75 +50,79 @@ def findQuietAgents(api: AirlockAPIWrapper):
             valid_range=(1, 150),
         )
 
+        confirm = Selector.confirm(f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N :  ")
         # Get execution history as a DataFrame
-        policy_exec_history = getPolicyInfo(
+        if confirm:
+            policy_exec_history = getPolicyInfo(
             api, selected_policy[0], [1, 2, 6, 7], history_days
-        )
+            )
         
        
-        if policy_exec_history.empty:
+            if policy_exec_history.empty:
-            logging.info("No execution history found for the selected policy and time range.")
+                logging.info("No execution history found for the selected policy and time range.")
-            return
+                get_sanitized_input("Press enter to continue")
+                return
 
 
-        # Convert 'datetime' column to timezone-aware datetime objects
+            # Convert 'datetime' column to timezone-aware datetime objects
-        policy_exec_history["datetime"] = pd.to_datetime(
+            policy_exec_history["datetime"] = pd.to_datetime(
-            policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
+                policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
-        )
+            )
 
-        # Get current UTC time
+            # Get current UTC time
-        now = datetime.datetime.now(datetime.timezone.utc)
+            now = datetime.datetime.now(datetime.timezone.utc)
 
-        # Calculate days ago
+            # Calculate days ago
-        policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
+            policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
-            lambda dt: (now - dt).days
+                lambda dt: (now - dt).days
-        )
+            )
 
-        # Count total executions per hostname
+            # Count total executions per hostname
-        hostname_counts = policy_exec_history["hostname"].value_counts()
+            hostname_counts = policy_exec_history["hostname"].value_counts()
 
-        # Map execution counts to agents
+            # Map execution counts to agents
-        agents["execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int)
+            agents["execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int)
 
-        # Find most recent execution per hostname
+            # Find most recent execution per hostname
-        most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates(
+            most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates(
-            subset="hostname", keep="first"
+                subset="hostname", keep="first"
-        )
+            )
 
-        # Map most recent execution age to agents
+            # Map most recent execution age to agents
-        agents["days_since"] = agents["hostname"].map(
+            agents["days_since"] = agents["hostname"].map(
-            most_recent_exec.set_index("hostname")["days_ago"]
+                most_recent_exec.set_index("hostname")["days_ago"]
-        )
+            )
 
-        # Check for enforcement readiness
+            # Check for enforcement readiness
-        agents["required_quiet"] = required_quiet
+            agents["required_quiet"] = required_quiet
-        agents["enforce_ready"] = agents["days_since"].apply(
+            agents["enforce_ready"] = agents["days_since"].apply(
-            lambda x: True if pd.isna(x) or x > required_quiet else False
+                lambda x: True if pd.isna(x) or x > required_quiet else False
-        )
+            )
 
-        # Sort agents by execution count and hostname
+            # Sort agents by execution count and hostname
-        agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True])
+            agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True])
 
-        # Save to CSV
+            # Save to CSV
-        filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
+            filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
-        logging.debug(f"Saving CSV to {filename}")
+            logging.debug(f"Saving CSV to {filename}")
-        print(colorText(f"Saving CSV to {filename}", "green"))
+            print(colorText(f"Saving CSV to {filename}", "green"))
-        agents.to_csv(filename, index=False)
+            agents.to_csv(filename, index=False)
 
-        # Summary statistics
+            # Summary statistics
-        total_agents = len(agents)
+            total_agents = len(agents)
-        ready_agents = agents["enforce_ready"].sum()
+            ready_agents = agents["enforce_ready"].sum()
-        not_ready_agents = total_agents - ready_agents
+            not_ready_agents = total_agents - ready_agents
-        ready_percentage = (ready_agents / total_agents) * 100
+            ready_percentage = (ready_agents / total_agents) * 100
 
-        # Print results
+            # Print results
+            
         
-    
+            message = (
-        message = (
+                f"Total agents: {total_agents}\n"
-            f"Total agents: {total_agents}\n"
+                f"Agents marked as 'enforce_ready': {ready_agents}\n"
-            f"Agents marked as 'enforce_ready': {ready_agents}\n"
+                f"Agents not ready: {not_ready_agents}\n"
-            f"Agents not ready: {not_ready_agents}\n"
+                f"Percentage ready for enforcement: {ready_percentage:.2f}%"
-            f"Percentage ready for enforcement: {ready_percentage:.2f}%"
+            )
-        )
+            logger.debug(message)
-        logger.debug(message)
+            colorText(message,"green")
-        colorText(message,"green")
+    get_sanitized_input("Press enter to continue")

services/agenthandler.py

@@ -24,6 +24,7 @@ from typing import List
 
 import pandas as pd
 
+from flows.prepPolicy import selectPolicies
 from models.agent import Agent
 from models.policy import Policy
 from services.API import AirlockAPIWrapper
@@ -241,6 +242,7 @@ def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']:
     enrich_agents(matched_agents, policies)
     return matched_agents
 
+
 def moveAgentToRelatedPolicy(
     api: AirlockAPIWrapper,
     agent: Agent,
@@ -284,4 +286,38 @@ def moveAgentToRelatedPolicy(
         return
 
     result = api.agent_move(agent.agentid, target_policy)
-    return result
+    return result
+
+
+def toggleEnforcement(api: AirlockAPIWrapper):
+    choices = ["Audit", "Enforcement", "Exit"]
+    print(colorText("Move devices to which state?:", "yellow"))
+    direction = Selector.select_string(choices, False, False)
+    if direction == "Exit":
+        pass
+    else:
+        devices = selectAgents(api)
+        for device in devices:
+            print(device.hostname)
+        confirm = Selector.confirm("Would you like to continue with these devices? Y/N: ")
+        if direction and devices and confirm:
+            for device in devices:
+                result = moveAgentToRelatedPolicy(api,device, str(direction).lower())
+                logger.info(f"{device.hostname}: result: {result}")
+    get_sanitized_input("Press enter to continue")
+
+def moveAgents(api: AirlockAPIWrapper):
+    devices = selectAgents(api)
+    for device in devices:
+        print(device.hostname)
+    confirm_devices = Selector.confirm("Would you like to continue with these devices? Y/N: ")
+    if devices and confirm_devices:
+        policies = selectPolicies(api, False)
+        confirm_move = Selector.confirm(f"Would you like to move these devices to {policies[0].name}?")
+        if confirm_move:
+            for device in devices:
+                result = api.agent_move(device.agentid, policies[0].groupid)
+                logger.info(f"{device.hostname}: result: {result}")
+        else:
+            logger.info("Exiting without change")
+    get_sanitized_input("Press enter to continue")

services/security.py

@@ -135,15 +135,15 @@ def getAPI(USERNAME, SERVICE_NAME):
         raise ValueError("Failed to retrieve API key after 3 incorrect attempts.")
     else:
         logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.")
-        api_key = getpass(f"🗝️ No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip()
+        api_key = getpass(f"No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip()
         print("Please exit and relaunch program after saving your credential to avoid errors")
 
         while True:
-            password = getpass("🔓 Create a password to encrypt your API key: ")
+            password = getpass("Create a password to encrypt your API key: ")
-            confirm_password = getpass("🔒 Confirm your password: ")
+            confirm_password = getpass("Confirm your password: ")
 
             if password != confirm_password:
-                logging.warning("❌ Passwords do not match. Try again.")
+                logging.warning("Passwords do not match. Try again.")
                 continue
 
             if check_password_complexity(password):
@@ -155,7 +155,7 @@ def getAPI(USERNAME, SERVICE_NAME):
                     logging.error(f"Failed to store API key: {e}")
                     break
             else:
-                logging.warning("❌ Password does not meet complexity requirements. Try again.")
+                logging.warning("Password does not meet complexity requirements. Try again.")
 
 
 class APIKeyManager:

utils/menus.py

@@ -29,10 +29,13 @@ from flows.prepPolicy import (
     testChange,
 )
 from flows.quietAgent import findQuietAgents
-from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAgents
+from services.agenthandler import (
+    findAgents,
+    moveAgents,
+    toggleEnforcement,
+)
 from services.API import AirlockAPIWrapper
 from utils.configmanager import load_env
-from utils.selector import Selector
 from utils.utils import (
     areYouSure,
     clear_screen,
@@ -42,6 +45,7 @@ from utils.utils import (
     locked,
     open_directory,
     printEnforceChecklist,
+    welcome,
 )
 
 logger = logging.getLogger(__name__)
@@ -54,45 +58,27 @@ def menu_main(api: AirlockAPIWrapper):
     while True:
         clear_screen()
         displayIntro()
-        # Add Settings, and give option to change working dir
+        welcome()
-        print(colorText("1. ✅  - Move Device(s) to local approval (Placeholder)", "yellow"))
+        print(colorText("1. 🔍  - Device Search", "yellow"))
-        print(colorText("2. 🎫  - OTP", "yellow"))
+        print(colorText("2. 🔀  - Move Device(s)", "yellow"))
-        print(colorText("3. 🔄  - Move to Audit/Enforcement", "yellow"))
+        print(colorText("3. 🎫  - One Time Pass (OTP)", "yellow"))
-        print(colorText("4. 🔍  - Device Search", "yellow"))
+        print(colorText("4. 🔇  - Find Quiet Hosts", "yellow"))
-        print(colorText("5. 🔇  - Find Quiet Hosts", "yellow"))
+        if extras == "POLICYPREP" : print(colorText("5. 🛡️   - Policy Enforcement Tools", "yellow"))
-        if extras == "POLICYPREP" : print(colorText("6. 🛡️  - Policy Enforcement Tools", "yellow"))
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
         print(colorText("F. 📂  - Open Working Directory", "yellow"))
         print(colorText("S. 🛠️  - Settings", "yellow"))
         print(colorText("Q. 🔚  - Quit", "yellow"))
 
         choice = get_sanitized_input("\nEnter Menu Item: ")
         if choice == "1":
-            print("This Feature is still in development")
+             findAgents(api,False)
-            get_sanitized_input("Press enter to continue")
         elif choice == "2":
-            menu_otp(api)
+            menu_move(api)
         elif choice == "3":
-            choices = ["audit", "enforcement", "Exit"]
+            menu_otp(api)
-            print(colorText("Move devices to which state?:", "yellow"))
-            direction = Selector.select_string(choices, False, False)
-            if direction == "Exit":
-                pass
-            else:
-                devices = selectAgents(api)
-                print(colorText("Would you like to continue with these devices?","white"))
-                for device in devices:
-                    print(device.hostname)
-                confirm = Selector.confirm()
-                if direction and devices and confirm:
-                    for device in devices:
-                        result = moveAgentToRelatedPolicy(api,device, str(direction))
-                        logger.info(f"{device.hostname}: result: {result}")
-            get_sanitized_input("Press enter to continue")
         elif choice == "4":
-            findAgents(api,False)
+           findQuietAgents(api)
         elif choice == "5":
-            findQuietAgents(api)
-        elif choice == "6":
             if extras == "POLICYPREP": menu_policymanagment(api)
         elif choice.upper() == "F":
             open_directory(working_dir)
@@ -214,14 +200,47 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
         else:
             print(colorText("Invalid choice. Please try again.", "red"))
 
+def menu_move(api: AirlockAPIWrapper):
+    working_dir = load_env("WORKING_DIR")
+    while True:
+        clear_screen()
+        displayIntro()
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= ↔️  Agent Movement ↔️  =-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
+        print(colorText("1. ✅  - Move to local approval (Placeholder)", "yellow"))
+        print(colorText("2. 🔄  - Move to Audit/Enforcement", "yellow"))
+        print(colorText("3. 🔀  - Move - Other", "yellow"))
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
+        print(colorText("F. 📂  - Open Working Directory", "yellow"))
+        print(colorText("S. 🛠️  - Settings", "yellow"))
+        print(colorText("B. 🔙  - Back", "yellow"))
+
+        choice = get_sanitized_input("Enter your choice: ")
+
+        if choice == "1":
+            print("This Feature is still in development")
+            get_sanitized_input("Press enter to continue")
+        elif choice == "2":
+            toggleEnforcement(api)
+        elif choice == "3":
+            moveAgents(api)
+        elif choice.upper() == "F":
+            open_directory(working_dir)
+        elif choice.upper() == "S":
+            menu_settings()
+        elif choice.upper() == "B":
+            break
+    
+
 def menu_otp(api: AirlockAPIWrapper):
     working_dir = load_env("WORKING_DIR")
     while True:
-
+        clear_screen()
-        print(colorText("\n--- 🎫 OTP Submenu 🎫 ---", "cyan"))
+        displayIntro()
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🎫  OTP 🎫 =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
         print(colorText("1. 🔐  -Generate OTPs", "cyan"))
         print(colorText("2. 📊  -OTP Activities By Agent", "cyan"))
         print(colorText("3. ❌  -Revoke OTPs", "cyan"))
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
         print(colorText("F. 📂  - Open Working Directory", "yellow"))
         print(colorText("S. 🛠️  - Settings", "yellow"))
         print(colorText("B. 🔙  - Back", "yellow"))
@@ -247,8 +266,12 @@ def menu_otp(api: AirlockAPIWrapper):
 def menu_policymanagment(api: AirlockAPIWrapper):
     working_dir = load_env("WORKING_DIR")
     while True:
+        clear_screen()
+        displayIntro()
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🛡️  Policy Tools 🛡️  =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))       
         print(colorText("1. 🔒  - Prepare Policy For Enforcement", "yellow"))
         print(colorText("2. 🔄  - Update Audit Policies from Enforcement Policies", "yellow"))
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
         print(colorText("F. 📂  - Open Working Directory", "yellow"))
         print(colorText("S. 🛠️  - Settings", "yellow"))
         print(colorText("B. 🔙  - Back", "yellow"))
@@ -272,8 +295,12 @@ def menu_policymanagment(api: AirlockAPIWrapper):
             print(colorText("Invalid choice. Please try again.", "red"))
 def menu_settings():
     while True:
-        print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
+        clear_screen()
+        displayIntro()
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🛠️  Settings 🛠️  =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
         print(colorText("This Feature is still in development, if you have ideas for options you would like to see, let us know.", "cyan"))
+        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
+
         # print(colorText("2. Sub-option B","cyan"))
         print(colorText("B. 🔙  - Back", "yellow"))
         choice = get_sanitized_input("Enter your choice: ")

utils/selector.py

@@ -316,5 +316,5 @@ class Selector:
             print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow"))
             return [pd.Series(row) for row in items if row not in selected]
         else:
-            print(colorText("⚠️ Invalid mode. Returning all rows.", "yellow"))
+            print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
-            return [pd.Series(row) for row in items]
+            return []

utils/utils.py

@@ -166,6 +166,7 @@ def displayIntro():
             "cyan",
         )
     )
+def welcome():
     print(
         colorText(
             "=================================================================================",